From 345667a8d545df2c338a242b5ca24226bd1a459a Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Mon, 3 Aug 2026 23:37:17 +0200 Subject: [PATCH] fix(sdk): stop discarding a cost that arrives without a token count `record_usage` returned as soon as the token total was falsy, and it wrote the cost only after that check. So a run that reports a cost but no trustworthy token split lost the cost entirely. That combination is deliberate upstream, not an edge case. The runner returns exactly `{total: 0, cost: ...}` to say "I have a billed figure but no split I trust", and it has a test pinning that shape. The SDK then threw the whole record away, so the cost lived on the runner's own spans and never reached the workflow root. Trace-focused analytics read root spans only, and the roll-up works per OTLP request, so nothing downstream could recover it. Now the early return fires only when neither a token count nor a cost is present, and the cost is written independently of the token total. Missing and measured zero stay distinct, and the two are drawn differently on purpose. A cost is decided by presence, so a reported 0.0 is stamped as a real measurement, a free model or a fully cached turn, while an absent cost stamps nothing. That matches the API's own convention, which tests for the attribute rather than its value. A token total of zero is read as absence, because it is the runner's own sentinel for "no trustworthy split" and no model call spends zero tokens; writing zeros there would assert a run that consumed nothing and would poison the token roll-up. Inside a reported split, a zero input or output is written as the measurement it is, since a split can honestly be one sided. The body also moved inside the existing error handling. The old code read the token total before the guard, so a malformed record raised into the caller instead of being logged. Tests: 2,396 passing in the SDK suite, and 100 passing in the service's own agent tests, which drive this recorder seam. Six new tests cover a cost-only record, a tokens-only record, records with neither, a reported zero cost against an absent one, a one-sided split, and a malformed record. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA --- sdks/python/agenta/sdk/agents/tracing.py | 65 +++++++-- sdks/python/agenta/sdk/agents/wire_models.py | 8 +- .../agents/test_record_usage_semantics.py | 135 ++++++++++++++++++ 3 files changed, 195 insertions(+), 13 deletions(-) create mode 100644 sdks/python/oss/tests/pytest/unit/agents/test_record_usage_semantics.py diff --git a/sdks/python/agenta/sdk/agents/tracing.py b/sdks/python/agenta/sdk/agents/tracing.py index 004b9b0af9..a5d34e86ad 100644 --- a/sdks/python/agenta/sdk/agents/tracing.py +++ b/sdks/python/agenta/sdk/agents/tracing.py @@ -29,6 +29,14 @@ log = get_module_logger(__name__) +# Declares which token contract a span's ``gen_ai.usage.input_tokens`` follows. ``False`` means +# EXCLUSIVE — uncached input only, with cache reads/writes counted beside it, which is what every +# harness aggregate reports. An ABSENT marker means the OpenTelemetry meaning (input already +# includes cached tokens), and ingest prices the two differently, so any span carrying an input +# count must declare its contract. Mirrors ``INPUT_TOKENS_INCLUDES_CACHE`` in the runner's +# ``services/runner/src/tracing/otel.ts``. +INPUT_TOKENS_INCLUDES_CACHE = "agenta.usage.input_tokens_includes_cache" + _CAPTURE_CONTENT = os.getenv( "AGENTA_AGENT_CONTENT_CAPTURE_ENABLED", "true" ).lower() not in ( @@ -227,20 +235,53 @@ def record_usage( 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. + + Tokens and cost are INDEPENDENT measurements. The runner keeps a cost it was billed even when + it has no trustworthy token split (``mergePromptAndStreamUsage`` in the runner returns + ``{input: 0, output: 0, total: 0, cost}`` for exactly that case), and dropping such a record + would leave the workflow span with nothing: the harness ships its own spans in a separate OTLP + batch, the roll-up runs per batch, and trace-focused analytics read root spans only. + + MISSING vs MEASURED ZERO. A cost is reported whenever the key carries a number, so a reported + ``0.0`` (a free model, a fully cached turn) is stamped as a measurement and a record with no + cost key stamps nothing. This is sound only because the producer can omit the key: the + runner's ``AgentUsage.cost`` is optional and left off when the harness reported no cost. A + token TOTAL of zero is the one value read as absence rather than measurement: it is the + runner's sentinel for "no trustworthy split", and no model call spends zero tokens — so a + zero total writes no token attributes at all, rather than asserting a run that consumed + nothing. Within a reported split, a zero input or output is written as the 0 it is (a split + can honestly be one-sided). + + INPUT TOKEN CONTRACT. Every input token count Agenta ingests must declare whether it already + includes cached tokens, or pricing derives ordinary input by subtracting cache buckets from a + count that never contained them. This record is the harness's aggregate, and both upstream + schemas that produce it report cache counts BESIDE the input count rather than inside it (Pi's + session ``tokens`` carries ``input``/``cacheRead``/``cacheWrite`` as siblings; ACP's ``Usage`` + carries ``inputTokens``/``cachedReadTokens``/``cachedWriteTokens`` as siblings). The aggregate + is therefore exclusive, and this span declares ``False`` alongside the counts — the same + declaration the runner stamps on the leaf span that carries these very numbers. """ - if not usage or not usage.get("total"): - return try: + if not usage: + return + raw_total = usage.get("total") + total_tokens = int(raw_total) if raw_total is not None else 0 + raw_cost = usage.get("cost") + has_tokens = total_tokens > 0 + has_cost = raw_cost is not None + if not has_tokens and not has_cost: + return 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) - span.set_attribute("gen_ai.usage.output_tokens", output_tokens) - span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) - span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) - span.set_attribute("gen_ai.usage.total_tokens", int(usage.get("total") or 0)) - cost = usage.get("cost") - if cost: - span.set_attribute("gen_ai.usage.cost", float(cost)) + if has_tokens: + input_tokens = int(usage.get("input") or 0) + output_tokens = int(usage.get("output") or 0) + span.set_attribute(INPUT_TOKENS_INCLUDES_CACHE, False) + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens) + span.set_attribute("gen_ai.usage.completion_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", total_tokens) + if has_cost: + span.set_attribute("gen_ai.usage.cost", float(raw_cost)) except Exception: # pylint: disable=broad-except log.warning("agent: failed to record usage on workflow span", exc_info=True) diff --git a/sdks/python/agenta/sdk/agents/wire_models.py b/sdks/python/agenta/sdk/agents/wire_models.py index 3c6f8ab27c..eab1c8cf09 100644 --- a/sdks/python/agenta/sdk/agents/wire_models.py +++ b/sdks/python/agenta/sdk/agents/wire_models.py @@ -363,7 +363,13 @@ class WireHarnessCapabilities(_WireModel): class WireAgentUsage(_WireModel): - """Token / cost usage rolled onto a workflow span.""" + """Token / cost usage rolled onto a workflow span. + + ``cost`` is OPTIONAL on the wire, and the distinction is load-bearing: absent means the cost + is UNKNOWN (the harness reported none), while a present ``0`` is a measured zero — a free + model or a fully cached turn. ``record_usage`` reads presence as evidence of a measurement, + so an absent cost must never be normalized to a zero anywhere on this path. + """ input: Optional[int] = None output: Optional[int] = None diff --git a/sdks/python/oss/tests/pytest/unit/agents/test_record_usage_semantics.py b/sdks/python/oss/tests/pytest/unit/agents/test_record_usage_semantics.py new file mode 100644 index 0000000000..41fd6fa785 --- /dev/null +++ b/sdks/python/oss/tests/pytest/unit/agents/test_record_usage_semantics.py @@ -0,0 +1,135 @@ +"""Which usage records reach the workflow span, and how a zero is read. + +``record_usage`` stamps the agent run's usage on the ``/invoke`` workflow span. Tokens and cost +are independent measurements there: the runner keeps a billed cost even when it has no +trustworthy token split, and that cost has nowhere else to land — the harness's own spans ship in +a separate OTLP batch, the cumulative roll-up runs per batch, and trace-focused analytics read +root spans only. A record carrying only a cost must therefore still be stamped. + +The other half is telling absence from a measurement. A cost is reported whenever its key carries +a number, so ``0.0`` is stamped as the measurement it is; a token total of zero is the runner's +sentinel for "no trustworthy split" and writes no token attributes at all. That reading is sound +only because the producer can express absence: the runner's ``AgentUsage.cost`` is optional and +omitted when the harness reported no cost. + +A token count also has to declare its contract. The harness aggregate counts uncached input only, +so the span says so with ``agenta.usage.input_tokens_includes_cache = False`` — without it, ingest +reads the OpenTelemetry (cache-inclusive) meaning and misprices the run. +""" + +from __future__ import annotations + +from typing import Any, Dict + +from opentelemetry.sdk.trace import TracerProvider + +from agenta.sdk.agents.tracing import INPUT_TOKENS_INCLUDES_CACHE, record_usage + + +def _workflow_span(): + return TracerProvider().get_tracer("agenta.tests").start_span("workflow") + + +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.") + } + + +def test_a_cost_without_a_token_split_is_still_stamped(): + # The runner's own shape for "billed, but no trustworthy token count". + span = _workflow_span() + + record_usage({"input": 0, "output": 0, "total": 0, "cost": 0.04}, span=span) + + assert _usage_attributes(span) == {"gen_ai.usage.cost": 0.04} + + +def test_a_token_split_without_a_cost_is_still_stamped(): + span = _workflow_span() + + record_usage({"input": 3, "output": 5, "total": 8}, span=span) + + 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, + } + + +def test_a_record_with_neither_tokens_nor_cost_is_skipped(): + for empty in (None, {}, {"input": 0, "output": 0, "total": 0}): + span = _workflow_span() + + record_usage(empty, span=span) + + assert _usage_attributes(span) == {} + + +def test_a_reported_zero_cost_is_a_measurement_but_an_absent_one_is_not(): + reported = _workflow_span() + absent = _workflow_span() + + record_usage({"input": 3, "output": 5, "total": 8, "cost": 0}, span=reported) + record_usage({"input": 3, "output": 5, "total": 8}, span=absent) + + # A free model or a fully cached turn really does cost 0.0; a record with no cost key + # measured nothing, and must not be reported as a zero-cost run. + assert _usage_attributes(reported)["gen_ai.usage.cost"] == 0.0 + assert "gen_ai.usage.cost" not in _usage_attributes(absent) + + +def test_a_one_sided_split_keeps_its_measured_zero(): + span = _workflow_span() + + record_usage({"input": 0, "output": 5, "total": 5, "cost": 0.01}, span=span) + + attributes = _usage_attributes(span) + assert attributes["gen_ai.usage.input_tokens"] == 0 + assert attributes["gen_ai.usage.total_tokens"] == 5 + + +def test_a_run_the_harness_never_priced_stamps_no_cost_at_all(): + # The runner omits `cost` when the harness reported none (codex reports a token split and + # no cost). Presence-means-measured only holds if such a record stamps nothing: a zero here + # would tell every downstream aggregate the run was free. + span = _workflow_span() + + record_usage({"input": 12, "output": 3, "total": 15}, span=span) + + assert "gen_ai.usage.cost" not in _usage_attributes(span) + assert _usage_attributes(span)["gen_ai.usage.total_tokens"] == 15 + + +def test_a_token_count_declares_that_it_excludes_cached_input(): + # The harness aggregate counts uncached input only (Pi and ACP both report cache reads and + # writes as siblings of the input count). Ingest assumes the OpenTelemetry cache-INCLUSIVE + # contract when the marker is absent, so the count has to declare itself. + span = _workflow_span() + + record_usage({"input": 3, "output": 5, "total": 8, "cost": 0.01}, span=span) + + assert dict(span.attributes or {})[INPUT_TOKENS_INCLUDES_CACHE] is False + + +def test_a_cost_only_record_declares_nothing_about_token_contracts(): + # No input token count means no contract to declare; a stray marker would describe a count + # this span never made. + span = _workflow_span() + + record_usage({"input": 0, "output": 0, "total": 0, "cost": 0.04}, span=span) + + assert INPUT_TOKENS_INCLUDES_CACHE not in dict(span.attributes or {}) + + +def test_a_malformed_record_never_raises_into_the_caller(): + span = _workflow_span() + + record_usage({"total": "not-a-number", "cost": 0.04}, span=span) + record_usage(["not", "a", "mapping"], span=span) # type: ignore[arg-type] + + assert _usage_attributes(span) == {}