From 10cd6e185bcf249062f6664c49eb3bb1217453ea Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Mon, 3 Aug 2026 22:37:42 +0200 Subject: [PATCH] fix(sdk): keep the workflow span current across a streamed agent run Since commit 8e99171aff (2026-07-07) an agent run on the streaming path has recorded no tokens and no cost on its workflow span. Measured: zero out of roughly 1,300 streaming runs across two independent stacks since 2026-07-08. Batch runs were never affected, which is why the bug stayed invisible. That commit added an SSE keepalive so a proxy would not drop a connection during a silent gap, such as a tool call that runs for minutes. To race a keepalive timer against each chunk, it began driving the response generator with one asyncio task per pull. Every task gets its own copy of the context. The instrumentation attaches the OpenTelemetry context and activates the workflow span from inside the generator body, so that activation landed in the first pull's private copy and every later pull started from a copy that had never seen it. Chunk one recorded. Everything after it, including the `finally` block where usage is written, ran against a non-recording span and was silently discarded. Two changes. The stream now captures one context when the generator is first driven and creates every pull task with it, so all pulls share a single context. That restores the property the old plain iteration had, keeps the existing lock-step backpressure with no read-ahead, and leaves cancellation and cleanup alone. It also removes the "Token created in a different Context" detach errors visible in the pre-fix logs. `record_usage` no longer depends on the ambient context. The handler captures the workflow span at entry and passes that reference, so this class of bug cannot silence usage again. The existing early return on a falsy total and the truthy-cost guard are unchanged, and an override that takes only the usage argument still works. Tests: 1,905 passing through the canonical runner, and 100 passing in the service's own agent tests, which confirms the recorder seam stays backward compatible. The two new test files fail against the unfixed code and pass against the fix. The stream test asserts the same recording span is current on every chunk, across a keepalive gap, and that no frame is dropped or duplicated. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA --- .../agenta/sdk/agents/adapters/vercel/sse.py | 19 +- sdks/python/agenta/sdk/agents/handler.py | 43 ++++- sdks/python/agenta/sdk/agents/tracing.py | 16 +- .../adapters/test_vercel_sse_context.py | 171 ++++++++++++++++++ .../unit/agents/test_usage_span_binding.py | 147 +++++++++++++++ 5 files changed, 390 insertions(+), 6 deletions(-) create mode 100644 sdks/python/oss/tests/pytest/unit/agents/adapters/test_vercel_sse_context.py create mode 100644 sdks/python/oss/tests/pytest/unit/agents/test_usage_span_binding.py diff --git a/sdks/python/agenta/sdk/agents/adapters/vercel/sse.py b/sdks/python/agenta/sdk/agents/adapters/vercel/sse.py index 7570d7d065..c161fd7161 100644 --- a/sdks/python/agenta/sdk/agents/adapters/vercel/sse.py +++ b/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 diff --git a/sdks/python/agenta/sdk/agents/handler.py b/sdks/python/agenta/sdk/agents/handler.py index 044bb77e4d..99726aac8d 100644 --- a/sdks/python/agenta/sdk/agents/handler.py +++ b/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 diff --git a/sdks/python/agenta/sdk/agents/tracing.py b/sdks/python/agenta/sdk/agents/tracing.py index 66a9b8e8c7..004b9b0af9 100644 --- a/sdks/python/agenta/sdk/agents/tracing.py +++ b/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) diff --git a/sdks/python/oss/tests/pytest/unit/agents/adapters/test_vercel_sse_context.py b/sdks/python/oss/tests/pytest/unit/agents/adapters/test_vercel_sse_context.py new file mode 100644 index 0000000000..707a63943e --- /dev/null +++ b/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}" diff --git a/sdks/python/oss/tests/pytest/unit/agents/test_usage_span_binding.py b/sdks/python/oss/tests/pytest/unit/agents/test_usage_span_binding.py new file mode 100644 index 0000000000..9b0e3669f5 --- /dev/null +++ b/sdks/python/oss/tests/pytest/unit/agents/test_usage_span_binding.py @@ -0,0 +1,147 @@ +"""Usage lands on the workflow span even when the run's ambient context is gone. + +``record_usage`` runs from the run's teardown: on the streaming path that is the ``finally`` of +a generator the server drives long after the handler frame returned, and whatever drives it may +carry only a copy of the context that made the workflow span current. Reading the ambient span +there is therefore not sound — a lost activation turns every usage write into a no-op on a +non-recording span, with no error anywhere. The handler pins the span it captured at run start +instead, which these tests hold it to by draining the stream under an adverse context. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Dict, List, Optional + +from opentelemetry import trace as otel_trace +from opentelemetry.sdk.trace import TracerProvider + +from agenta.sdk.agents import AgentResult +from agenta.sdk.agents.handler import AgentComposition, make_agent_handler +from agenta.sdk.agents.tracing import record_usage +from agenta.sdk.models.workflows import WorkflowServiceRequest + +_USAGE = {"input": 3, "output": 5, "total": 8, "cost": 0.25} + + +def _workflow_span(): + return TracerProvider().get_tracer("agenta.tests").start_span("workflow") + + +def _params() -> Dict[str, Any]: + return {"agent": {"harness": {"kind": "pi_core"}}} + + +def _messages() -> List[Dict[str, Any]]: + return [{"role": "user", "content": "hi"}] + + +async def _drain_one_task_per_pull(stream) -> List[Any]: + """Drain the way the SSE framing does when it races a keepalive: a task per pull, each on + its own copy of this context.""" + iterator = stream.__aiter__() + out: List[Any] = [] + while True: + pull = asyncio.ensure_future(iterator.__anext__()) + try: + out.append(await pull) + except StopAsyncIteration: + break + return out + + +def _usage_attributes(span) -> Dict[str, Any]: + return { + key: value + for key, value in dict(span.attributes or {}).items() + if key.startswith("gen_ai.usage.") + } + + +async def test_streaming_usage_lands_on_the_span_captured_at_run_start(make_backend): + backend = make_backend(result=AgentResult(output="ok", usage=dict(_USAGE))) + handler = make_agent_handler( + AgentComposition(select_backend=lambda template: backend) + ) + span = _workflow_span() + + # The span is current only for the handler call — exactly the instrumentation's shape. + with otel_trace.use_span(span, end_on_exit=False): + stream = await handler( + request=WorkflowServiceRequest(flags={"stream": True}), + messages=_messages(), + parameters=_params(), + ) + + assert not otel_trace.get_current_span().is_recording() + await _drain_one_task_per_pull(stream) + + assert _usage_attributes(span) == { + "gen_ai.usage.input_tokens": 3, + "gen_ai.usage.output_tokens": 5, + "gen_ai.usage.prompt_tokens": 3, + "gen_ai.usage.completion_tokens": 5, + "gen_ai.usage.total_tokens": 8, + "gen_ai.usage.cost": 0.25, + } + + +async def test_batch_usage_lands_on_the_span_captured_at_run_start(make_backend): + backend = make_backend(result=AgentResult(output="ok", usage=dict(_USAGE))) + handler = make_agent_handler( + AgentComposition(select_backend=lambda template: backend) + ) + span = _workflow_span() + + with otel_trace.use_span(span, end_on_exit=False): + await handler( + request=WorkflowServiceRequest(), + messages=_messages(), + parameters=_params(), + ) + + assert _usage_attributes(span)["gen_ai.usage.total_tokens"] == 8 + + +async def test_composition_recorder_without_a_span_parameter_still_sees_it( + make_backend, +): + # The composition seam predates the span argument: a `(usage)`-only recorder must keep + # working AND must still read the workflow span from the ambient context. + seen: List[Any] = [] + + def legacy_recorder(usage: Optional[Dict[str, Any]]) -> None: + seen.append((usage, otel_trace.get_current_span())) + + backend = make_backend(result=AgentResult(output="ok", usage=dict(_USAGE))) + handler = make_agent_handler( + AgentComposition( + select_backend=lambda template: backend, + record_usage=legacy_recorder, + ) + ) + span = _workflow_span() + + with otel_trace.use_span(span, end_on_exit=False): + stream = await handler( + request=WorkflowServiceRequest(flags={"stream": True}), + messages=_messages(), + parameters=_params(), + ) + await _drain_one_task_per_pull(stream) + + assert len(seen) == 1 + usage, observed_span = seen[0] + assert usage == _USAGE + assert observed_span is span + + +def test_record_usage_stamps_the_given_span_over_the_ambient_one(): + span = _workflow_span() + + # No span is current here: the ambient read would write to a NonRecordingSpan and vanish. + assert not otel_trace.get_current_span().is_recording() + record_usage(_USAGE, span=span) + + assert _usage_attributes(span)["gen_ai.usage.total_tokens"] == 8 + assert _usage_attributes(span)["gen_ai.usage.cost"] == 0.25