|
| 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