fix(sdk): keep the workflow span current across a streamed agent run

sdks/python/agenta/sdk/agents/adapters/vercel/sse.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import os
+from contextvars import copy_context
 from json import dumps
 from typing import Any, AsyncGenerator
 
@@ -53,11 +54,27 @@ async def gen():
         # pull times out, we emit a comment, and re-await the SAME pending pull (never dropping or
         # reordering a real part).
         iterator = aiter.__aiter__()
+
+        async def pull():
+            return await iterator.__anext__()
+
+        # INVARIANT: every pull runs in ONE context, captured when this generator is first
+        # driven — the same context the pre-keepalive `async for` ran the upstream in. A task
+        # otherwise starts on a fresh COPY of that context per pull, so whatever the upstream
+        # attaches while producing a chunk (above all the OpenTelemetry activation of the
+        # workflow span, which the instrumentation installs INSIDE the streamed generator's
+        # body) dies with the copy: every later chunk, and the upstream's `finally` where the
+        # run's token/cost usage is stamped, would then run with no workflow span current and
+        # write to a NonRecordingSpan. Pulls are strictly sequential — the next task is created
+        # only after the previous one produced its chunk — so one context is never entered
+        # concurrently.
+        context = copy_context()
+        loop = asyncio.get_running_loop()
         pending: asyncio.Task | None = None
         try:
             while True:
                 if pending is None:
-                    pending = asyncio.ensure_future(iterator.__anext__())
+                    pending = loop.create_task(pull(), context=context)
                 try:
                     chunk = await asyncio.wait_for(
                         asyncio.shield(pending), timeout=interval

sdks/python/agenta/sdk/agents/handler.py

@@ -11,8 +11,11 @@
 
 import os
 from dataclasses import dataclass, field
+from inspect import signature
 from typing import Any, Awaitable, Callable, Dict, List, Optional
 
+from opentelemetry import trace as otel_trace
+
 from agenta.sdk.agents.dtos import AgentTemplate, SessionConfig, to_messages
 from agenta.sdk.agents.interfaces import Backend, Environment
 from agenta.sdk.agents.capabilities import (
@@ -244,6 +247,35 @@ def _agent_model_ref(agent_template: AgentTemplate) -> Optional[ModelRef]:
     return None
 
 
+def _bind_workflow_span(record_usage: RecordUsageFn, span: Any) -> RecordUsageFn:
+    """Pin the run's workflow span onto the usage recorder.
+
+    INVARIANT: usage lands on the span that was current where the run BEGAN, not on whatever is
+    current when the totals are known. The write happens from the run's teardown — after the
+    stream has been driven, possibly by another task holding only a copy of this context — so
+    reading the ambient span there is not sound. A recorder that takes a ``span`` gets the
+    captured reference; a composition-supplied recorder with the plain ``(usage)`` signature
+    keeps working, with the span re-activated around the call instead.
+    """
+    try:
+        accepts_span = "span" in signature(record_usage).parameters
+    except (TypeError, ValueError):  # builtins / C callables expose no signature
+        accepts_span = False
+
+    if accepts_span:
+
+        def _record_with_span(usage: Optional[Dict[str, Any]]) -> None:
+            record_usage(usage, span=span)  # type: ignore[call-arg]
+
+        return _record_with_span
+
+    def _record_under_span(usage: Optional[Dict[str, Any]]) -> None:
+        with otel_trace.use_span(span, end_on_exit=False):
+            record_usage(usage)
+
+    return _record_under_span
+
+
 def make_agent_handler(composition: Optional[AgentComposition] = None):
     """Build the `agent_v0`-shaped handler bound to `composition` (defaults if omitted)."""
 
@@ -261,6 +293,13 @@ async def _agent(
         stream = flags.stream
         session_id = request.session_id
 
+        # Captured HERE, in the handler frame the instrumentation runs with the workflow span
+        # current — the streaming teardown that reports usage no longer can (see
+        # `_bind_workflow_span`).
+        record_usage = _bind_workflow_span(
+            comp.record_usage, otel_trace.get_current_span()
+        )
+
         params = parameters or {}
         agent_template = AgentTemplate.from_params(
             params, defaults=comp.default_template()
@@ -318,14 +357,14 @@ async def _agent(
 
         if stream:
             return agent_event_stream(
-                harness, session_config, msgs, record_usage=comp.record_usage
+                harness, session_config, msgs, record_usage=record_usage
             )
         return await agent_batch(
             harness,
             session_config,
             msgs,
             trim=flags.trim,
-            record_usage=comp.record_usage,
+            record_usage=record_usage,
         )
 
     return _agent

sdks/python/agenta/sdk/agents/tracing.py

@@ -210,18 +210,28 @@ def run_context() -> Optional[RunContext]:
     return RunContext(workflow=workflow, trace=trace)
 
 
-def record_usage(usage: Optional[Dict[str, Any]]) -> None:
-    """Stamp the agent's token/cost totals onto the active ``/invoke`` workflow span.
+def record_usage(
+    usage: Optional[Dict[str, Any]],
+    *,
+    span: Optional[Any] = None,
+) -> None:
+    """Stamp the agent's token/cost totals onto the ``/invoke`` workflow span.
 
     The harness emits its own span tree (turns, LLM, tools) in a separate OTLP batch, so
     Agenta's per-batch cumulative roll-up cannot bridge the totals onto the workflow span.
     Setting ``gen_ai.usage.*`` here records them directly on that span (the root of its
     batch), so the trace shows the run's tokens and cost. Best-effort.
+
+    ``span`` pins the workflow span captured where the run began. Prefer it: the write happens
+    from the run's teardown, arbitrarily far from the frame that made the span current, and any
+    task driving the stream in between carries only a COPY of that context — so the ambient span
+    at write time is not reliably the workflow span, and a write to a non-recording one is
+    silently discarded. Omitting it falls back to the ambient span for standalone callers.
     """
     if not usage or not usage.get("total"):
         return
     try:
-        span = otel_trace.get_current_span()
+        span = span if span is not None else otel_trace.get_current_span()
         input_tokens = int(usage.get("input") or 0)
         output_tokens = int(usage.get("output") or 0)
         span.set_attribute("gen_ai.usage.input_tokens", input_tokens)

sdks/python/oss/tests/pytest/unit/agents/adapters/test_vercel_sse_context.py

@@ -0,0 +1,171 @@
+"""The Vercel SSE framing must not sever the context the stream was entered with.
+
+The framing races each upstream pull against a keepalive tick. Racing it with one asyncio task
+per pull is what breaks tracing: a task runs on a COPY of this generator's context, so anything
+the upstream attaches while producing the first chunk — above all the workflow span that the
+``instrument`` decorator activates INSIDE the streamed generator's body — dies with that copy.
+From the second chunk on, and in the upstream's ``finally``, the current span is then a
+non-recording one and every attribute written there is silently dropped, which is how streaming
+runs stopped recording token/cost usage.
+
+These tests pin the property directly: the workflow span stays current for the whole stream, at
+the DEFAULT keepalive interval and across an interval that actually fires keepalives.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import importlib
+from typing import Any, AsyncIterator, Dict, List
+
+from opentelemetry import context as otel_context
+from opentelemetry import trace as otel_trace
+from opentelemetry.sdk.trace import TracerProvider
+
+import agenta.sdk.agents.adapters.vercel.sse as sse_module
+from agenta.sdk.agents.tracing import record_usage
+
+_USAGE = {"input": 3, "output": 5, "total": 8, "cost": 0.25}
+
+
+def _workflow_span():
+    """A real (recording) SDK span, isolated from whatever global provider the suite installed."""
+    return TracerProvider().get_tracer("agenta.tests").start_span("workflow")
+
+
+def _instrumented_parts(
+    span,
+    seen: List[Any],
+    *,
+    count: int = 4,
+    gap: float = 0.0,
+) -> AsyncIterator[Dict[str, Any]]:
+    """Mirror the SDK instrumentation's streamed-generator wrapper.
+
+    ``instrument`` re-attaches the captured otel context and activates the workflow span from
+    INSIDE the generator body (see ``_wrap_returned_gen``), and the agent handler stamps usage
+    from the generator's ``finally``. Both only work if the framing keeps driving this generator
+    on one context.
+    """
+    captured = otel_context.get_current()
+
+    async def parts() -> AsyncIterator[Dict[str, Any]]:
+        token = otel_context.attach(captured)
+        try:
+            with otel_trace.use_span(span, end_on_exit=False):
+                try:
+                    for index in range(count):
+                        if gap:
+                            await asyncio.sleep(gap)
+                        seen.append(otel_trace.get_current_span())
+                        yield {"type": "text-delta", "delta": str(index)}
+                finally:
+                    seen.append(otel_trace.get_current_span())
+                    record_usage(_USAGE)
+        finally:
+            with contextlib.suppress(Exception):
+                otel_context.detach(token)
+
+    return parts()
+
+
+async def _collect(aiter) -> List[str]:
+    return [chunk async for chunk in aiter]
+
+
+def _assert_span_stayed_current(span, seen: List[Any]) -> None:
+    assert seen, "the upstream never ran"
+    assert all(observed is span for observed in seen), (
+        "the workflow span stopped being current mid-stream: "
+        f"{[getattr(o, 'name', type(o).__name__) for o in seen]}"
+    )
+    assert all(observed.is_recording() for observed in seen)
+
+
+def _assert_usage_landed(span) -> None:
+    attributes = dict(span.attributes or {})
+    assert attributes.get("gen_ai.usage.input_tokens") == 3
+    assert attributes.get("gen_ai.usage.output_tokens") == 5
+    assert attributes.get("gen_ai.usage.total_tokens") == 8
+    assert attributes.get("gen_ai.usage.cost") == 0.25
+
+
+async def test_workflow_span_stays_current_across_the_stream():
+    # Default keepalive interval, no silent gap: no keepalive frame is due, yet the context must
+    # already survive — the regression was in how the pull is driven, not in the keepalive frame.
+    span = _workflow_span()
+    seen: List[Any] = []
+
+    chunks = await _collect(
+        sse_module.vercel_sse_stream(_instrumented_parts(span, seen))
+    )
+
+    assert chunks[-1] == "data: [DONE]\n\n"
+    assert len([c for c in chunks if c.startswith("data: ")]) == 5  # 4 parts + [DONE]
+    _assert_span_stayed_current(span, seen)
+    _assert_usage_landed(span)
+
+
+async def test_workflow_span_survives_keepalive_ticks(monkeypatch):
+    # The same property while keepalives actually fire: every gap times out a pull, and the
+    # resumed pull must land back on the same context.
+    monkeypatch.setenv("AGENTA_AGENT_SSE_KEEPALIVE_SECONDS", "0.02")
+    mod = importlib.reload(sse_module)
+    try:
+        span = _workflow_span()
+        seen: List[Any] = []
+
+        chunks = await _collect(
+            mod.vercel_sse_stream(_instrumented_parts(span, seen, count=3, gap=0.05))
+        )
+    finally:
+        monkeypatch.delenv("AGENTA_AGENT_SSE_KEEPALIVE_SECONDS", raising=False)
+        importlib.reload(sse_module)
+
+    assert [c for c in chunks if c == ": keepalive\n\n"], "no keepalive rode the gaps"
+    payloads = [c for c in chunks if c.startswith("data: ")]
+    assert len(payloads) == 4  # 3 parts + [DONE], none dropped or duplicated
+    _assert_span_stayed_current(span, seen)
+    _assert_usage_landed(span)
+
+
+async def test_disconnect_tears_down_the_in_flight_pull(monkeypatch):
+    # A client that walks away mid-stream must not strand the pull that is still outstanding.
+    monkeypatch.setenv("AGENTA_AGENT_SSE_KEEPALIVE_SECONDS", "0.02")
+    mod = importlib.reload(sse_module)
+    before = asyncio.all_tasks()
+    torn_down = asyncio.Event()
+    try:
+
+        async def parts() -> AsyncIterator[Dict[str, Any]]:
+            yield {"type": "start"}
+            try:
+                await asyncio.sleep(30)  # a part that never arrives
+            except asyncio.CancelledError:
+                torn_down.set()
+                raise
+            yield {"type": "finish"}
+
+        frames: List[str] = []
+        stream = mod.vercel_sse_stream(parts())
+
+        async def consume() -> None:
+            async for frame in stream:
+                frames.append(frame)
+
+        consumer = asyncio.create_task(consume())
+        await asyncio.sleep(0.1)  # first part out, keepalives now riding the silent gap
+        consumer.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await consumer
+        await asyncio.sleep(0.05)
+    finally:
+        monkeypatch.delenv("AGENTA_AGENT_SSE_KEEPALIVE_SECONDS", raising=False)
+        importlib.reload(sse_module)
+
+    assert frames[0] == 'data: {"type": "start"}\n\n'
+    assert ": keepalive\n\n" in frames
+    assert torn_down.is_set(), "the outstanding pull was never cancelled"
+    leaked = [task for task in asyncio.all_tasks() - before if not task.done()]
+    assert not leaked, f"pull left running after disconnect: {leaked}"