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