Skip to content

Commit e591454

Browse files
seratchseratch
committed
fix: #2540 reattach resumed traces without duplicate trace starts
1 parent 4581307 commit e591454

6 files changed

Lines changed: 519 additions & 29 deletions

File tree

‎src/agents/run.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,8 @@
107107
from .run_state import RunState
108108
from .tool import dispose_resolved_computers
109109
from .tool_guardrails import ToolInputGuardrailResult, ToolOutputGuardrailResult
110-
from .tracing import Span, SpanError, agent_span, get_current_trace, trace
111-
from .tracing.context import TraceCtxManager
110+
from .tracing import Span, SpanError, agent_span, get_current_trace
111+
from .tracing.context import TraceCtxManager, create_trace_for_run
112112
from .tracing.span_data import AgentSpanData
113113
from .util import _error_tracing
114114

@@ -549,6 +549,8 @@ async def run(
549549
metadata=trace_metadata,
550550
tracing=trace_config,
551551
disabled=run_config.tracing_disabled,
552+
trace_state=run_state._trace_state if run_state is not None else None,
553+
reattach_resumed_trace=is_resumed_state,
552554
):
553555
if is_resumed_state and run_state is not None:
554556
run_state.set_trace(get_current_trace())
@@ -1519,17 +1521,15 @@ def run_streamed(
15191521
# If there's already a trace, we don't create a new one. In addition, we can't end the
15201522
# trace here, because the actual work is done in `stream_events` and this method ends
15211523
# before that.
1522-
new_trace = (
1523-
None
1524-
if get_current_trace()
1525-
else trace(
1526-
workflow_name=trace_workflow_name,
1527-
trace_id=trace_id,
1528-
group_id=trace_group_id,
1529-
metadata=trace_metadata,
1530-
tracing=trace_config,
1531-
disabled=run_config.tracing_disabled,
1532-
)
1524+
new_trace = create_trace_for_run(
1525+
workflow_name=trace_workflow_name,
1526+
trace_id=trace_id,
1527+
group_id=trace_group_id,
1528+
metadata=trace_metadata,
1529+
tracing=trace_config,
1530+
disabled=run_config.tracing_disabled,
1531+
trace_state=run_state._trace_state if run_state is not None else None,
1532+
reattach_resumed_trace=is_resumed_state,
15331533
)
15341534
if run_state is not None:
15351535
run_state.set_trace(new_trace or get_current_trace())

‎src/agents/tracing/context.py‎

Lines changed: 79 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,70 @@
44

55
from .config import TracingConfig
66
from .create import get_current_trace, trace
7-
from .traces import Trace
7+
from .traces import Trace, TraceState, _trace_id_was_started, reattach_trace
8+
9+
10+
def _trace_state_matches_effective_settings(
11+
*,
12+
trace_state: TraceState,
13+
workflow_name: str,
14+
trace_id: str | None,
15+
group_id: str | None,
16+
metadata: dict[str, Any] | None,
17+
tracing: TracingConfig | None,
18+
) -> bool:
19+
if trace_state.trace_id is None or trace_state.trace_id != trace_id:
20+
return False
21+
if trace_state.workflow_name != workflow_name:
22+
return False
23+
if trace_state.group_id != group_id:
24+
return False
25+
if trace_state.metadata != metadata:
26+
return False
27+
tracing_api_key = tracing.get("api_key") if tracing is not None else None
28+
return trace_state.tracing_api_key == tracing_api_key
29+
30+
31+
def create_trace_for_run(
32+
*,
33+
workflow_name: str,
34+
trace_id: str | None,
35+
group_id: str | None,
36+
metadata: dict[str, Any] | None,
37+
tracing: TracingConfig | None,
38+
disabled: bool,
39+
trace_state: TraceState | None = None,
40+
reattach_resumed_trace: bool = False,
41+
) -> Trace | None:
42+
"""Return a trace object for this run when one is not already active."""
43+
current_trace = get_current_trace()
44+
if current_trace:
45+
return None
46+
47+
if (
48+
reattach_resumed_trace
49+
and not disabled
50+
and trace_state is not None
51+
and _trace_state_matches_effective_settings(
52+
trace_state=trace_state,
53+
workflow_name=workflow_name,
54+
trace_id=trace_id,
55+
group_id=group_id,
56+
metadata=metadata,
57+
tracing=tracing,
58+
)
59+
and _trace_id_was_started(trace_state.trace_id)
60+
):
61+
return reattach_trace(trace_state)
62+
63+
return trace(
64+
workflow_name=workflow_name,
65+
trace_id=trace_id,
66+
group_id=group_id,
67+
metadata=metadata,
68+
tracing=tracing,
69+
disabled=disabled,
70+
)
871

972

1073
class TraceCtxManager:
@@ -18,6 +81,8 @@ def __init__(
1881
metadata: dict[str, Any] | None,
1982
tracing: TracingConfig | None,
2083
disabled: bool,
84+
trace_state: TraceState | None = None,
85+
reattach_resumed_trace: bool = False,
2186
):
2287
self.trace: Trace | None = None
2388
self.workflow_name = workflow_name
@@ -26,18 +91,21 @@ def __init__(
2691
self.metadata = metadata
2792
self.tracing = tracing
2893
self.disabled = disabled
94+
self.trace_state = trace_state
95+
self.reattach_resumed_trace = reattach_resumed_trace
2996

3097
def __enter__(self) -> TraceCtxManager:
31-
current_trace = get_current_trace()
32-
if not current_trace:
33-
self.trace = trace(
34-
workflow_name=self.workflow_name,
35-
trace_id=self.trace_id,
36-
group_id=self.group_id,
37-
metadata=self.metadata,
38-
tracing=self.tracing,
39-
disabled=self.disabled,
40-
)
98+
self.trace = create_trace_for_run(
99+
workflow_name=self.workflow_name,
100+
trace_id=self.trace_id,
101+
group_id=self.group_id,
102+
metadata=self.metadata,
103+
tracing=self.tracing,
104+
disabled=self.disabled,
105+
trace_state=self.trace_state,
106+
reattach_resumed_trace=self.reattach_resumed_trace,
107+
)
108+
if self.trace:
41109
assert self.trace is not None
42110
self.trace.start(mark_as_current=True)
43111
return self

‎src/agents/tracing/traces.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import abc
44
import contextvars
5+
import threading
56
from collections.abc import Mapping
67
from dataclasses import dataclass, field
78
from typing import Any
@@ -219,6 +220,119 @@ def to_json(self, *, include_tracing_api_key: bool = False) -> dict[str, Any] |
219220
return payload
220221

221222

223+
_started_trace_ids: set[str] = set()
224+
_started_trace_ids_lock = threading.Lock()
225+
226+
227+
def _mark_trace_id_started(trace_id: str | None) -> None:
228+
if not trace_id or trace_id == "no-op":
229+
return
230+
with _started_trace_ids_lock:
231+
_started_trace_ids.add(trace_id)
232+
233+
234+
def _trace_id_was_started(trace_id: str | None) -> bool:
235+
if not trace_id or trace_id == "no-op":
236+
return False
237+
with _started_trace_ids_lock:
238+
return trace_id in _started_trace_ids
239+
240+
241+
class ReattachedTrace(Trace):
242+
"""A trace context rebuilt from persisted state without re-emitting trace start events."""
243+
244+
__slots__ = (
245+
"_name",
246+
"_trace_id",
247+
"_tracing_api_key",
248+
"group_id",
249+
"metadata",
250+
"_prev_context_token",
251+
"_started",
252+
)
253+
254+
def __init__(
255+
self,
256+
*,
257+
name: str,
258+
trace_id: str,
259+
group_id: str | None,
260+
metadata: dict[str, Any] | None,
261+
tracing_api_key: str | None,
262+
) -> None:
263+
self._name = name
264+
self._trace_id = trace_id
265+
self._tracing_api_key = tracing_api_key
266+
self.group_id = group_id
267+
self.metadata = metadata
268+
self._prev_context_token: contextvars.Token[Trace | None] | None = None
269+
self._started = False
270+
271+
@property
272+
def trace_id(self) -> str:
273+
return self._trace_id
274+
275+
@property
276+
def name(self) -> str:
277+
return self._name
278+
279+
@property
280+
def tracing_api_key(self) -> str | None:
281+
return self._tracing_api_key
282+
283+
def start(self, mark_as_current: bool = False):
284+
if self._started:
285+
return
286+
287+
self._started = True
288+
_mark_trace_id_started(self.trace_id)
289+
290+
if mark_as_current:
291+
self._prev_context_token = Scope.set_current_trace(self)
292+
293+
def finish(self, reset_current: bool = False):
294+
if not self._started:
295+
return
296+
297+
if reset_current and self._prev_context_token is not None:
298+
Scope.reset_current_trace(self._prev_context_token)
299+
self._prev_context_token = None
300+
301+
def __enter__(self) -> Trace:
302+
if self._started:
303+
if not self._prev_context_token:
304+
logger.error("Trace already started but no context token set")
305+
return self
306+
307+
self.start(mark_as_current=True)
308+
return self
309+
310+
def __exit__(self, exc_type, exc_val, exc_tb):
311+
self.finish(reset_current=exc_type is not GeneratorExit)
312+
313+
def export(self) -> dict[str, Any] | None:
314+
return {
315+
"object": "trace",
316+
"id": self.trace_id,
317+
"workflow_name": self.name,
318+
"group_id": self.group_id,
319+
"metadata": self.metadata,
320+
}
321+
322+
323+
def reattach_trace(trace_state: TraceState) -> Trace | None:
324+
"""Build a live trace context from persisted state without notifying processors."""
325+
if trace_state.trace_id is None:
326+
return None
327+
return ReattachedTrace(
328+
name=trace_state.workflow_name or "Agent workflow",
329+
trace_id=trace_state.trace_id,
330+
group_id=trace_state.group_id,
331+
metadata=dict(trace_state.metadata) if trace_state.metadata is not None else None,
332+
tracing_api_key=trace_state.tracing_api_key,
333+
)
334+
335+
222336
class NoOpTrace(Trace):
223337
"""A no-op implementation of Trace that doesn't record any data.
224338
@@ -347,6 +461,7 @@ def start(self, mark_as_current: bool = False):
347461

348462
self._started = True
349463
self._processor.on_trace_start(self)
464+
_mark_trace_id_started(self.trace_id)
350465

351466
if mark_as_current:
352467
self._prev_context_token = Scope.set_current_trace(self)

0 commit comments

Comments
 (0)