From 22da869b877b6e5fa8688f224690397ca8cb064e Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Mon, 3 Aug 2026 22:31:26 +0200 Subject: [PATCH] fix(api): make an agent run's cost and tokens reach the trace An agent trace shows no cost. `ag.metrics.costs.cumulative.*` was present on 0 of 48,005 spans measured on a real stack, root spans and child spans alike, while the playground showed a correct cost for the same run. The playground reads the number the harness streams to the browser. The trace reads a number the platform recomputes, and the recompute produces nothing for an agent run. Three causes, fixed here. The platform never accepted a reported cost, it only recomputed one. The runner already stamps the harness's own figure as `gen_ai.usage.cost`, but no ingest adapter mapped it into the `ag.*` namespace, so it sat at a raw path no product surface reads. Map it to `ag.metrics.costs.cumulative.total`. Cumulative, not incremental: the value is the run's aggregate total, so calling it incremental would double count against the runner's own model-call spans. The roll-up then overwrote it. Measured on the real pipeline: a span reporting $0.42 with one priceable child came out at $0.00021. The two numbers describe the same spend, not additive spend, and the recomputed one is the lossier estimate because it prices cached tokens at the full uncached rate. So the roll-up now fills only spans that report nothing. The roll-up did not run at all on the runner's telemetry. One run arrives in two OTLP requests: the SDK sends the workflow root span, and the runner sends its own subtree headed by a span whose parent lives in the other request. The tree builder seeded only from spans with no parent, so the runner's request produced an empty tree and nothing was cumulated. Measured: 0 of 8,719 child agent spans carried `ag.metrics.tokens.cumulative`. A span whose parent is absent from the batch now seeds a tree, because a dangling parent id means "not in this batch" rather than "no parent". Roll-up stays batch local and still does not cross requests. Tests: 1,608 passing in OSS and 1,867 in EE, including six new tests that drive the real ingest sequence rather than mocking it, and four covering the widened seeding rule against disconnected subtrees and a parent cycle. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA --- .../extractors/adapters/logfire_adapter.py | 5 + api/oss/src/core/tracing/utils/trees.py | 52 +++- .../pytest/unit/otlp/test_logfire_adapter.py | 11 + .../unit/otlp/test_reported_cost_ingest.py | 251 ++++++++++++++++++ .../pytest/unit/tracing/utils/test_trees.py | 161 +++++++++++ 5 files changed, 468 insertions(+), 12 deletions(-) create mode 100644 api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py diff --git a/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py b/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py index 32ec482ea2..435ddfc899 100644 --- a/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py +++ b/api/oss/src/apis/fastapi/otlp/extractors/adapters/logfire_adapter.py @@ -193,6 +193,11 @@ def _split_agent_messages( "gen_ai.usage.cache_creation.input_tokens", "ag.metrics.unit.tokens.cache_creation", ), + # Producer-reported cost. This is the aggregate total for the span's whole run + # (see the agent SDK's record_usage), so it maps to `cumulative`, not + # `incremental`: as an incremental value the tree roll-up would add it on top of + # the run's own model-call spans and double count. + ("gen_ai.usage.cost", "ag.metrics.costs.cumulative.total"), ] OPERATION_TO_NODETYPE = { diff --git a/api/oss/src/core/tracing/utils/trees.py b/api/oss/src/core/tracing/utils/trees.py index adf7181524..fbf93c4f6f 100644 --- a/api/oss/src/core/tracing/utils/trees.py +++ b/api/oss/src/core/tracing/utils/trees.py @@ -31,8 +31,8 @@ def calculate_and_propagate_metrics( """ Calculate and propagate costs/tokens/errors for a list of span DTOs. - This must be called BEFORE batching to ensure complete trace trees. - If called after batching, partial traces will fail to propagate correctly. + Roll-up is batch-local: a span only ever sums the children present in this call, + so totals are complete only when the whole trace is passed in. Args: span_dtos: List of span DTOs (should be from a complete trace) @@ -175,19 +175,39 @@ def promote_identity_by_trace( def parse_span_idx_to_span_id_tree( span_idx: Dict[str, OTelFlatSpan], ) -> OrderedDict: - span_id_tree = OrderedDict() - index = {} + """ + Build the forest of span trees for one batch of spans. + + A span whose parent is missing from the batch is a root here: a trace is split + across OTLP requests (the agent runner ships its own subtree, headed by a span + whose parent lives in the SDK's request), so a dangling parent id means "not in + this batch", not "no parent". Without this, such a batch yields an empty tree and + nothing is cumulated. Roll-up stays batch-local; it never crosses requests. - def push(span_dto: OTelFlatSpan) -> None: - if span_dto.parent_id is None: - span_id_tree[span_dto.span_id] = OrderedDict() - index[span_dto.span_id] = span_id_tree[span_dto.span_id] - elif span_dto.parent_id in index: - index[span_dto.parent_id][span_dto.span_id] = OrderedDict() - index[span_dto.span_id] = index[span_dto.parent_id][span_dto.span_id] + Every span has at most one parent and every parent is expanded at most once, so + each span appears at most once in the forest. Spans in a parent cycle are reachable + from no root and are simply left out. + """ + span_id_tree = OrderedDict() + children_by_parent_id: Dict[str, List[OTelFlatSpan]] = {} + roots: List[OTelFlatSpan] = [] for span_dto in sorted(span_idx.values(), key=lambda span_dto: span_dto.start_time): - push(span_dto) + if span_dto.parent_id is None or span_dto.parent_id not in span_idx: + roots.append(span_dto) + else: + children_by_parent_id.setdefault(span_dto.parent_id, []).append(span_dto) + + stack = [(span_dto, span_id_tree) for span_dto in reversed(roots)] + + while stack: + span_dto, siblings = stack.pop() + + children = OrderedDict() + siblings[span_dto.span_id] = children + + for child_span_dto in reversed(children_by_parent_id.get(span_dto.span_id, [])): + stack.append((child_span_dto, children)) return span_id_tree @@ -315,6 +335,14 @@ def _set_cumulative(span: OTelFlatSpan, costs: dict): if span.attributes is None: span.attributes = {} + # A cumulative total already on the span was reported by the producer (an + # agent harness's gen_ai.usage.cost, mapped at ingest) and is the billed + # aggregate for this subtree. Child costs recomputed from token counts + # re-estimate the same spend, so overwriting here would swap a billed figure + # for a lossier one. The roll-up only fills spans that report nothing. + if _get_cumulative(span).get("total", 0.0) != 0.0: + return + if ( costs.get("prompt", 0.0) != 0.0 or costs.get("completion", 0.0) != 0.0 diff --git a/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py b/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py index 90ac86a5c8..7dcc55c662 100644 --- a/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py +++ b/api/oss/tests/pytest/unit/otlp/test_logfire_adapter.py @@ -708,6 +708,17 @@ def test_cache_read_tokens_mapped(self, adapter): assert features.metrics["unit.tokens.cache_read"] == 2304 + def test_reported_cost_mapped_to_cumulative_not_incremental(self, adapter): + # gen_ai.usage.cost is the aggregate for the whole run, so it lands on + # `cumulative`; as an `incremental` value the tree roll-up would add it on top + # of the run's own model-call spans. + bag = _make_bag({"gen_ai.usage.cost": 0.4237}) + features = SpanFeatures() + adapter.process(bag, features) + + assert features.metrics["costs.cumulative.total"] == 0.4237 + assert "costs.incremental.total" not in features.metrics + def test_provider_name_mapped(self, adapter): bag = _make_bag({"gen_ai.provider.name": "openai"}) features = SpanFeatures() diff --git a/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py b/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py new file mode 100644 index 0000000000..549ccd012d --- /dev/null +++ b/api/oss/tests/pytest/unit/otlp/test_reported_cost_ingest.py @@ -0,0 +1,251 @@ +"""Ingest-path tests for a producer-reported agent run cost. + +An agent harness reports what the run actually cost as `gen_ai.usage.cost`. The +platform never priced agent runs itself (its litellm recompute is keyed on +`ag.meta.response.model`, which agent spans do not carry), so the trace showed no +cost at all. These tests exercise the real ingest sequence — adapters first, then +the tree roll-up — and pin how the two interact: + + router.otlp_ingest -> parse_from_otel_span_dto (adapters) [router.py:189] + router.otlp_ingest -> TracingService.ingest_span_dtos [router.py:256] + -> calculate_and_propagate_metrics_by_trace [service.py:147] +""" + +from datetime import datetime, timedelta, timezone + +import pytest + +from oss.src.apis.fastapi.otlp.utils.processing import parse_from_otel_span_dto +from oss.src.core.otel.dtos import ( + OTelContextDTO, + OTelSpanDTO, + OTelSpanKind, + OTelStatusCode, +) +from oss.src.core.tracing.utils.trees import calculate_and_propagate_metrics_by_trace + + +TRACE_ID = "a" * 32 +ROOT_SPAN_ID = "b" * 16 +AGENT_SPAN_ID = "c" * 16 +MODEL_SPAN_ID = "d" * 16 + +START = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def _otel_span( + *, + span_id: str, + span_name: str, + attributes: dict, + parent_id: str = None, + start_offset_s: int = 0, +) -> OTelSpanDTO: + start = START + timedelta(seconds=start_offset_s) + + return OTelSpanDTO( + context=OTelContextDTO(trace_id=f"0x{TRACE_ID}", span_id=f"0x{span_id}"), + parent=( + OTelContextDTO(trace_id=f"0x{TRACE_ID}", span_id=f"0x{parent_id}") + if parent_id + else None + ), + name=span_name, + kind=OTelSpanKind.SPAN_KIND_INTERNAL, + start_time=start, + end_time=start + timedelta(seconds=1), + status_code=OTelStatusCode.STATUS_CODE_OK, + attributes=attributes, + ) + + +def _ingest(otel_spans): + """Run the ingest sequence: adapters + builder, then the metric roll-up.""" + span_dtos = [parse_from_otel_span_dto(otel_span) for otel_span in otel_spans] + + return { + span_dto.span_name: span_dto + for span_dto in calculate_and_propagate_metrics_by_trace(span_dtos) + } + + +def _cumulative_costs(span_dto) -> dict: + return ( + (span_dto.attributes or {}) + .get("ag", {}) + .get("metrics", {}) + .get("costs", {}) + .get("cumulative", {}) + ) + + +def _agent_span(**overrides) -> OTelSpanDTO: + attributes = { + "gen_ai.operation.name": "invoke_agent", + "gen_ai.usage.cost": 0.4237, + } + attributes.update(overrides.pop("attributes", {})) + + return _otel_span( + span_id=AGENT_SPAN_ID, + span_name="invoke_agent", + attributes=attributes, + **overrides, + ) + + +def _unpriceable_model_span(**overrides) -> OTelSpanDTO: + # The agent runner reports the request model but not the response model, which is + # the key the litellm recompute needs. This is what every agent chat span looks + # like today, and why the recompute produces nothing. + return _otel_span( + span_id=MODEL_SPAN_ID, + span_name="chat", + attributes={ + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gpt-4o-mini", + "gen_ai.usage.input_tokens": 1000, + "gen_ai.usage.output_tokens": 100, + }, + **overrides, + ) + + +def _priceable_model_span(**overrides) -> OTelSpanDTO: + return _otel_span( + span_id=MODEL_SPAN_ID, + span_name="chat", + attributes={ + "gen_ai.operation.name": "chat", + "gen_ai.request.model": "gpt-4o-mini", + "gen_ai.response.model": "gpt-4o-mini", + "gen_ai.usage.input_tokens": 1000, + "gen_ai.usage.output_tokens": 100, + }, + **overrides, + ) + + +@pytest.fixture +def fixed_token_pricing(monkeypatch): + """Price every model span at a flat, tiny amount, independent of litellm's catalog.""" + + def _cost_per_token(*, model, prompt_tokens, completion_tokens, **_kwargs): + if not model: + raise ValueError("model is required") + return (0.00015, 0.00006) + + monkeypatch.setattr( + "oss.src.core.tracing.utils.trees.cost_calculator.cost_per_token", + _cost_per_token, + ) + + +def test_reported_cost_becomes_cumulative_total_after_ingest(): + span_idx = _ingest([_agent_span()]) + + agent_span = span_idx["invoke_agent"] + + assert _cumulative_costs(agent_span)["total"] == 0.4237 + + +def test_reported_cost_survives_rollup_when_no_model_span_is_priceable( + fixed_token_pricing, +): + # The roll-up writes `cumulative` only when it computes a non-zero total, so an + # agent run whose model spans price to nothing must keep the reported figure. + span_idx = _ingest( + [ + _agent_span(), + _unpriceable_model_span(parent_id=AGENT_SPAN_ID, start_offset_s=1), + ] + ) + + agent_span = span_idx["invoke_agent"] + model_span = span_idx["chat"] + + assert _cumulative_costs(agent_span)["total"] == 0.4237 + assert _cumulative_costs(model_span) == {} + + +def test_reported_cost_wins_over_recomputed_child_costs(fixed_token_pricing): + # Both figures describe the same spend. The reported one is what the harness was + # billed; the recomputed one re-derives it from token counts and undercounts + # (cached prompt tokens are priced as fresh ones). Summing children over the + # reporting span would replace the billed figure with the lossier estimate. + span_idx = _ingest( + [ + _agent_span(), + _priceable_model_span(parent_id=AGENT_SPAN_ID, start_offset_s=1), + ] + ) + + agent_span = span_idx["invoke_agent"] + model_span = span_idx["chat"] + + assert _cumulative_costs(agent_span)["total"] == 0.4237 + assert _cumulative_costs(model_span)["total"] == pytest.approx(0.00021) + + +def test_reported_cost_propagates_to_an_ancestor_that_reports_nothing( + fixed_token_pricing, +): + # The SDK's workflow root does not always report a cost; it must still show the + # agent subtree's, which is what the trace list and playground read. + root = _otel_span( + span_id=ROOT_SPAN_ID, + span_name="workflow", + attributes={"ag.type.node": "workflow"}, + ) + span_idx = _ingest( + [ + root, + _agent_span(parent_id=ROOT_SPAN_ID, start_offset_s=1), + _unpriceable_model_span(parent_id=AGENT_SPAN_ID, start_offset_s=2), + ] + ) + + root_span = span_idx["workflow"] + agent_span = span_idx["invoke_agent"] + + assert _cumulative_costs(root_span)["total"] == 0.4237 + assert _cumulative_costs(agent_span)["total"] == 0.4237 + + +def test_reported_cost_is_not_double_counted_when_parent_and_child_both_report(): + # The SDK stamps the run total on the workflow root and the runner stamps the same + # total on its own agent span. The reported value lands on `cumulative`, never on + # `incremental`, so the roll-up has nothing to add it to. + span_idx = _ingest( + [ + _otel_span( + span_id=ROOT_SPAN_ID, + span_name="workflow", + attributes={"gen_ai.usage.cost": 0.4237}, + ), + _agent_span(parent_id=ROOT_SPAN_ID, start_offset_s=1), + ] + ) + + root_span = span_idx["workflow"] + + assert _cumulative_costs(root_span)["total"] == 0.4237 + + +def test_rollup_still_owns_cost_for_traces_that_report_nothing(fixed_token_pricing): + # Guard on the change to the roll-up: spans that carry no reported cumulative must + # keep summing their children exactly as before. + span_idx = _ingest( + [ + _otel_span( + span_id=ROOT_SPAN_ID, + span_name="workflow", + attributes={"ag.type.node": "workflow"}, + ), + _priceable_model_span(parent_id=ROOT_SPAN_ID, start_offset_s=1), + ] + ) + + root_span = span_idx["workflow"] + + assert _cumulative_costs(root_span)["total"] == pytest.approx(0.00021) diff --git a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py index be94c669da..f413c8a0e4 100644 --- a/api/oss/tests/pytest/unit/tracing/utils/test_trees.py +++ b/api/oss/tests/pytest/unit/tracing/utils/test_trees.py @@ -28,6 +28,11 @@ ROOT_UUID = "31d6cfe0-4b90-11ec-31d6-cfe04b9011ec" CHILD_A_UUID = "41d6cfe0-4b90-11ec-41d6-cfe04b9011ec" CHILD_B_UUID = "51d6cfe0-4b90-11ec-51d6-cfe04b9011ec" +ABSENT_PARENT_UUID = "61d6cfe0-4b90-11ec-61d6-cfe04b9011ec" +ORPHAN_UUID = "71d6cfe0-4b90-11ec-71d6-cfe04b9011ec" +ORPHAN_CHILD_UUID = "81d6cfe0-4b90-11ec-81d6-cfe04b9011ec" +CYCLE_A_UUID = "91d6cfe0-4b90-11ec-91d6-cfe04b9011ec" +CYCLE_B_UUID = "a1d6cfe0-4b90-11ec-a1d6-cfe04b9011ec" def _span( @@ -114,6 +119,162 @@ def test_parse_span_dtos_to_span_idx_and_tree_hierarchy(): assert list(tree[ROOT_UUID].keys()) == [CHILD_A_UUID] +def test_parentless_root_still_seeds_its_own_tree_only(): + # A span with parent_id=None keeps seeding exactly as before; widening the rule to + # dangling parents must not turn its children into extra roots. + root = _span(span_id=ROOT_UUID, span_name="root", start_offset_s=0) + child = _span( + span_id=CHILD_A_UUID, + parent_id=ROOT_UUID, + span_name="child", + start_offset_s=1, + ) + grandchild = _span( + span_id=CHILD_B_UUID, + parent_id=CHILD_A_UUID, + span_name="grandchild", + start_offset_s=2, + ) + + span_idx = parse_span_dtos_to_span_idx([grandchild, child, root]) + tree = parse_span_idx_to_span_id_tree(span_idx) + + assert list(tree.keys()) == [ROOT_UUID] + assert list(tree[ROOT_UUID].keys()) == [CHILD_A_UUID] + assert list(tree[ROOT_UUID][CHILD_A_UUID].keys()) == [CHILD_B_UUID] + + +def test_span_with_parent_absent_from_batch_seeds_and_cumulates(): + # The agent runner ships its subtree in its own OTLP request; the top span points at + # a parent that arrived in the SDK's request. That span must still root a tree. + invoke_agent = _span( + span_id=ORPHAN_UUID, + parent_id=ABSENT_PARENT_UUID, + span_name="invoke_agent", + errors=1, + start_offset_s=0, + ) + llm_call = _span( + span_id=ORPHAN_CHILD_UUID, + parent_id=ORPHAN_UUID, + span_name="llm_call", + prompt_tokens=10, + completion_tokens=20, + prompt_cost=0.1, + completion_cost=0.2, + errors=2, + start_offset_s=1, + ) + + span_idx = parse_span_dtos_to_span_idx([invoke_agent, llm_call]) + tree = parse_span_idx_to_span_id_tree(span_idx) + + assert list(tree.keys()) == [ORPHAN_UUID] + assert list(tree[ORPHAN_UUID].keys()) == [ORPHAN_CHILD_UUID] + + cumulate_tokens(tree, span_idx) + cumulate_costs(tree, span_idx) + cumulate_errors(tree, span_idx) + + metrics = span_idx[ORPHAN_UUID].attributes["ag"]["metrics"] + + assert metrics["tokens"]["cumulative"] == { + "prompt": 10.0, + "completion": 20.0, + "total": 30.0, + } + assert metrics["costs"]["cumulative"]["total"] == pytest.approx(0.3) + assert metrics["errors"]["cumulative"] == 3 + + +def test_disconnected_subtrees_cumulate_independently(): + root = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=1, + start_offset_s=0, + ) + root_child = _span( + span_id=CHILD_A_UUID, + parent_id=ROOT_UUID, + span_name="root-child", + prompt_tokens=2, + start_offset_s=1, + ) + orphan = _span( + span_id=ORPHAN_UUID, + parent_id=ABSENT_PARENT_UUID, + span_name="orphan", + prompt_tokens=4, + start_offset_s=2, + ) + orphan_child = _span( + span_id=ORPHAN_CHILD_UUID, + parent_id=ORPHAN_UUID, + span_name="orphan-child", + prompt_tokens=8, + start_offset_s=3, + ) + + span_idx = parse_span_dtos_to_span_idx([root, root_child, orphan, orphan_child]) + tree = parse_span_idx_to_span_id_tree(span_idx) + + assert list(tree.keys()) == [ROOT_UUID, ORPHAN_UUID] + + cumulate_tokens(tree, span_idx) + + def _cumulative_prompt(span_id: str) -> float: + return span_idx[span_id].attributes["ag"]["metrics"]["tokens"]["cumulative"][ + "prompt" + ] + + assert _cumulative_prompt(ROOT_UUID) == 3.0 + assert _cumulative_prompt(ORPHAN_UUID) == 12.0 + assert _cumulative_prompt(CHILD_A_UUID) == 2.0 + assert _cumulative_prompt(ORPHAN_CHILD_UUID) == 8.0 + + +def test_parent_cycle_terminates_and_is_left_out_of_the_forest(): + cycle_a = _span( + span_id=CYCLE_A_UUID, + parent_id=CYCLE_B_UUID, + span_name="cycle-a", + prompt_tokens=1, + start_offset_s=0, + ) + cycle_b = _span( + span_id=CYCLE_B_UUID, + parent_id=CYCLE_A_UUID, + span_name="cycle-b", + prompt_tokens=2, + start_offset_s=1, + ) + root = _span( + span_id=ROOT_UUID, + span_name="root", + prompt_tokens=4, + start_offset_s=2, + ) + + span_idx = parse_span_dtos_to_span_idx([cycle_a, cycle_b, root]) + tree = parse_span_idx_to_span_id_tree(span_idx) + + assert list(tree.keys()) == [ROOT_UUID] + assert tree[ROOT_UUID] == {} + + cumulate_tokens(tree, span_idx) + + assert span_idx[ROOT_UUID].attributes["ag"]["metrics"]["tokens"]["cumulative"][ + "prompt" + ] == pytest.approx(4.0) + assert ( + "cumulative" not in span_idx[CYCLE_A_UUID].attributes["ag"]["metrics"]["tokens"] + ) + assert ( + "cumulative" not in span_idx[CYCLE_B_UUID].attributes["ag"]["metrics"]["tokens"] + ) + + def test_cumulate_tokens_and_costs_propagate_from_children_to_parent(): root = _span( span_id=ROOT_UUID,