Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion src/agents/tracing/traces.py
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,10 @@ def __enter__(self) -> Trace:
return self

def __exit__(self, exc_type, exc_val, exc_tb):
# Same-task generator close should still reset the context var so the
# NoOpTrace doesn't linger as the current trace. Cross-context closes
# (different task / event loop) raise ValueError from contextvars,
# which the finish() helper swallows below.
self.finish(reset_current=True)

def start(self, mark_as_current: bool = False):
Expand All @@ -407,7 +411,14 @@ def start(self, mark_as_current: bool = False):

def finish(self, reset_current: bool = False):
if reset_current and self._prev_context_token is not None:
Scope.reset_current_trace(self._prev_context_token)
try:
Scope.reset_current_trace(self._prev_context_token)
except ValueError:
# The context token was created in a different Context (e.g.
# the trace was entered in another task and the generator is
# closing from the parent's context). Skipping reset here is
# safe: that other Context owns its own contextvar copy.
pass
self._prev_context_token = None

@property
Expand Down
36 changes: 36 additions & 0 deletions tests/tracing/test_traces_impl.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import contextvars
import logging
from typing import Any, cast

Expand Down Expand Up @@ -42,6 +43,41 @@ def test_no_op_trace_double_enter_logs_error(caplog) -> None:
trace.__exit__(None, None, None)


def test_no_op_trace_resets_context_on_same_task_generator_exit() -> None:
"""NoOpTrace must reset the current trace on a same-task GeneratorExit.

A previous version skipped reset unconditionally on GeneratorExit, which
handles the cross-task GC case but leaves NoOpTrace as the current trace
after a normal same-context generator close, suppressing later tracing.
Now the reset is attempted and any ValueError (from a token created in a
different Context) is swallowed.
"""
Scope.set_current_trace(None)
trace = NoOpTrace()
trace.__enter__()
assert trace._prev_context_token is not None
trace.__exit__(GeneratorExit, GeneratorExit(), None)
# Same-task close: reset succeeded, token consumed, current trace cleared.
assert trace._prev_context_token is None
assert Scope.get_current_trace() is None


def test_no_op_trace_swallows_cross_context_reset_error() -> None:
"""A token created in a different Context raises ValueError on reset; swallow it."""
Scope.set_current_trace(None)
trace = NoOpTrace()

other_context = contextvars.copy_context()
other_context.run(trace.__enter__)
token = trace._prev_context_token
assert token is not None

# Resetting from our context (not the one that set it) raises ValueError;
# the helper must swallow that and clear the stored token.
trace.__exit__(GeneratorExit, GeneratorExit(), None)
assert trace._prev_context_token is None


def test_trace_impl_lifecycle_sets_scope() -> None:
Scope.set_current_trace(None)
processor = DummyProcessor()
Expand Down