diff --git a/src/agents/tracing/traces.py b/src/agents/tracing/traces.py index 591d1f3980..a703acf317 100644 --- a/src/agents/tracing/traces.py +++ b/src/agents/tracing/traces.py @@ -26,15 +26,26 @@ def _finish_on_generator_exit(trace: Trace) -> None: finalizing cannot rewrite the caller's context, and the caller keeps seeing this trace as current until its own scope ends. Raising would only add a crash on top of that. - The tolerance is deliberately limited to this path. An explicit ``finish`` from the - wrong context is a context-ownership violation rather than an unavoidable one, so it - still raises. + The tolerance is deliberately limited to this path, and within it to the reset itself. + An explicit ``finish`` from the wrong context is a context-ownership violation rather + than an unavoidable one, so it still raises, and a processor that fails during + ``finish`` still surfaces rather than being mistaken for a foreign token. + + The reset runs in a ``finally`` so that a failing ``finish`` still releases the scope + instead of leaving the finished trace current, and the token is cleared only once the + reset has either succeeded or hit the foreign-context ``ValueError`` it expects. An + unexpected reset failure therefore leaves the handle in place rather than losing it. """ try: - trace.finish(reset_current=True) - except ValueError: - logger.debug("Skipping trace context reset, token belongs to another context") - trace._prev_context_token = None # type: ignore[attr-defined] + trace.finish(reset_current=False) + finally: + token: contextvars.Token[Trace | None] | None = trace._prev_context_token # type: ignore[attr-defined] + if token is not None: + try: + Scope.reset_current_trace(token) + except ValueError: + logger.debug("Skipping trace context reset, token belongs to another context") + trace._prev_context_token = None # type: ignore[attr-defined] class Trace(abc.ABC): diff --git a/tests/tracing/test_traces_impl.py b/tests/tracing/test_traces_impl.py index fc24580dea..dbf4a5f82e 100644 --- a/tests/tracing/test_traces_impl.py +++ b/tests/tracing/test_traces_impl.py @@ -237,3 +237,42 @@ def test_reattached_trace_restores_scope_without_reemitting_processor_events() - assert processor.started == ["trace-123"] assert processor.ended == ["trace-123"] assert Scope.get_current_trace() is None + + +async def test_generator_close_surfaces_processor_failure() -> None: + """A processor failing during close must not be mistaken for a foreign token. + + ``finish`` calls ``on_trace_end`` before resetting the scope, so catching every + ``ValueError`` around the whole call would swallow a processor failure, drop the saved + token, and leave the finished trace current for everything that ran afterwards. + """ + Scope.set_current_trace(None) + + class FailingProcessor(DummyProcessor): + def on_trace_end(self, trace: Trace) -> None: + raise ValueError("processor exploded") + + trace = TraceImpl( + name="processor-failure", + trace_id="trace-processor-failure", + group_id=None, + metadata=None, + processor=cast(Any, FailingProcessor()), + ) + + async def stream() -> AsyncGenerator[int, None]: + with trace: + yield 1 + + generator = stream() + assert await generator.asend(None) == 1 + + with pytest.raises(ValueError, match="processor exploded"): + await generator.aclose() + + # The processor failure is the one that propagates, and the scope is still released: + # running the reset in a finally keeps a failing finish from leaving the trace current. + assert Scope.get_current_trace() is None + assert trace._prev_context_token is None + + Scope.set_current_trace(None)