diff --git a/.claude/notes/agents.md b/.claude/notes/agents.md index 1135be43..9857d839 100644 --- a/.claude/notes/agents.md +++ b/.claude/notes/agents.md @@ -306,6 +306,15 @@ cost bills it at the input rate and the cache buckets separately. id), and the `usage` snapshot the last resort. Cost is backfilled from the rate card when a turn was timed out or killed, so there is no terminal `ResultMessage` — the tokens are already captured, so this is pure pricing. + `model_usage` and `total_cost_usd` count the whole SESSION: a CLI resumed with + `resume=` restores the running totals it saved on exit, so a dialog's turn N reports + turns 1..N (`usage` and `num_turns` stay per-turn). `_turn_usage_slice` subtracts what + the same session last reported, from any ResultMessage, error or not, since the CLI + saves on every exit. Summing unsliced snapshots shows as a reconciliation row equal to + the previous turn's total (200/200 multi-turn tasks of a 2026-10-07 Haiku 5.5 run), + and inflates the simulator's tokens the same way. A counter that moved + backwards means no restore, and the snapshot is booked whole. A killed turn reports + nothing, so if its CLI saved totals, the next turn absorbs its usage. - **Codex/OpenAI** report `input_tokens` INCLUSIVE of the cached prefix, and bill no separate cache-write fee. So the fresh slice is `input - cached`, `cache_creation` is 0, and `cache_read` is `cached`. diff --git a/docs/agents/CLAUDE_CODE.md b/docs/agents/CLAUDE_CODE.md index 84e2146c..e3473089 100644 --- a/docs/agents/CLAUDE_CODE.md +++ b/docs/agents/CLAUDE_CODE.md @@ -212,9 +212,11 @@ resolution for any agent that does not declare `supports_cooperative_stop`. See Claude Code produces the richest telemetry of the agents: -- **Authoritative billing** comes from the SDK's cumulative `model_usage` on the +- **Authoritative billing** comes from the SDK's `model_usage` on the terminal result message and reconciles to `total_cost_usd` — this already includes - sub-agent consumption that the per-message stream under-reports. + sub-agent consumption that the per-message stream under-reports. Both count the + whole session, so a resumed turn books only what changed since the session's + previous result message. - **Sub-agent accounting** is derived by grouping `parent_tool_use_id`-tagged assistant messages; there is no separate per-sub-agent field. The terminal sub-agent generation (delivered as the Agent tool result, never streamed) is diff --git a/src/coder_eval/agents/claude_code_agent.py b/src/coder_eval/agents/claude_code_agent.py index ecde1c38..3a626769 100644 --- a/src/coder_eval/agents/claude_code_agent.py +++ b/src/coder_eval/agents/claude_code_agent.py @@ -10,7 +10,7 @@ from contextlib import suppress from datetime import datetime, timedelta from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, NamedTuple from claude_agent_sdk import ( ClaudeAgentOptions, @@ -178,6 +178,79 @@ def _is_sdk_result_message(message: Any) -> bool: return hasattr(message, "session_id") and hasattr(message, "usage") and not _is_task_notification(message) +_CUMULATIVE_MODEL_USAGE_KEYS = ( + "inputTokens", + "outputTokens", + "cacheReadInputTokens", + "cacheCreationInputTokens", + "thinkingTokens", + "webSearchRequests", + "costUSD", +) + + +class _SessionUsage(NamedTuple): + """The session-cumulative totals a CLI session last reported on its ResultMessage.""" + + session_id: str + model_usage: dict[str, Any] + cost: float | None + + +def _model_usage_since(cumulative: dict[str, Any], baseline: dict[str, Any]) -> dict[str, Any] | None: + """Per-model ``model_usage`` minus an earlier snapshot of the same session. + + ``None`` when a model or counter moved backwards: the snapshot does not continue + ``baseline``. + """ + if any(model not in cumulative for model in baseline): + return None + sliced: dict[str, Any] = {} + for model, entry in cumulative.items(): + prev = baseline.get(model) + if not isinstance(entry, dict) or not isinstance(prev, dict): + sliced[model] = entry + continue + turn = dict(entry) + for key in _CUMULATIVE_MODEL_USAGE_KEYS: + now, before = entry.get(key), prev.get(key) + if not isinstance(now, int | float) or not isinstance(before, int | float): + continue + if now < before - 1e-9: + return None + turn[key] = max(now - before, 0) + sliced[model] = turn + return sliced + + +def _turn_usage_slice( + model_usage: Any, cost: Any, baseline: _SessionUsage | None, resumed_from: str | None +) -> tuple[Any, Any]: + """This turn's own ``model_usage`` and ``total_cost_usd`` from a ResultMessage. + + A resumed CLI restores its session's running totals, so a resumed turn's + ResultMessage also counts every earlier turn. Subtract what that session last + reported. A total that moved backwards means the CLI did not restore them, so the + snapshot is already this turn's and is returned whole. + + Rationale: .claude/notes/agents.md § Token accounting, per harness + """ + if baseline is None or resumed_from is None or baseline.session_id != resumed_from: + return model_usage, cost + if not isinstance(model_usage, dict) or not model_usage: + return model_usage, cost + sliced = _model_usage_since(model_usage, baseline.model_usage) + if sliced is None: + logger.debug("model_usage of resumed session %s moved backwards; booking it whole", resumed_from) + return model_usage, cost + turn_cost = ( + max(cost - baseline.cost, 0.0) + if isinstance(cost, int | float) and isinstance(baseline.cost, int | float) + else None + ) + return sliced, turn_cost + + _JSON_START_SEARCH_LIMIT = 200 @@ -253,6 +326,8 @@ def __init__( self.first_output_seen: bool = False # SDK ResultMessage capture. + self.resumed_from = agent._session_id + self.session_usage_baseline = agent._session_usage self.sdk_result_usage: dict[str, Any] | None = None self.sdk_result_model_usage: dict[str, Any] | None = None self.sdk_result_cost: float | None = None @@ -471,12 +546,20 @@ def on_task_notification(self, message: Message) -> None: def on_result_message(self, message: Message) -> None: """Capture the SDK ResultMessage usage/cost/session + the id-less backfill.""" self.sdk_result_usage = getattr(message, "usage", None) - self.sdk_result_model_usage = getattr(message, "model_usage", None) - self.sdk_result_cost = getattr(message, "total_cost_usd", None) + model_usage = getattr(message, "model_usage", None) + cost = getattr(message, "total_cost_usd", None) + self.sdk_result_model_usage, self.sdk_result_cost = _turn_usage_slice( + model_usage, cost, self.session_usage_baseline, self.resumed_from + ) self.num_turns = getattr(message, "num_turns", None) self.sdk_result_summary = self._agent._summarize_result(message) - # Only advance session_id on clean turns. new_session_id = getattr(message, "session_id", None) + # Every ResultMessage, error or not: the CLI saves its running totals on exit. + if isinstance(model_usage, dict) and model_usage and isinstance(new_session_id, str): + self._agent._session_usage = _SessionUsage( + new_session_id, model_usage, cost if isinstance(cost, int | float) else None + ) + # Only advance session_id on clean turns. if self.sdk_result_summary is not None and self.sdk_result_summary.is_error: self.log.debug( "is_error ResultMessage; not advancing session_id (kept %s)", @@ -743,6 +826,7 @@ def __init__( # Turn-lifecycle bookkeeping lives on the Agent base class. self._sdk_options_dump: dict[str, Any] | None = None self._session_id: str | None = None + self._session_usage: _SessionUsage | None = None # Held only while a communicate() call is in flight, so kill() can reach # the CLI subprocess when the SDK swallows asyncio cancellation. self._active_transport: SubprocessCLITransport | None = None @@ -1373,10 +1457,10 @@ def _finalize_commands( @staticmethod def _aggregate_model_usage(model_usage: dict[str, Any] | None) -> TokenUsage | None: - """Sum the SDK ResultMessage ``model_usage`` into a cumulative TokenUsage. + """Sum the turn's ``model_usage`` (see ``_turn_usage_slice``) into a TokenUsage. - Maps each model id to its cumulative session billing (camelCase, unlike - ``usage``). The SDK's authoritative cost breakdown: summed and priced it + Maps each model id to its billing (camelCase, unlike ``usage``). The SDK's + authoritative cost breakdown: summed and priced it reconciles to ``total_cost_usd`` exactly, and it INCLUDES sub-agent consumption the stream under-reports. None when absent, so the caller can fall back. @@ -1413,12 +1497,12 @@ def _build_token_usage( sdk_result_model_usage: dict[str, Any] | None = None, model: str | None = None, ) -> TokenUsage | None: - """Build the run's cumulative TokenUsage, or None if unavailable. + """Build the turn's TokenUsage, or None if unavailable. Source-of-truth order: - 1. ``ResultMessage.model_usage`` — the SDK's cumulative per-model billing, - authoritative and inclusive of sub-agent consumption. Prefer it. + 1. ``ResultMessage.model_usage`` — the SDK's per-model billing, sliced to + this turn, authoritative and inclusive of sub-agent consumption. Prefer it. 2. Per-call telemetry stream (sum) — used when ``model_usage`` is absent. Exact only when EVERY token-bearing emission carries a ``message_id``, since that is what the dedup keys on. diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py index ed78af8b..13d98a2d 100644 --- a/tests/test_token_usage.py +++ b/tests/test_token_usage.py @@ -11,6 +11,9 @@ ClaudeCodeAgent, _is_sdk_result_message, _is_task_notification, + _model_usage_since, + _SessionUsage, + _turn_usage_slice, ) from coder_eval.errors import AgentCrashError from coder_eval.models import ( @@ -736,6 +739,138 @@ async def mock_query(*args, **kwargs): assert sub[0].content_blocks[0].text == "5050" +# --- Resumed-session usage slicing --- + + +def _haiku_usage(inp: int, out: int, cache_read: int, cache_write: int, cost: float) -> dict: + """One ``model_usage`` entry shaped as CLI 2.1.281 reports it.""" + return { + "eu.anthropic.claude-haiku-5-5": { + "inputTokens": inp, + "outputTokens": out, + "cacheReadInputTokens": cache_read, + "cacheCreationInputTokens": cache_write, + "webSearchRequests": 0, + "costUSD": cost, + "contextWindow": 200000, + "maxOutputTokens": 32000, + } + } + + +# Three resumed turns of one session, as captured live on Bedrock (CLI 2.1.281). +_LIVE_SESSION = [ + (_haiku_usage(6, 468, 24103, 25347, 0.1409396), 0.1409396), + (_haiku_usage(12, 829, 75727, 26920, 0.1663734), 0.1663734), + (_haiku_usage(16, 877, 102647, 28016, 0.1782134), 0.1782134), +] + + +def _result_message(session_id: str, model_usage: dict, cost: float, *, is_error: bool = False) -> MagicMock: + result = MagicMock(spec=["session_id", "usage", "num_turns", "total_cost_usd", "model_usage", "is_error"]) + result.session_id = session_id + result.usage = {} + result.num_turns = 1 + result.total_cost_usd = cost + result.model_usage = model_usage + result.is_error = is_error + return result + + +class TestResumedSessionUsage: + """A resumed CLI reports session totals; each turn books only its own slice.""" + + def test_slice_subtracts_previous_totals(self): + (first, first_cost), (second, second_cost) = _LIVE_SESSION[:2] + sliced, cost = _turn_usage_slice(second, second_cost, _SessionUsage("s1", first, first_cost), "s1") + entry = sliced["eu.anthropic.claude-haiku-5-5"] + assert (entry["inputTokens"], entry["outputTokens"]) == (6, 361) + assert (entry["cacheReadInputTokens"], entry["cacheCreationInputTokens"]) == (51624, 1573) + assert entry["costUSD"] == pytest.approx(0.0254338) + assert entry["contextWindow"] == 200000 + assert cost == pytest.approx(0.0254338) + + @pytest.mark.parametrize( + ("baseline_session", "resumed_from"), + [("s1", None), ("s0", "s1")], + ids=["fresh-session", "other-session"], + ) + def test_no_slice_without_a_resume_of_the_same_session(self, baseline_session, resumed_from): + (first, first_cost), (second, second_cost) = _LIVE_SESSION[:2] + baseline = _SessionUsage(baseline_session, first, first_cost) + assert _turn_usage_slice(second, second_cost, baseline, resumed_from) == (second, second_cost) + + def test_counter_moving_backwards_books_the_snapshot_whole(self): + (first, first_cost), (second, second_cost) = _LIVE_SESSION[:2] + baseline = _SessionUsage("s1", second, second_cost) + assert _turn_usage_slice(first, first_cost, baseline, "s1") == (first, first_cost) + + def test_model_dropped_from_the_snapshot_is_not_a_continuation(self): + first, _ = _LIVE_SESSION[0] + assert _model_usage_since({}, first) is None + + def test_model_new_this_turn_is_booked_whole(self): + (first, _), (second, _) = _LIVE_SESSION[:2] + helper = {"claude-sonnet-x": {"inputTokens": 7, "outputTokens": 3, "costUSD": 0.01}} + sliced = _model_usage_since({**second, **helper}, first) + assert sliced is not None + assert sliced["claude-sonnet-x"] == helper["claude-sonnet-x"] + assert sliced["eu.anthropic.claude-haiku-5-5"]["outputTokens"] == 361 + + @pytest.mark.asyncio + async def test_dialog_books_each_turn_once(self): + """Three ``communicate()`` calls on one session sum to the session total, not to the sum of totals.""" + config = parse_agent_config(type=AgentKind.CLAUDE_CODE, permission_mode="acceptEdits", allowed_tools=["Read"]) + agent = ClaudeCodeAgent(config) + agent.working_directory = MagicMock() + agent.working_directory.rglob.return_value = [] + resumes: list[str | None] = [] + + records = [] + for model_usage, cost in _LIVE_SESSION: + + async def mock_query(*args, _usage=model_usage, _cost=cost, **kwargs): + resumes.append(kwargs["options"].resume) + yield _result_message("s1", _usage, _cost) + + with patch("coder_eval.agents.claude_code_agent.query", side_effect=mock_query): + records.append(await agent.communicate("next")) + + assert resumes == [None, "s1", "s1"] + turns = [r.token_usage for r in records] + assert all(t is not None for t in turns) + assert [t.output_tokens for t in turns] == [468, 361, 48] + assert [t.cache_read_input_tokens for t in turns] == [24103, 51624, 26920] + assert sum(t.cache_creation_input_tokens for t in turns) == 28016 + assert sum(t.total_cost_usd for t in turns) == pytest.approx(0.1782134) + + @pytest.mark.asyncio + async def test_error_turn_still_moves_the_baseline(self): + """The CLI saves totals on every exit, so an is_error turn's totals are the next turn's baseline.""" + config = parse_agent_config(type=AgentKind.CLAUDE_CODE, permission_mode="acceptEdits", allowed_tools=["Read"]) + agent = ClaudeCodeAgent(config) + agent.working_directory = MagicMock() + agent.working_directory.rglob.return_value = [] + + results = [ + _result_message("s1", *_LIVE_SESSION[0]), + _result_message("s1", *_LIVE_SESSION[1], is_error=True), + _result_message("s1", *_LIVE_SESSION[2]), + ] + records = [] + for result in results: + + async def mock_query(*args, _result=result, **kwargs): + yield _result + + with patch("coder_eval.agents.claude_code_agent.query", side_effect=mock_query): + records.append(await agent.communicate("next")) + + assert records[2].token_usage is not None + assert records[2].token_usage.output_tokens == 48 + assert records[2].token_usage.cache_read_input_tokens == 26920 + + # --- Orchestrator aggregation tests ---