fix: #2540 reattach resumed traces without duplicate trace starts (#2547)

Author: seratch

src/agents/run.py

@@ -107,8 +107,8 @@
 from .run_state import RunState
 from .tool import dispose_resolved_computers
 from .tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
-from .tracing import Span, SpanError, agent_span, get_current_trace, trace
-from .tracing.context import TraceCtxManager
+from .tracing import Span, SpanError, agent_span, get_current_trace
+from .tracing.context import TraceCtxManager, create_trace_for_run
 from .tracing.span_data import AgentSpanData
 from .util import _error_tracing
 
@@ -549,6 +549,8 @@ async def run(
             metadata=trace_metadata,
             tracing=trace_config,
             disabled=run_config.tracing_disabled,
+            trace_state=run_state._trace_state if run_state is not None else None,
+            reattach_resumed_trace=is_resumed_state,
         ):
             if is_resumed_state and run_state is not None:
                 run_state.set_trace(get_current_trace())
@@ -1519,17 +1521,15 @@ def run_streamed(
         # If there's already a trace, we don't create a new one. In addition, we can't end the
         # trace here, because the actual work is done in `stream_events` and this method ends
         # before that.
-        new_trace = (
-            None
-            if get_current_trace()
-            else trace(
-                workflow_name=trace_workflow_name,
-                trace_id=trace_id,
-                group_id=trace_group_id,
-                metadata=trace_metadata,
-                tracing=trace_config,
-                disabled=run_config.tracing_disabled,
-            )
+        new_trace = create_trace_for_run(
+            workflow_name=trace_workflow_name,
+            trace_id=trace_id,
+            group_id=trace_group_id,
+            metadata=trace_metadata,
+            tracing=trace_config,
+            disabled=run_config.tracing_disabled,
+            trace_state=run_state._trace_state if run_state is not None else None,
+            reattach_resumed_trace=is_resumed_state,
         )
         if run_state is not None:
             run_state.set_trace(new_trace or get_current_trace())

src/agents/run_state.py

@@ -99,8 +99,8 @@
 # 2. Keep older readable versions in SUPPORTED_SCHEMA_VERSIONS for backward reads.
 # 3. to_json() always emits CURRENT_SCHEMA_VERSION.
 # 4. Forward compatibility is intentionally fail-fast (older SDKs reject newer versions).
-CURRENT_SCHEMA_VERSION = "1.2"
-SUPPORTED_SCHEMA_VERSIONS = frozenset({"1.0", "1.1", CURRENT_SCHEMA_VERSION})
+CURRENT_SCHEMA_VERSION = "1.3"
+SUPPORTED_SCHEMA_VERSIONS = frozenset({"1.0", "1.1", "1.2", CURRENT_SCHEMA_VERSION})
 
 _FUNCTION_OUTPUT_ADAPTER: TypeAdapter[FunctionCallOutput] = TypeAdapter(FunctionCallOutput)
 _COMPUTER_OUTPUT_ADAPTER: TypeAdapter[ComputerCallOutput] = TypeAdapter(ComputerCallOutput)

src/agents/tracing/context.py

@@ -4,7 +4,88 @@
 
 from .config import TracingConfig
 from .create import get_current_trace, trace
-from .traces import Trace
+from .traces import (
+    Trace,
+    TraceState,
+    _hash_tracing_api_key,
+    _trace_id_was_started,
+    reattach_trace,
+)
+
+
+def _get_tracing_api_key(tracing: TracingConfig | None) -> str | None:
+    return tracing.get("api_key") if tracing is not None else None
+
+
+def _trace_state_matches_effective_settings(
+    *,
+    trace_state: TraceState,
+    workflow_name: str,
+    trace_id: str | None,
+    group_id: str | None,
+    metadata: dict[str, Any] | None,
+    tracing: TracingConfig | None,
+) -> bool:
+    if trace_state.trace_id is None or trace_state.trace_id != trace_id:
+        return False
+    if trace_state.workflow_name != workflow_name:
+        return False
+    if trace_state.group_id != group_id:
+        return False
+    if trace_state.metadata != metadata:
+        return False
+    tracing_api_key = _get_tracing_api_key(tracing)
+    if trace_state.tracing_api_key is not None:
+        return trace_state.tracing_api_key == tracing_api_key
+    if trace_state.tracing_api_key_hash is not None:
+        # A fingerprint lets stripped RunState snapshots prove the caller
+        # re-supplied the same explicit key.
+        return trace_state.tracing_api_key_hash == _hash_tracing_api_key(tracing_api_key)
+    return tracing_api_key is None
+
+
+def create_trace_for_run(
+    *,
+    workflow_name: str,
+    trace_id: str | None,
+    group_id: str | None,
+    metadata: dict[str, Any] | None,
+    tracing: TracingConfig | None,
+    disabled: bool,
+    trace_state: TraceState | None = None,
+    reattach_resumed_trace: bool = False,
+) -> Trace | None:
+    """Return a trace object for this run when one is not already active."""
+    current_trace = get_current_trace()
+    if current_trace:
+        return None
+
+    if (
+        reattach_resumed_trace
+        and not disabled
+        and trace_state is not None
+        and _trace_id_was_started(trace_state.trace_id)
+        and _trace_state_matches_effective_settings(
+            trace_state=trace_state,
+            workflow_name=workflow_name,
+            trace_id=trace_id,
+            group_id=group_id,
+            metadata=metadata,
+            tracing=tracing,
+        )
+    ):
+        # Reuse the live key because secure snapshots may persist only the
+        # fingerprint, not the secret itself.
+        return reattach_trace(trace_state, tracing_api_key=_get_tracing_api_key(tracing))
+
+    return trace(
+        workflow_name=workflow_name,
+        trace_id=trace_id,
+        group_id=group_id,
+        metadata=metadata,
+        tracing=tracing,
+        disabled=disabled,
+    )
 
 
 class TraceCtxManager:
@@ -18,6 +99,8 @@ def __init__(
         metadata: dict[str, Any] | None,
         tracing: TracingConfig | None,
         disabled: bool,
+        trace_state: TraceState | None = None,
+        reattach_resumed_trace: bool = False,
     ):
         self.trace: Trace | None = None
         self.workflow_name = workflow_name
@@ -26,18 +109,21 @@ def __init__(
         self.metadata = metadata
         self.tracing = tracing
         self.disabled = disabled
+        self.trace_state = trace_state
+        self.reattach_resumed_trace = reattach_resumed_trace
 
     def __enter__(self) -> TraceCtxManager:
-        current_trace = get_current_trace()
-        if not current_trace:
-            self.trace = trace(
-                workflow_name=self.workflow_name,
-                trace_id=self.trace_id,
-                group_id=self.group_id,
-                metadata=self.metadata,
-                tracing=self.tracing,
-                disabled=self.disabled,
-            )
+        self.trace = create_trace_for_run(
+            workflow_name=self.workflow_name,
+            trace_id=self.trace_id,
+            group_id=self.group_id,
+            metadata=self.metadata,
+            tracing=self.tracing,
+            disabled=self.disabled,
+            trace_state=self.trace_state,
+            reattach_resumed_trace=self.reattach_resumed_trace,
+        )
+        if self.trace:
             assert self.trace is not None
             self.trace.start(mark_as_current=True)
         return self