diff --git a/src/agents/tracing/scope.py b/src/agents/tracing/scope.py index 1d31c1bd1d..66146c4184 100644 --- a/src/agents/tracing/scope.py +++ b/src/agents/tracing/scope.py @@ -31,8 +31,25 @@ def set_current_span(cls, span: "Span[Any] | None") -> "contextvars.Token[Span[A return _current_span.set(span) @classmethod - def reset_current_span(cls, token: "contextvars.Token[Span[Any] | None]") -> None: - _current_span.reset(token) + def reset_current_span( + cls, + token: "contextvars.Token[Span[Any] | None]", + prev_span: "Span[Any] | None" = None, + ) -> None: + try: + _current_span.reset(token) + except ValueError: + # Token was created in a different Context. This can happen when multiple + # Runner.run() calls execute concurrently via asyncio.gather(). + # Fall back to setting the previous value directly. + # See: https://github.com/openai/openai-agents-python/issues/2246 + logger.warning( + "Tracing context mismatch detected during concurrent execution. " + "Span context was reset using fallback. This may affect trace hierarchy " + "in concurrent scenarios. Consider using asyncio.create_task() for concurrent " + "Runner.run() calls to ensure proper context isolation." + ) + _current_span.set(prev_span) @classmethod def get_current_trace(cls) -> "Trace | None": @@ -44,6 +61,29 @@ def set_current_trace(cls, trace: "Trace | None") -> "contextvars.Token[Trace | return _current_trace.set(trace) @classmethod - def reset_current_trace(cls, token: "contextvars.Token[Trace | None]") -> None: + def reset_current_trace( + cls, + token: "contextvars.Token[Trace | None]", + prev_trace: "Trace | None" = None, + ) -> None: + """Reset the current trace to its previous value. + + Uses token-based reset when possible, with fallback to direct set for + concurrent execution scenarios where Context objects may differ. + See: https://github.com/openai/openai-agents-python/issues/2246 + """ logger.debug("Resetting current trace") - _current_trace.reset(token) + try: + _current_trace.reset(token) + except ValueError: + # Token was created in a different Context. This can happen when multiple + # Runner.run() calls execute concurrently via asyncio.gather(). + # Fall back to setting the previous value directly. + # See: https://github.com/openai/openai-agents-python/issues/2246 + logger.warning( + "Tracing context mismatch detected during concurrent execution. " + "Trace context was reset using fallback. This may affect trace hierarchy " + "in concurrent scenarios. Consider using asyncio.create_task() for concurrent " + "Runner.run() calls to ensure proper context isolation." + ) + _current_trace.set(prev_trace) diff --git a/src/agents/tracing/spans.py b/src/agents/tracing/spans.py index dbde6f9ecf..478f52fea5 100644 --- a/src/agents/tracing/spans.py +++ b/src/agents/tracing/spans.py @@ -182,11 +182,12 @@ class NoOpSpan(Span[TSpanData]): span_data: The operation-specific data for this span. """ - __slots__ = ("_span_data", "_prev_span_token") + __slots__ = ("_span_data", "_prev_span_token", "_prev_span") def __init__(self, span_data: TSpanData): self._span_data = span_data self._prev_span_token: contextvars.Token[Span[TSpanData] | None] | None = None + self._prev_span: Span[Any] | None = None @property def trace_id(self) -> str: @@ -206,12 +207,14 @@ def parent_id(self) -> str | None: def start(self, mark_as_current: bool = False): if mark_as_current: + self._prev_span = Scope.get_current_span() self._prev_span_token = Scope.set_current_span(self) def finish(self, reset_current: bool = False) -> None: if reset_current and self._prev_span_token is not None: - Scope.reset_current_span(self._prev_span_token) + Scope.reset_current_span(self._prev_span_token, self._prev_span) self._prev_span_token = None + self._prev_span = None def __enter__(self) -> Span[TSpanData]: self.start(mark_as_current=True) @@ -253,6 +256,7 @@ class SpanImpl(Span[TSpanData]): "_ended_at", "_error", "_prev_span_token", + "_prev_span", "_processor", "_span_data", ) @@ -273,6 +277,7 @@ def __init__( self._processor = processor self._error: SpanError | None = None self._prev_span_token: contextvars.Token[Span[TSpanData] | None] | None = None + self._prev_span: Span[Any] | None = None self._span_data = span_data @property @@ -299,6 +304,7 @@ def start(self, mark_as_current: bool = False): self._started_at = util.time_iso() self._processor.on_span_start(self) if mark_as_current: + self._prev_span = Scope.get_current_span() self._prev_span_token = Scope.set_current_span(self) def finish(self, reset_current: bool = False) -> None: @@ -309,8 +315,9 @@ def finish(self, reset_current: bool = False) -> None: self._ended_at = util.time_iso() self._processor.on_span_end(self) if reset_current and self._prev_span_token is not None: - Scope.reset_current_span(self._prev_span_token) + Scope.reset_current_span(self._prev_span_token, self._prev_span) self._prev_span_token = None + self._prev_span = None def __enter__(self) -> Span[TSpanData]: self.start(mark_as_current=True) diff --git a/src/agents/tracing/traces.py b/src/agents/tracing/traces.py index ff286de4ff..2cba78b49f 100644 --- a/src/agents/tracing/traces.py +++ b/src/agents/tracing/traces.py @@ -143,6 +143,7 @@ class NoOpTrace(Trace): def __init__(self): self._started = False self._prev_context_token: contextvars.Token[Trace | None] | None = None + self._prev_trace: Trace | None = None def __enter__(self) -> Trace: if self._started: @@ -160,12 +161,14 @@ def __exit__(self, exc_type, exc_val, exc_tb): def start(self, mark_as_current: bool = False): if mark_as_current: + self._prev_trace = Scope.get_current_trace() self._prev_context_token = Scope.set_current_trace(self) 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) + Scope.reset_current_trace(self._prev_context_token, self._prev_trace) self._prev_context_token = None + self._prev_trace = None @property def trace_id(self) -> str: @@ -208,6 +211,7 @@ class TraceImpl(Trace): "group_id", "metadata", "_prev_context_token", + "_prev_trace", "_processor", "_started", ) @@ -225,6 +229,7 @@ def __init__( self.group_id = group_id self.metadata = metadata self._prev_context_token: contextvars.Token[Trace | None] | None = None + self._prev_trace: Trace | None = None self._processor = processor self._started = False @@ -244,6 +249,7 @@ def start(self, mark_as_current: bool = False): self._processor.on_trace_start(self) if mark_as_current: + self._prev_trace = Scope.get_current_trace() self._prev_context_token = Scope.set_current_trace(self) def finish(self, reset_current: bool = False): @@ -253,8 +259,9 @@ def finish(self, reset_current: bool = False): self._processor.on_trace_end(self) if reset_current and self._prev_context_token is not None: - Scope.reset_current_trace(self._prev_context_token) + Scope.reset_current_trace(self._prev_context_token, self._prev_trace) self._prev_context_token = None + self._prev_trace = None def __enter__(self) -> Trace: if self._started: