Skip to content

Commit 035be54

Browse files
feat(tracing): span each tool call on the sync agent path
The sync path only spans Model.stream_response, the inference call. The Runner executes tools outside it, so tool time falls into gaps between spans and nothing accounts for it. On dev-sgp an agent with slow tools reported a third of its real wall-clock time. SyncTracingHooks ports the span lifecycle the Temporal plugin already uses, minus the activity plumbing, so both paths emit the same trace shape.
1 parent 55930f2 commit 035be54

2 files changed

Lines changed: 299 additions & 0 deletions

File tree

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
"""Tool-call tracing for the sync (non-Temporal) OpenAI Agents path.
2+
3+
The sync path only spans `Model.stream_response`, the inference call.
4+
The ``Runner`` executes tools *outside* that call, so tool time lands in gaps
5+
between spans and no span accounts for it. Measured on dev-sgp, an agent whose
6+
tools are slow reported a third of its real wall-clock time.
7+
8+
The Temporal plugin already solves this with ``TemporalStreamingHooks``. This is
9+
the same span lifecycle without the Temporal activity plumbing, so both paths
10+
produce the same trace shape.
11+
"""
12+
13+
from __future__ import annotations
14+
15+
from typing import Any
16+
from datetime import timedelta
17+
18+
from agents import Tool, Agent, RunHooks
19+
from agents.run_context import RunContextWrapper
20+
21+
from agentex.lib.utils.logging import make_logger
22+
23+
logger = make_logger(__name__)
24+
25+
_TRACE_TIMEOUT = timedelta(seconds=5)
26+
# Cap tool-result span output so a large payload can't bloat the trace.
27+
_MAX_SPAN_OUTPUT_CHARS = 2000
28+
29+
30+
def _get_adk() -> Any:
31+
"""Lazily import the adk facade so this module stays cheap to import."""
32+
from agentex.lib import adk
33+
34+
return adk
35+
36+
37+
class SyncTracingHooks(RunHooks):
38+
"""Opens one span per tool call, closed when the tool returns.
39+
40+
Span width is therefore the tool's real execution time, which is the whole
41+
point: a duration recorded as an attribute on a zero-width span is invisible
42+
on a timeline.
43+
44+
Every tracing call is best-effort. A tracing failure must never break a turn.
45+
"""
46+
47+
def __init__(
48+
self,
49+
trace_id: str | None = None,
50+
parent_span_id: str | None = None,
51+
task_id: str | None = None,
52+
) -> None:
53+
self.trace_id = trace_id
54+
self.parent_span_id = parent_span_id
55+
self.task_id = task_id
56+
# tool_call_id -> open span, so on_tool_end closes the right one.
57+
self._tool_spans: dict[str, Any] = {}
58+
59+
@staticmethod
60+
def _tool_call_id(context: RunContextWrapper, tool: Tool) -> str:
61+
return getattr(context, "tool_call_id", None) or tool.name
62+
63+
@staticmethod
64+
def _tool_arguments(context: RunContextWrapper) -> dict[str, Any]:
65+
raw = getattr(context, "tool_arguments", None)
66+
if isinstance(raw, dict):
67+
return raw
68+
if isinstance(raw, str) and raw:
69+
import json
70+
71+
try:
72+
parsed = json.loads(raw)
73+
except ValueError:
74+
return {"raw": raw[:_MAX_SPAN_OUTPUT_CHARS]}
75+
return parsed if isinstance(parsed, dict) else {"raw": parsed}
76+
return {}
77+
78+
async def on_tool_start(self, context: RunContextWrapper, agent: Agent, tool: Tool) -> None: # noqa: ARG002
79+
if not self.trace_id:
80+
return
81+
try:
82+
span = await _get_adk().tracing.start_span(
83+
trace_id=self.trace_id,
84+
parent_id=self.parent_span_id,
85+
task_id=self.task_id,
86+
name=tool.name,
87+
input={"arguments": self._tool_arguments(context)},
88+
start_to_close_timeout=_TRACE_TIMEOUT,
89+
)
90+
if span is not None:
91+
self._tool_spans[self._tool_call_id(context, tool)] = span
92+
except Exception as e: # noqa: BLE001 - tracing is best-effort
93+
logger.warning(f"[tracing] tool start_span failed (non-fatal): {e}")
94+
95+
async def on_tool_end(
96+
self,
97+
context: RunContextWrapper,
98+
agent: Agent, # noqa: ARG002
99+
tool: Tool,
100+
result: str,
101+
) -> None:
102+
span = self._tool_spans.pop(self._tool_call_id(context, tool), None)
103+
if span is None or not self.trace_id:
104+
return
105+
try:
106+
span.output = {"result": str(result)[:_MAX_SPAN_OUTPUT_CHARS]}
107+
await _get_adk().tracing.end_span(
108+
trace_id=self.trace_id,
109+
span=span,
110+
start_to_close_timeout=_TRACE_TIMEOUT,
111+
)
112+
except Exception as e: # noqa: BLE001 - tracing is best-effort
113+
logger.warning(f"[tracing] tool end_span failed (non-fatal): {e}")
114+
115+
async def close_open_tool_spans(self) -> None:
116+
"""Drain spans whose ``on_tool_end`` never fired.
117+
118+
A runner that dies mid-tool (max-turns, cancellation, an SDK error) never
119+
fires the matching end hook, which would orphan the span. Call this from a
120+
``finally`` around the run.
121+
"""
122+
if not self._tool_spans:
123+
return
124+
orphaned = list(self._tool_spans.items())
125+
self._tool_spans.clear()
126+
for tool_call_id, span in orphaned:
127+
logger.warning(
128+
f"[tracing] tool span for {tool_call_id} left open "
129+
"(on_tool_end never fired); closing as incomplete"
130+
)
131+
try:
132+
span.output = {"result": None, "incomplete": True}
133+
await _get_adk().tracing.end_span(
134+
trace_id=self.trace_id,
135+
span=span,
136+
start_to_close_timeout=_TRACE_TIMEOUT,
137+
)
138+
except Exception as e: # noqa: BLE001 - tracing is best-effort
139+
logger.warning(f"[tracing] draining tool span failed (non-fatal): {e}")
Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
"""Tests for per-tool-call spans on the sync OpenAI-Agents path.
2+
3+
The defect these guard: a span emitted *after* a tool returns has ~zero width,
4+
so the tool's real duration is invisible on a timeline even when it is recorded
5+
as an attribute. These assert the span is opened before the tool runs and closed
6+
after, and that tracing failures never propagate.
7+
"""
8+
9+
from __future__ import annotations
10+
11+
from types import SimpleNamespace
12+
from unittest.mock import AsyncMock, MagicMock
13+
14+
import pytest
15+
from agents.tool_context import ToolContext
16+
17+
from agentex.lib.adk._modules import _sync_tracing_hooks as hooks_mod
18+
19+
SyncTracingHooks = hooks_mod.SyncTracingHooks
20+
21+
22+
def _tool_context(args: str = '{"query": "hi"}') -> ToolContext:
23+
return ToolContext(context=None, tool_name="search", tool_call_id="call_abc", tool_arguments=args)
24+
25+
26+
def _tool(name: str = "search") -> MagicMock:
27+
tool = MagicMock()
28+
tool.name = name
29+
return tool
30+
31+
32+
def _adk(span=None):
33+
adk = SimpleNamespace(tracing=SimpleNamespace(start_span=AsyncMock(return_value=span), end_span=AsyncMock()))
34+
return adk
35+
36+
37+
# --------------------------------------------------------------------------- #
38+
# Argument parsing
39+
# --------------------------------------------------------------------------- #
40+
41+
42+
def test_tool_arguments_valid_dict():
43+
assert SyncTracingHooks._tool_arguments(_tool_context('{"a": 1}')) == {"a": 1}
44+
45+
46+
def test_tool_arguments_garbage_is_preserved_raw():
47+
assert SyncTracingHooks._tool_arguments(_tool_context("not json")) == {"raw": "not json"}
48+
49+
50+
def test_tool_arguments_missing_is_empty():
51+
assert SyncTracingHooks._tool_arguments(SimpleNamespace()) == {}
52+
53+
54+
# --------------------------------------------------------------------------- #
55+
# Span lifecycle
56+
# --------------------------------------------------------------------------- #
57+
58+
59+
@pytest.mark.asyncio
60+
async def test_no_trace_id_is_a_no_op(monkeypatch):
61+
adk = _adk()
62+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
63+
64+
hooks = SyncTracingHooks()
65+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
66+
67+
adk.tracing.start_span.assert_not_awaited()
68+
69+
70+
@pytest.mark.asyncio
71+
async def test_span_opens_with_arguments_and_parent(monkeypatch):
72+
span = MagicMock()
73+
adk = _adk(span)
74+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
75+
76+
hooks = SyncTracingHooks(trace_id="tr1", parent_span_id="root1", task_id="task1")
77+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
78+
79+
kwargs = adk.tracing.start_span.await_args.kwargs
80+
assert kwargs["name"] == "search"
81+
assert kwargs["parent_id"] == "root1"
82+
assert kwargs["input"] == {"arguments": {"query": "hi"}}
83+
84+
85+
@pytest.mark.asyncio
86+
async def test_span_closes_with_result(monkeypatch):
87+
span = MagicMock()
88+
adk = _adk(span)
89+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
90+
91+
hooks = SyncTracingHooks(trace_id="tr1")
92+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
93+
await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "42 rows")
94+
95+
assert span.output == {"result": "42 rows"}
96+
adk.tracing.end_span.assert_awaited_once()
97+
assert hooks._tool_spans == {}
98+
99+
100+
@pytest.mark.asyncio
101+
async def test_result_is_truncated(monkeypatch):
102+
span = MagicMock()
103+
adk = _adk(span)
104+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
105+
106+
hooks = SyncTracingHooks(trace_id="tr1")
107+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
108+
await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "x" * 10_000)
109+
110+
assert len(span.output["result"]) == hooks_mod._MAX_SPAN_OUTPUT_CHARS
111+
112+
113+
@pytest.mark.asyncio
114+
async def test_end_without_start_is_a_no_op(monkeypatch):
115+
adk = _adk()
116+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
117+
118+
hooks = SyncTracingHooks(trace_id="tr1")
119+
await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "result")
120+
121+
adk.tracing.end_span.assert_not_awaited()
122+
123+
124+
@pytest.mark.asyncio
125+
async def test_start_span_failure_does_not_propagate(monkeypatch):
126+
adk = _adk()
127+
adk.tracing.start_span.side_effect = RuntimeError("tracing backend down")
128+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
129+
130+
hooks = SyncTracingHooks(trace_id="tr1")
131+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
132+
133+
assert hooks._tool_spans == {}
134+
135+
136+
@pytest.mark.asyncio
137+
async def test_end_span_failure_does_not_propagate(monkeypatch):
138+
span = MagicMock()
139+
adk = _adk(span)
140+
adk.tracing.end_span.side_effect = RuntimeError("tracing backend down")
141+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
142+
143+
hooks = SyncTracingHooks(trace_id="tr1")
144+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
145+
await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "result")
146+
147+
148+
@pytest.mark.asyncio
149+
async def test_orphaned_spans_are_drained(monkeypatch):
150+
span = MagicMock()
151+
adk = _adk(span)
152+
monkeypatch.setattr(hooks_mod, "_get_adk", lambda: adk)
153+
154+
hooks = SyncTracingHooks(trace_id="tr1")
155+
await hooks.on_tool_start(_tool_context(), MagicMock(), _tool())
156+
await hooks.close_open_tool_spans()
157+
158+
assert span.output["incomplete"] is True
159+
adk.tracing.end_span.assert_awaited_once()
160+
assert hooks._tool_spans == {}

0 commit comments

Comments
 (0)