From 64214f155550cc460d750010c9a938283cf16bbc Mon Sep 17 00:00:00 2001 From: Eric Lee Date: Sun, 2 Aug 2026 07:23:23 -0700 Subject: [PATCH] fix(usage): carry cumulative cache tokens in result.usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `agent_loop_compat`'s aggregator summed only `input_tokens` and `output_tokens`. `input_tokens` counts only cache MISSES — the shared convention across providers (Anthropic natively; OpenAI-compatible ones since #785's cache-read split) — so the cumulative total silently omitted every cached token. That total is not a display nicety. `agent_server` passes this dict straight to `compute_cost`, which then bills the cached portion at nothing: a real turn (2613-token prompt, 2560 served from cache) lost 71.7% of its cost. `eval/harbor/clawcodex_agent.py` had already diagnosed the same defect from the other end — 6 tokens reported against a real ~412 K on a prompt-cached opus run — and routes around it by reading the session cost block. The fix is deliberately ADDITIVE: `cache_read_input_tokens` and `cache_creation_input_tokens` are initialized and summed alongside the existing keys, and `input_tokens` keeps its meaning. Redefining it to include the cache would have been a stream-json contract change; adding counters beside it is not, so every existing reader is untouched and only gains the ability to see the rest. The harbor fallback already computed `input + cache_read + cache_creation` from this payload — it was written for exactly this convention and was simply starved of the data. It becomes correct rather than double-counting. The cost block stays preferred (it also supplies the real dollar cost, and older clawcodex builds still need it); three comments there that described the missing counters as a permanent property are updated. Two consumers in `entrypoints.headless` came along with it: * `/goal`'s token budget summed input+output, so on a warm prefix cache it saw about 2.3% of what had been spent. It now counts every billed token. * The multi-prompt accumulator summed EVERY key generically, including the `last_*` snapshot. Those are last-wins — the live-context measure — and adding a snapshot to itself across prompts produces a number that measures nothing, which then ships in the stream-json `result` payload. They now replace. Both were inline in a long function, so they are extracted as `_accumulate_usage` and `_billed_token_total`. That is not cosmetic: the first version of these tests re-implemented the loop and passed with BOTH headless fixes reverted. The cumulative sum still must NOT be used as a live-context measure — it double-counts context because every turn re-sends the whole conversation. That half of the `C3a` comment stands; only the "drops the cache keys" half is now stale and was corrected. Completing that total made a SECOND defect reachable, fixed here too. `agent_server` priced the turn by calling `compute_cost` on the cumulative dict, but `get_pricing` selects a tier from a PER-REQUEST threshold (gpt-5.6-luna at 272K, MiniMax-M3 at 512K). Cache reads sum to roughly turns x conversation-size, so a 25-turn loop over a 15K conversation reaches 385K and prices at the long-context rate no single request came near — 1.76x the true cost, where before the change it under-reported at 0.60x. Accurate tokens are what push the aggregate over the line, so the two changes could not ship apart. The turn's cost is now a delta of `cost_tracker`'s running total, which prices each response as it arrives with the tier chosen from that one request. That also removes a duplicate cost implementation rather than patching one. Only tiered models were affected; every linear-priced model aggregated exactly. Mutation-tested: both cache sums, the dict initialization, the `last_*` branch, the billed-key tuple, both headless call sites and the server's cost path. Four mutants initially survived, and each exposed a real gap: * every fixture used `cache_creation_input_tokens: 0`, which cannot distinguish summing a value from dropping it; * every fixture ended after ONE API call, so `+=` and `=` were indistinguishable — turn count is the axis that tests accumulation, and it is the same axis the tier-crossing lives on; * the headless assertions called the extracted helpers directly, so reverting the call sites changed nothing; * a single-prompt headless test cannot separate summing from replacing, because both agree on one accumulation. Co-Authored-By: Claude Opus 5 --- eval/harbor/clawcodex_agent.py | 39 ++- src/entrypoints/headless.py | 51 ++- src/query/agent_loop_compat.py | 43 ++- src/server/agent_server.py | 29 +- tests/server/test_agent_server_workflows.py | 90 +++++ tests/test_headless_cli.py | 125 +++++++ tests/test_usage_aggregation_cache_keys.py | 344 ++++++++++++++++++++ 7 files changed, 686 insertions(+), 35 deletions(-) create mode 100644 tests/test_usage_aggregation_cache_keys.py diff --git a/eval/harbor/clawcodex_agent.py b/eval/harbor/clawcodex_agent.py index 8b606ff1..e15365d2 100644 --- a/eval/harbor/clawcodex_agent.py +++ b/eval/harbor/clawcodex_agent.py @@ -686,8 +686,13 @@ def populate_context_post_run(self, context: AgentContext) -> None: # Leaderboard token/cost columns. Prefer the session's authoritative # BILLING totals (input + cache, summed per turn — matches - # claude-code); fall back to the stream-json usage (incomplete: no - # cumulative cache) only when no session cost block was synced. + # claude-code); fall back to the stream-json usage only when no + # session cost block was synced. That fallback now carries cumulative + # cache counters, so the sum below is no longer short by the cached + # portion — but it is still not COMPLETE: result.usage only ever sees + # the main loop's assistant messages, so subagent and compaction + # tokens never reach it. The cost block includes them, which is one + # reason it stays preferred. if totals is not None: context.n_input_tokens = totals["prompt"] context.n_cache_tokens = totals["cached"] @@ -748,8 +753,10 @@ def _final_metrics( totals: dict[str, Any] | None, ) -> FinalMetrics | None: """``totals`` = authoritative session billing totals (preferred). - Falls back to the stream-json usage (incomplete — no cumulative - cache) only when no session cost block was synced.""" + Falls back to the stream-json usage when no session cost block was + synced; that fallback now carries cumulative cache counters, so it is + no longer short by the cached portion — though it still omits + subagent and compaction tokens, which only the cost block sees.""" usage = (result_event or {}).get("usage") if not isinstance(usage, dict): usage = {} @@ -814,15 +821,21 @@ def _billing_totals_from_cost_block( ) -> dict[str, Any] | None: """Authoritative BILLING token totals from a session ``cost`` block. - clawcodex's stream-json ``result.usage`` is built for live-CONTEXT - measurement, not accounting: ``input_tokens`` is the running sum of - NON-cached input and it drops cumulative cache tokens (only a - ``last_*`` snapshot survives). For a heavily prompt-cached opus run - that under-reports total prompt tokens by orders of magnitude - (observed: 6 vs the real ~412 K). The session cost block's - ``model_usage`` accumulates the real per-turn billing counters - (input + cache_read + cache_creation, summed across turns), which is - exactly the convention the built-in claude-code agent reports. + Historically clawcodex's stream-json ``result.usage`` summed only + input+output, and ``input_tokens`` is the running sum of NON-cached + input, so cumulative cache tokens were dropped entirely (only a + ``last_*`` snapshot survived). For a heavily prompt-cached opus run + that under-reported total prompt tokens by orders of magnitude — + observed here as 6 against the real ~412 K, which is why this cost + block path exists. ``result.usage`` now carries the cumulative cache + counters, so the fallback is sound; this remains PREFERRED because + the cost block also supplies the real dollar cost, and because a run + against an older clawcodex build still needs it. + + The session cost block's ``model_usage`` accumulates the real + per-turn billing counters (input + cache_read + cache_creation, + summed across turns), which is exactly the convention the built-in + claude-code agent reports. Returns ``{prompt, completion, cached, cost}`` or ``None``. """ model_usage = cost.get("model_usage") diff --git a/src/entrypoints/headless.py b/src/entrypoints/headless.py index 3d550daa..fcb32dec 100644 --- a/src/entrypoints/headless.py +++ b/src/entrypoints/headless.py @@ -64,6 +64,49 @@ from src.utils.abort_controller import AbortController, AbortError +def _accumulate_usage(total: dict[str, int], usage: dict[str, int]) -> None: + """Fold one prompt's usage into the running totals, in place. + + Two different kinds of key live in the same dict and must not be treated + alike. The cumulative counters (``input_tokens``, ``output_tokens`` and + the two cache counters) ADD across prompts — together they are the + billing total. The ``last_*`` keys are a snapshot of the most recent + response, used as the live-context measure, so they REPLACE: adding a + snapshot to itself across prompts yields a number that measures nothing, + and it ships in the stream-json ``result`` payload. + """ + for key, value in usage.items(): + # Provider usage dicts carry non-numeric entries — MinimaxProvider + # puts a ``service_tier`` string in its. The aggregator upstream + # whitelists keys today, so none arrive; skipping rather than + # raising means the natural next edit there ("stop dropping keys") + # cannot take down a headless run with ``int("standard")``. + try: + numeric = int(value) + except (TypeError, ValueError): + continue + if key.startswith("last_"): + total[key] = numeric + else: + total[key] = int(total.get(key, 0) or 0) + numeric + + +#: The cumulative counters that together make up what was actually billed. +#: ``input_tokens`` is only the NON-cached part of each prompt, so any total +#: built from input+output alone silently omits every cached token. +_BILLED_TOKEN_KEYS = ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) + + +def _billed_token_total(usage_total: dict[str, int]) -> int: + """Every billed token across the run — what ``/goal`` spends against.""" + return int(sum(int(usage_total.get(key, 0) or 0) for key in _BILLED_TOKEN_KEYS)) + + def _json_result_subtype( exit_code: int, early_stop_subtype: str | None ) -> tuple[str, bool]: @@ -824,8 +867,7 @@ def _persist(msg: Any) -> None: num_turns_total += result.num_turns if result.usage: - for key, value in result.usage.items(): - usage_total[key] = usage_total.get(key, 0) + int(value) + _accumulate_usage(usage_total, result.usage) if writer is not None: writer.write(AssistantEvent(text=result.response_text)) @@ -841,10 +883,7 @@ def _persist(msg: Any) -> None: evidence = collect_turn_evidence( list(session.conversation.messages) ) or (result.response_text or "") - tokens_now = int( - usage_total.get("input_tokens", 0) - + usage_total.get("output_tokens", 0) - ) + tokens_now = _billed_token_total(usage_total) decision = goal_mgr.evaluate_after_turn( evidence, tokens_now=tokens_now, ) diff --git a/src/query/agent_loop_compat.py b/src/query/agent_loop_compat.py index c04458e9..b9b229df 100644 --- a/src/query/agent_loop_compat.py +++ b/src/query/agent_loop_compat.py @@ -664,7 +664,12 @@ async def run_query_as_agent_loop( holder = TerminalHolder() response_text_parts: list[str] = [] - usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0} + usage: dict[str, int] = { + "input_tokens": 0, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + } num_turns = 0 last_assistant_text = "" last_api_error_text = "" @@ -717,16 +722,36 @@ async def run_query_as_agent_loop( num_turns += 1 # Sum usage across turns. mu = getattr(msg, "usage", None) or {} - usage["input_tokens"] += mu.get("input_tokens", 0) - usage["output_tokens"] += mu.get("output_tokens", 0) + usage["input_tokens"] += mu.get("input_tokens", 0) or 0 + usage["output_tokens"] += mu.get("output_tokens", 0) or 0 + # The cache counters are summed too, because ``input_tokens`` is + # only the NON-cached part of each prompt — that is the shared + # convention across providers (Anthropic reports it natively; + # OpenAI-compatible providers since the cache-read split). Summing + # input+output alone therefore drops every cached token from the + # total, and ``agent_server`` feeds this very dict to + # ``compute_cost``, which then bills the cached portion at nothing: + # a real 2613-token turn with a 2560-token hit lost 71.7% of its + # cost. ``eval/harbor/clawcodex_agent.py`` measured the same defect + # from the other end — 6 tokens reported against a real ~412K on a + # prompt-cached opus run — and works around it by reading the + # session cost block instead. + # + # Purely additive: ``input_tokens`` keeps its meaning, so existing + # readers are untouched and only gain the ability to see the rest. + usage["cache_read_input_tokens"] += mu.get("cache_read_input_tokens", 0) or 0 + usage["cache_creation_input_tokens"] += ( + mu.get("cache_creation_input_tokens", 0) or 0 + ) # C3a: also keep the LAST response's FULL usage (all four # keys, last-wins — TS getCurrentUsage, utils/tokens.ts: - # 152-171). The cumulative sum above double-counts context - # across a multi-tool-call run and drops the cache keys, so - # it must NOT be used as a live-context measure - # (tokens.ts:407-420 warning). Last-wins also covers the - # parallel-tool-call case where N assistant records share - # one usage object. + # 152-171). The cumulative sum above is a BILLING total and + # still must NOT be used as a live-context measure: it + # double-counts context across a multi-tool-call run, because + # each turn re-sends the whole conversation (tokens.ts:407-420 + # warning). Last-wins is the context measure, and it also covers + # the parallel-tool-call case where N assistant records share one + # usage object. if mu: usage["last_input_tokens"] = int(mu.get("input_tokens", 0) or 0) usage["last_output_tokens"] = int(mu.get("output_tokens", 0) or 0) diff --git a/src/server/agent_server.py b/src/server/agent_server.py index b91cd415..7ef0b5a1 100644 --- a/src/server/agent_server.py +++ b/src/server/agent_server.py @@ -4782,6 +4782,12 @@ def on_message(message: Any) -> None: # subagents spawn from the Agent tool's captured FULL registry. from src.coordinator.mode import coordinator_main_loop_registry + # Cost is read as a DELTA of the tracker's running total rather + # than recomputed from ``result.usage`` below — see the note at + # the ``_cost`` assignment for why the aggregate cannot be priced. + from src.bootstrap.state import get_total_cost_usd + + _cost_before = get_total_cost_usd() result = asyncio.run(run_query_as_agent_loop( initial_messages=list(self.session.conversation.messages), provider=turn_provider, @@ -4848,14 +4854,23 @@ def on_message(message: Any) -> None: msgs.extend(_btw_snapshot) _usage = result.usage if result.num_turns > 0 else None + # The turn's cost is the tracker's delta, NOT ``compute_cost`` over + # ``result.usage``. That dict is the sum across every loop turn, and + # ``get_pricing`` selects a tier from a PER-REQUEST threshold + # (gpt-5.6-luna at 272K, MiniMax-M3 at 512K). Cache reads sum to + # roughly turns x conversation-size, so a long loop of small requests + # crosses a boundary no single request came near: a 25-turn loop over + # a 15K conversation prices at 1.76x the true cost. + # + # ``cost_tracker.record_api_usage`` already prices each response as it + # arrives, with the tier chosen from that one request, so the running + # total is correct by construction. Reading its delta also drops a + # duplicate cost implementation rather than patching one. _cost = 0.0 - if _usage: - try: - from src.services.pricing import compute_cost - - _cost = compute_cost(getattr(self.provider, "model", None) or self.config.model or "", _usage) - except Exception: # noqa: BLE001 — cost is best-effort, never break the turn - _cost = 0.0 + try: + _cost = max(0.0, get_total_cost_usd() - _cost_before) + except Exception: # noqa: BLE001 — cost is best-effort, never break the turn + _cost = 0.0 # One more completed user turn. Internal (notification) and btw # (ephemeral, rolled-back) turns don't move the odometer — same rule # as the deleted REPL, which only counted real prompt→response rounds. diff --git a/tests/server/test_agent_server_workflows.py b/tests/server/test_agent_server_workflows.py index e28c9a1d..e39ccc67 100644 --- a/tests/server/test_agent_server_workflows.py +++ b/tests/server/test_agent_server_workflows.py @@ -921,3 +921,93 @@ def _stop(*a, **k): {"subtype": "error_during_execution", "response_text": "[Stopped: x]"}, ) assert not sess2._inbox.empty(), "control failed — nothing ever enqueues" + + +class _TieredCacheProvider: + """Many small requests whose cache reads only cross a tier once summed. + + Each request is ~15.4K prompt tokens — far below gpt-5.6-luna's 272K + long-context threshold — but the loop's cumulative total reaches 385K. + """ + + PER_REQUEST_USAGE = { + "input_tokens": 400, + "output_tokens": 300, + "cache_read_input_tokens": 15_000, + "cache_creation_input_tokens": 0, + } + TURNS = 25 + + def __init__(self, api_key=None, base_url=None, model=None): + self.model = "openai/gpt-5.6-luna" + self._calls = 0 + + def chat(self, messages, tools=None, **kw): + from src.providers.base import ChatResponse + + self._calls += 1 + last = self._calls >= self.TURNS + return ChatResponse( + content="done" if last else "thinking", + model=self.model, + usage=dict(self.PER_REQUEST_USAGE), + finish_reason="end_turn" if last else "tool_use", + tool_uses=None + if last + else [ + { + "id": f"tool_{self._calls}", + "name": "Bash", + "input": {"command": "true", "description": "noop"}, + } + ], + ) + + def chat_stream_response(self, *a, **kw): # force fallback to chat() + raise NotImplementedError + + +async def test_turn_cost_is_not_priced_from_the_cumulative_usage(tmp_path): + """A long loop must not be billed at the long-context rate. + + ``get_pricing`` picks a tier from a PER-REQUEST threshold, but + ``result.usage`` is the sum across every loop turn. Once the cumulative + dict carried cache reads — which sum to roughly turns x conversation + size — pricing it crossed a boundary no single request came near, and + this turn billed at 1.76x its true cost. + + The cost is now a delta of ``cost_tracker``'s running total, which prices + each response as it arrives with the tier chosen from that one request. + + BEHAVIOURAL: asserted on the emitted ``total_cost_usd``, because the + failure mode is a plausible-looking number rather than an error. + """ + from src.services.pricing import compute_cost + + per_request = _TieredCacheProvider.PER_REQUEST_USAGE + turns = _TieredCacheProvider.TURNS + model = "openai/gpt-5.6-luna" + + truth = sum(compute_cost(model, per_request) for _ in range(turns)) + aggregate = {k: v * turns for k, v in per_request.items()} + if_priced_as_aggregate = compute_cost(model, aggregate) + assert if_priced_as_aggregate > truth * 1.5, "fixture must cross the tier" + + async with _spawned(tmp_path, _TieredCacheProvider) as (handle, gen): + await handle.send_to_agent( + {"type": "user", "message": {"role": "user", "content": "go"}} + ) + result = None + for _ in range(200): + msg = await asyncio.wait_for(gen.__anext__(), timeout=20) + if msg.get("type") == "result": + result = msg + break + + assert result is not None, "no result frame emitted" + reported = float(result.get("total_cost_usd") or 0.0) + assert reported > 0, "the turn reported no cost at all" + assert reported < if_priced_as_aggregate * 0.8, ( + f"turn billed at the long-context rate: {reported} vs aggregate " + f"{if_priced_as_aggregate} (true {truth})" + ) diff --git a/tests/test_headless_cli.py b/tests/test_headless_cli.py index 86dfcba9..679d0e45 100644 --- a/tests/test_headless_cli.py +++ b/tests/test_headless_cli.py @@ -789,3 +789,128 @@ def test_subtype_and_is_error_never_disagree(exit_code, early, expected): a failure. """ assert headless_mod._json_result_subtype(exit_code, early) == expected + + +# --------------------------------------------------------------------------- +# usage accounting in the emitted result + + +def _cached_response(text: str, *, cache_read: int) -> ChatResponse: + """A response whose prompt was mostly served from the cache.""" + return ChatResponse( + content=text, + model="fake-model", + usage={ + "input_tokens": 5, + "output_tokens": 3, + "cache_read_input_tokens": cache_read, + "cache_creation_input_tokens": 2, + }, + finish_reason="end_turn", + tool_uses=None, + ) + + +def test_headless_result_usage_sums_cache_tokens(fake_wiring, tmp_path): + """The emitted `result.usage` must be a complete billing total. + + Drives `run_headless` rather than the accumulation helper: an earlier + version of these assertions called the helper directly and passed with + the call site reverted to the old inline loop. + """ + fake_wiring.append(_cached_response("hi", cache_read=1000)) + + stdout = io.StringIO() + code = run_headless( + HeadlessOptions( + prompt="hi", + output_format="json", + stdout=stdout, + stderr=io.StringIO(), + workspace_root=tmp_path, + ) + ) + + assert code == 0 + payload = json.loads(stdout.getvalue()) + usage = payload["usage"] + assert usage["cache_read_input_tokens"] == 1000, ( + "cumulative cache tokens never reached the emitted result" + ) + assert usage["cache_creation_input_tokens"] == 2 + # input_tokens keeps its meaning — cache misses only + assert usage["input_tokens"] == 5 + + +def test_headless_result_usage_keeps_last_snapshot_unsummed(fake_wiring, tmp_path): + """`last_*` are a snapshot of the most recent response, not a running sum. + + Summing them across turns produces a number that measures nothing, and it + ships in this payload. + """ + fake_wiring.append(_cached_response("hi", cache_read=1000)) + + stdout = io.StringIO() + run_headless( + HeadlessOptions( + prompt="hi", + output_format="json", + stdout=stdout, + stderr=io.StringIO(), + workspace_root=tmp_path, + ) + ) + + usage = json.loads(stdout.getvalue())["usage"] + assert usage["last_cache_read_input_tokens"] == 1000 + assert usage["last_input_tokens"] == 5 + + +def test_headless_multi_prompt_sums_cache_but_not_the_last_snapshot( + fake_wiring, tmp_path +): + """Two prompts: cumulative keys ADD, `last_*` keys REPLACE. + + One prompt cannot tell the two apart — the generic + `total = total + value` loop and the split accumulator agree on a single + accumulation, which is why a mutant reverting the call site survived a + single-prompt test. The divergence only appears from the second prompt on. + """ + fake_wiring.append(_cached_response("A", cache_read=1000)) + fake_wiring.append(_cached_response("B", cache_read=3000)) + + stdin = io.StringIO( + "\n".join( + [ + json.dumps({"type": "user", "message": {"content": "one"}}), + json.dumps({"type": "user", "message": {"content": "two"}}), + ] + ) + + "\n" + ) + stdout = io.StringIO() + code = run_headless( + HeadlessOptions( + output_format="stream-json", + input_format="stream-json", + stdin=stdin, + stdout=stdout, + stderr=io.StringIO(), + workspace_root=tmp_path, + ) + ) + + assert code == 0 + parsed = [json.loads(l) for l in stdout.getvalue().splitlines() if l.strip()] + usage = parsed[-1]["usage"] + + # cumulative: both prompts' cache reads + assert usage["cache_read_input_tokens"] == 4000 + assert usage["cache_creation_input_tokens"] == 4 + assert usage["input_tokens"] == 10 + + # snapshot: the SECOND prompt's value alone, not 1000 + 3000 + assert usage["last_cache_read_input_tokens"] == 3000, ( + "last_* was summed across prompts, which measures nothing" + ) + assert usage["last_input_tokens"] == 5 diff --git a/tests/test_usage_aggregation_cache_keys.py b/tests/test_usage_aggregation_cache_keys.py new file mode 100644 index 00000000..79578727 --- /dev/null +++ b/tests/test_usage_aggregation_cache_keys.py @@ -0,0 +1,344 @@ +"""`result.usage` must be a complete BILLING total, not input+output. + +``input_tokens`` counts only cache MISSES — that is the shared convention +across providers (Anthropic reports it natively; OpenAI-compatible providers +since the cache-read split). Summing input+output alone therefore dropped +every cached token from the total, and `agent_server` feeds that same dict +straight to `compute_cost`. +""" + +from __future__ import annotations + +import asyncio +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +from src.providers.base import ChatResponse +from src.entrypoints.headless import _accumulate_usage, _billed_token_total +from src.query.agent_loop_compat import run_query_as_agent_loop +from src.services.pricing import compute_cost +from src.tool_system.context import ToolContext +from src.tool_system.defaults import build_default_registry +from src.types.messages import UserMessage + + +def _run(coro): + return asyncio.run(coro) + + +def _provider(usage: dict) -> MagicMock: + provider = MagicMock() + provider.chat_stream_response.side_effect = NotImplementedError() + provider.chat.return_value = ChatResponse( + content="done", + model="test", + usage=usage, + finish_reason="end_turn", + tool_uses=None, + ) + return provider + + +class TestUsageAggregationCarriesCacheTokens(unittest.TestCase): + def setUp(self) -> None: + self.temp_dir = tempfile.TemporaryDirectory() + self.registry = build_default_registry() + self.context = ToolContext(workspace_root=Path(self.temp_dir.name)) + + def tearDown(self) -> None: + self.temp_dir.cleanup() + + def _usage_for(self, usage: dict) -> dict: + result = _run( + run_query_as_agent_loop( + initial_messages=[UserMessage(content="Hi")], + provider=_provider(usage), + tool_registry=self.registry, + tool_context=self.context, + system_prompt="You are helpful.", + max_turns=5, + ) + ) + return result.usage + + def test_cache_tokens_survive_the_cumulative_sum(self): + """Both counters carry through, with distinct non-zero values. + + Both matter and they are priced differently — Anthropic charges a + premium to WRITE the cache and a discount to READ it. Testing either + one at zero would not distinguish summing it from dropping it. + """ + usage = self._usage_for( + { + "input_tokens": 53, + "output_tokens": 8, + "cache_read_input_tokens": 2560, + "cache_creation_input_tokens": 384, + } + ) + self.assertEqual(usage["input_tokens"], 53) + self.assertEqual(usage["output_tokens"], 8) + self.assertEqual(usage["cache_read_input_tokens"], 2560) + self.assertEqual(usage["cache_creation_input_tokens"], 384) + + def test_the_aggregated_dict_prices_a_cached_turn_correctly(self): + """`agent_server` passes this dict straight to `compute_cost`. + + Numbers are a real turn measured against DeepSeek: 2613-token prompt, + 2560 of it served from the prefix cache. Without the cache counters + the cached portion was billed at nothing — 71.7% of the turn's cost + simply absent from the total. + """ + per_turn = { + "input_tokens": 53, + "output_tokens": 8, + "cache_read_input_tokens": 2560, + "cache_creation_input_tokens": 384, + } + aggregated = self._usage_for(per_turn) + + model = "openai/gpt-5.6-luna" + truth = compute_cost(model, per_turn) + self.assertGreater(truth, 0) + self.assertAlmostEqual(compute_cost(model, aggregated), truth, places=12) + + # The identity above is the assertion. A ratio against input+output + # alone would pin a live pricing rate, breaking on a rate edit for a + # reason that has nothing to do with aggregation. + + def test_cache_tokens_ACCUMULATE_across_loop_turns(self): + """The counters must ADD across turns, not just survive one. + + Every other fixture here ends after a single API call, so `+=` and a + plain `=` are indistinguishable — a mutant swapping them passed the + whole file. Accumulation IS the claim ("result.usage is a billing + total"), and turn count is the axis that tests it. Distinct values + per turn so a last-wins assignment cannot coincide with the sum. + """ + provider = MagicMock() + provider.chat_stream_response.side_effect = NotImplementedError() + provider.chat.side_effect = [ + ChatResponse( + content="running a command", + model="test", + usage={ + "input_tokens": 100, + "output_tokens": 10, + "cache_read_input_tokens": 1000, + "cache_creation_input_tokens": 64, + }, + finish_reason="tool_use", + tool_uses=[ + { + "id": "tool_1", + "name": "Bash", + "input": {"command": "true", "description": "noop"}, + } + ], + ), + ChatResponse( + content="done", + model="test", + usage={ + "input_tokens": 200, + "output_tokens": 20, + "cache_read_input_tokens": 3000, + "cache_creation_input_tokens": 128, + }, + finish_reason="end_turn", + tool_uses=None, + ), + ] + + result = _run( + run_query_as_agent_loop( + initial_messages=[UserMessage(content="Hi")], + provider=provider, + tool_registry=self.registry, + tool_context=self.context, + system_prompt="You are helpful.", + max_turns=5, + ) + ) + + usage = result.usage + self.assertEqual(usage["input_tokens"], 300) + self.assertEqual(usage["output_tokens"], 30) + self.assertEqual(usage["cache_read_input_tokens"], 4000) + self.assertEqual(usage["cache_creation_input_tokens"], 192) + # and the snapshot is still the LAST turn, not the sum + self.assertEqual(usage["last_cache_read_input_tokens"], 3000) + + def test_a_provider_reporting_no_cache_still_gets_zeroed_counters(self): + """Keys are always present, so consumers need no `.get` defaults.""" + usage = self._usage_for({"input_tokens": 10, "output_tokens": 5}) + self.assertEqual(usage["cache_read_input_tokens"], 0) + self.assertEqual(usage["cache_creation_input_tokens"], 0) + self.assertEqual(usage["input_tokens"], 10) + + def test_input_tokens_keeps_its_meaning(self): + """The change is additive — existing readers must be untouched. + + Redefining `input_tokens` to include the cache would have been a + stream-json contract change; adding counters alongside it is not. + """ + usage = self._usage_for( + { + "input_tokens": 53, + "output_tokens": 8, + "cache_read_input_tokens": 2560, + "cache_creation_input_tokens": 7, + } + ) + self.assertEqual(usage["input_tokens"], 53) + + def test_the_last_snapshot_still_measures_live_context(self): + """The cumulative sum is billing; `last_*` remains the context view. + + They must not be conflated: the sum double-counts context because + every turn re-sends the whole conversation. + """ + usage = self._usage_for( + { + "input_tokens": 53, + "output_tokens": 8, + "cache_read_input_tokens": 2560, + "cache_creation_input_tokens": 0, + } + ) + self.assertEqual(usage["last_input_tokens"], 53) + self.assertEqual(usage["last_cache_read_input_tokens"], 2560) + + +class TestHeadlessUsageTotals(unittest.TestCase): + """The multi-prompt accumulator in `entrypoints.headless`. + + Calls the real helpers. An earlier version of this class re-implemented + the loop inline, which passed with BOTH headless fixes reverted — the + accumulation lived inside a long function, and mirroring it tested the + mirror. The helpers were extracted so these assertions reach the code + that actually runs. + """ + + @staticmethod + def _accumulate(per_prompt: list[dict]) -> dict: + usage_total: dict[str, int] = {} + for usage in per_prompt: + _accumulate_usage(usage_total, usage) + return usage_total + + def test_cumulative_keys_add_and_last_keys_replace(self): + totals = self._accumulate( + [ + {"input_tokens": 10, "cache_read_input_tokens": 100, "last_input_tokens": 10}, + {"input_tokens": 20, "cache_read_input_tokens": 200, "last_input_tokens": 20}, + ] + ) + self.assertEqual(totals["input_tokens"], 30) + self.assertEqual(totals["cache_read_input_tokens"], 300) + # last-wins: summing a snapshot across prompts measures nothing, and + # the result rides out in the stream-json payload. + self.assertEqual(totals["last_input_tokens"], 20) + + def test_the_goal_budget_counts_every_billed_token(self): + """`/goal` spends against this; `input_tokens` alone is cache misses. + + On a warm prefix cache the budget saw a small fraction of what had + actually been spent. + """ + totals = self._accumulate( + [ + { + "input_tokens": 53, + "output_tokens": 8, + "cache_read_input_tokens": 2560, + "cache_creation_input_tokens": 0, + } + ] + ) + # 53 miss + 8 out + 2560 cached. The old computation saw 61 of these + # 2621 — the assertion that matters is the identity, not the ratio. + self.assertEqual(_billed_token_total(totals), 2621) + + +if __name__ == "__main__": + unittest.main() + + +class TestTurnCostIsNotPricedFromTheAggregate(unittest.TestCase): + """Cost must be priced per REQUEST, never over the cumulative dict. + + `get_pricing` selects a tier from a per-request threshold + (`gpt-5.6-luna` at 272K, `MiniMax-M3` at 512K). `result.usage` is the sum + across every loop turn, and cache reads sum to roughly + turns x conversation-size — so a long loop of small requests crosses a + boundary no single request came near. + + This is the trap that makes completing the cumulative dict actively + dangerous for anything that then prices it: making the token total + accurate is exactly what pushes the aggregate over the threshold. + """ + + MODEL = "openai/gpt-5.6-luna" + PER_REQUEST = { + "input_tokens": 400, + "output_tokens": 300, + "cache_read_input_tokens": 15_000, + "cache_creation_input_tokens": 0, + } + TURNS = 25 + + def test_pricing_the_aggregate_crosses_a_per_request_tier(self): + """Documents WHY the server must not call compute_cost on the sum.""" + truth = sum( + compute_cost(self.MODEL, self.PER_REQUEST) for _ in range(self.TURNS) + ) + aggregate = {k: v * self.TURNS for k, v in self.PER_REQUEST.items()} + priced_as_aggregate = compute_cost(self.MODEL, aggregate) + + prompt_total = ( + aggregate["input_tokens"] + aggregate["cache_read_input_tokens"] + ) + self.assertGreater(prompt_total, 272_000, "fixture must cross the tier") + self.assertGreater( + priced_as_aggregate, + truth * 1.5, + "pricing the aggregate should over-bill once the tier is crossed", + ) + + def test_below_the_tier_the_two_agree(self): + """The hazard is the threshold, not aggregation itself.""" + turns = 5 + truth = sum(compute_cost(self.MODEL, self.PER_REQUEST) for _ in range(turns)) + aggregate = {k: v * turns for k, v in self.PER_REQUEST.items()} + self.assertAlmostEqual(compute_cost(self.MODEL, aggregate), truth, places=12) + + def test_the_running_total_prices_each_request_separately(self): + """The mechanism the server now reads: a delta of this total. + + `record_api_usage` prices each response as it arrives, so the tier is + always chosen from one request and the accumulated total equals the + per-request sum even across a tier-crossing run. + """ + from src.bootstrap.state import get_total_cost_usd, reset_cost_state + from src.cost_tracker import record_api_usage + + reset_cost_state() + try: + before = get_total_cost_usd() + for _ in range(self.TURNS): + record_api_usage(self.MODEL, dict(self.PER_REQUEST)) + delta = get_total_cost_usd() - before + + truth = sum( + compute_cost(self.MODEL, self.PER_REQUEST) for _ in range(self.TURNS) + ) + self.assertAlmostEqual(delta, truth, places=10) + + # and it is materially BELOW what pricing the aggregate would give + aggregate = {k: v * self.TURNS for k, v in self.PER_REQUEST.items()} + self.assertLess(delta, compute_cost(self.MODEL, aggregate) * 0.7) + finally: + reset_cost_state()