diff --git a/headroom/subscription/session_tracking.py b/headroom/subscription/session_tracking.py index e9a1cacfd..9bcb557bb 100644 --- a/headroom/subscription/session_tracking.py +++ b/headroom/subscription/session_tracking.py @@ -134,6 +134,14 @@ def compute_window_tokens(start_ts: float, end_ts: float) -> WindowTokens: totals = WindowTokens() by_model: dict[str, WindowTokens] = {} unattributed = WindowTokens() + # Claude Code can store one assistant response across multiple transcript + # lines (e.g. one entry per content block), each carrying the SAME + # request-level ``message.usage``. Summing per line therefore multiplies a + # single response's tokens by its block count (observed 19x for one 420K + # response, #2340). Count each response's usage once, keyed by the unique + # Anthropic ``message.id``. Entries without an id keep the per-line + # behavior, so this only ever removes true duplicates. + seen_message_ids: set[str] = set() for path in find_transcript_files(): # Skip transcripts that cannot contain entries inside the window. @@ -171,6 +179,12 @@ def compute_window_tokens(start_ts: float, end_ts: float) -> WindowTokens: if not usage: continue + msg_id = msg.get("id") + if isinstance(msg_id, str) and msg_id: + if msg_id in seen_message_ids: + continue + seen_message_ids.add(msg_id) + _add_usage_to_tokens(totals, usage) model_id: str | None = msg.get("model") diff --git a/tests/test_subscription_session_tracking.py b/tests/test_subscription_session_tracking.py index 058846ced..9074c7f58 100644 --- a/tests/test_subscription_session_tracking.py +++ b/tests/test_subscription_session_tracking.py @@ -87,3 +87,44 @@ def test_read_transcript_lines_preserves_small_transcript( '{"marker":"first"}', '{"marker":"second"}', ] + + +def test_compute_window_tokens_dedups_usage_by_message_id( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """A single response stored across multiple transcript lines (one per content + block, same message.usage) must be counted once, not per line (#2340).""" + timestamp = "2026-01-01T00:00:00Z" + dup = { + "timestamp": timestamp, + "message": { + "id": "msg_dup", + "model": "claude-opus-4-1", + "usage": {"input_tokens": 100, "output_tokens": 10}, + }, + } + other = { + "timestamp": timestamp, + "message": { + "id": "msg_other", + "model": "claude-opus-4-1", + "usage": {"input_tokens": 5, "output_tokens": 2}, + }, + } + noid = { + "timestamp": timestamp, + "message": {"model": "claude-opus-4-1", "usage": {"input_tokens": 1, "output_tokens": 1}}, + } + lines = [dup, dup, dup, other, noid] + transcript = tmp_path / "session.jsonl" + transcript.write_text("\n".join(json.dumps(e) for e in lines) + "\n", encoding="utf-8") + monkeypatch.setattr(session_tracking, "find_transcript_files", lambda: [transcript]) + + entry_ts = datetime.fromisoformat(timestamp.replace("Z", "+00:00")).timestamp() + tokens = session_tracking.compute_window_tokens(entry_ts - 1, entry_ts + 1) + + # msg_dup counted once (100), msg_other (5), id-less line (1) -> 106, NOT 306. + assert tokens.input == 106 + assert tokens.output == 13 + assert tokens.by_model["claude-opus-4-1"]["input"] == 106