fix(api): make an agent run's cost and tokens reach the trace

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 = {

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

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()

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)