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,