feat(tracing): wrap a sync agent turn in a root span

Nothing owned the turn on the sync path, so no span covered it and every span
was emitted parentless. A trace was a flat list of fragments whose durations
did not add up to the turn.

run_turn_streamed owns it: one root span, tool spans parented beneath it,
orphans drained, root closed when the stream ends. The root's width is the
turn's real latency, and the UI can draw a waterfall.

Author: mohammadatallah-scale

src/agentex/lib/adk/_modules/_sync_turn.py

@@ -0,0 +1,95 @@
+"""One traced turn on the sync (non-Temporal) OpenAI Agents path.
+
+Agents on this path call ``Runner.run_streamed`` themselves, so nothing owns the
+turn and no span covers it: the provider only spans the inference call, and every
+span is emitted with no parent. A trace is therefore a flat list of fragments
+whose durations do not add up to the turn, and tool time is missing entirely.
+
+``run_turn_streamed`` owns the turn instead. It opens one root span, hangs the
+tool spans off it, drains anything left open, and closes the root when the stream
+is done, so the root's width is the turn's real latency.
+"""
+
+from __future__ import annotations
+
+from typing import Any, AsyncIterator
+from datetime import timedelta
+
+from agents import Agent, Runner, RunConfig
+
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.adk._modules._sync_tracing_hooks import SyncTracingHooks
+
+logger = make_logger(__name__)
+
+_TRACE_TIMEOUT = timedelta(seconds=5)
+
+
+def _get_adk() -> Any:
+    from agentex.lib import adk
+
+    return adk
+
+
+async def run_turn_streamed(
+    starting_agent: Agent,
+    input: Any,
+    *,
+    trace_id: str | None = None,
+    task_id: str | None = None,
+    run_config: RunConfig | None = None,
+    max_turns: int | None = None,
+    span_name: str = "turn",
+) -> AsyncIterator[Any]:
+    """Stream one agent turn, wrapped in a root span with tool spans beneath it.
+
+    Yields the SDK's raw stream events untouched, so callers keep whatever event
+    conversion they already do.
+
+    Tracing is best-effort throughout: if the backend is unreachable the turn
+    still runs and still streams.
+    """
+    root_span = None
+    if trace_id:
+        try:
+            root_span = await _get_adk().tracing.start_span(
+                trace_id=trace_id,
+                task_id=task_id,
+                name=span_name,
+                input={"agent": starting_agent.name},
+                start_to_close_timeout=_TRACE_TIMEOUT,
+            )
+        except Exception as e:  # noqa: BLE001 - tracing is best-effort
+            logger.warning(f"[tracing] turn start_span failed (non-fatal): {e}")
+
+    hooks = SyncTracingHooks(
+        trace_id=trace_id,
+        # Tool spans nest under the turn, which is what lets the UI draw a
+        # waterfall instead of a flat list.
+        parent_span_id=getattr(root_span, "id", None),
+        task_id=task_id,
+    )
+
+    run_kwargs: dict[str, Any] = {"hooks": hooks}
+    if run_config is not None:
+        run_kwargs["run_config"] = run_config
+    if max_turns is not None:
+        run_kwargs["max_turns"] = max_turns
+
+    try:
+        result = Runner.run_streamed(starting_agent, input, **run_kwargs)
+        async for event in result.stream_events():
+            yield event
+    finally:
+        # A turn that dies mid-tool (max-turns, cancellation, an SDK error) never
+        # fires the matching end hook, so drain before closing the root.
+        await hooks.close_open_tool_spans()
+        if root_span is not None and trace_id:
+            try:
+                await _get_adk().tracing.end_span(
+                    trace_id=trace_id,
+                    span=root_span,
+                    start_to_close_timeout=_TRACE_TIMEOUT,
+                )
+            except Exception as e:  # noqa: BLE001 - tracing is best-effort
+                logger.warning(f"[tracing] turn end_span failed (non-fatal): {e}")

tests/lib/adk/test_sync_turn.py

@@ -0,0 +1,124 @@
+"""Tests for the root turn span on the sync OpenAI-Agents path.
+
+The defect these guard: with no root span and no parenting, a trace is a flat
+list of fragments whose durations do not sum to the turn, so tool time simply
+goes missing. These assert a root span exists, that tool spans hang off it, and
+that tracing failures never stop the turn from streaming.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from agentex.lib.adk._modules import _sync_turn as turn_mod
+
+
+def _adk(root_span=None):
+    return SimpleNamespace(
+        tracing=SimpleNamespace(
+            start_span=AsyncMock(return_value=root_span),
+            end_span=AsyncMock(),
+        )
+    )
+
+
+def _runner_yielding(*events):
+    async def _stream():
+        for e in events:
+            yield e
+
+    result = MagicMock()
+    result.stream_events = _stream
+    return MagicMock(return_value=result)
+
+
+def _agent(name="analyst"):
+    agent = MagicMock()
+    agent.name = name
+    return agent
+
+
+async def _drain(gen):
+    return [e async for e in gen]
+
+
+@pytest.mark.asyncio
+async def test_opens_a_root_span_and_closes_it(monkeypatch):
+    root = SimpleNamespace(id="root-1")
+    adk = _adk(root)
+    monkeypatch.setattr(turn_mod, "_get_adk", lambda: adk)
+    monkeypatch.setattr(turn_mod.Runner, "run_streamed", _runner_yielding("a", "b"))
+
+    events = await _drain(turn_mod.run_turn_streamed(_agent(), [], trace_id="tr1", task_id="task1"))
+
+    assert events == ["a", "b"]
+    assert adk.tracing.start_span.await_args.kwargs["name"] == "turn"
+    adk.tracing.end_span.assert_awaited_once()
+    assert adk.tracing.end_span.await_args.kwargs["span"] is root
+
+
+@pytest.mark.asyncio
+async def test_tool_spans_are_parented_to_the_root(monkeypatch):
+    root = SimpleNamespace(id="root-1")
+    monkeypatch.setattr(turn_mod, "_get_adk", lambda: _adk(root))
+
+    captured = {}
+
+    class _Hooks:
+        def __init__(self, **kwargs):
+            captured.update(kwargs)
+
+        async def close_open_tool_spans(self):
+            return None
+
+    monkeypatch.setattr(turn_mod, "SyncTracingHooks", _Hooks)
+    monkeypatch.setattr(turn_mod.Runner, "run_streamed", _runner_yielding("a"))
+
+    await _drain(turn_mod.run_turn_streamed(_agent(), [], trace_id="tr1"))
+
+    assert captured["parent_span_id"] == "root-1"
+
+
+@pytest.mark.asyncio
+async def test_no_trace_id_still_streams(monkeypatch):
+    adk = _adk()
+    monkeypatch.setattr(turn_mod, "_get_adk", lambda: adk)
+    monkeypatch.setattr(turn_mod.Runner, "run_streamed", _runner_yielding("a", "b"))
+
+    events = await _drain(turn_mod.run_turn_streamed(_agent(), []))
+
+    assert events == ["a", "b"]
+    adk.tracing.start_span.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_start_span_failure_still_streams(monkeypatch):
+    adk = _adk()
+    adk.tracing.start_span.side_effect = RuntimeError("tracing backend down")
+    monkeypatch.setattr(turn_mod, "_get_adk", lambda: adk)
+    monkeypatch.setattr(turn_mod.Runner, "run_streamed", _runner_yielding("a"))
+
+    assert await _drain(turn_mod.run_turn_streamed(_agent(), [], trace_id="tr1")) == ["a"]
+
+
+@pytest.mark.asyncio
+async def test_root_span_closes_when_the_run_raises(monkeypatch):
+    root = SimpleNamespace(id="root-1")
+    adk = _adk(root)
+    monkeypatch.setattr(turn_mod, "_get_adk", lambda: adk)
+
+    async def _boom():
+        yield "a"
+        raise RuntimeError("max turns exceeded")
+
+    result = MagicMock()
+    result.stream_events = _boom
+    monkeypatch.setattr(turn_mod.Runner, "run_streamed", MagicMock(return_value=result))
+
+    with pytest.raises(RuntimeError):
+        await _drain(turn_mod.run_turn_streamed(_agent(), [], trace_id="tr1"))
+
+    adk.tracing.end_span.assert_awaited_once()