fix(subscription): dedup transcript usage by message id (#2340 token inflation) (#2408)

## Description

Addresses the usage-inflation part of #2340.

`compute_window_tokens` (`headroom/subscription/session_tracking.py`)
sums `message.usage` for every transcript line whose timestamp falls in
the window:

```python
for line in _read_transcript_lines(path):
    ...
    usage = msg.get("usage")
    if not usage:
        continue
    _add_usage_to_tokens(totals, usage)
```

But Claude Code can store a single assistant response across **multiple
transcript lines** (e.g. one entry per content block), and each of those
lines carries the **same request-level `message.usage`**. Summing per
line therefore multiplies that one response's tokens by its block count.
#2340 observed a single 420,609-input-token response counted **19
times** (~8M attributed input tokens from one record), which is most of
the reported window-total inflation.

## Fix

Count each response's usage once, keyed by the Anthropic `message.id`
(unique per response):

```python
seen_message_ids: set[str] = set()
...
    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)
```

Entries without a `message.id` keep the previous per-line behavior, so
this only ever removes true duplicates: a response is de-duplicated only
when the exact same (unique) message id appears more than once, and
distinct responses are unaffected.

Scope: this fixes the token-accounting inflation only. The separate
retry-amplification / `tool_search_tool_result` SSE-502 behavior
described in the same issue is a different code path and is not touched
here.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/subscription/session_tracking.py`: dedup usage accumulation
by `message.id` in `compute_window_tokens`.
- `tests/test_subscription_session_tracking.py`: add a test where one
response is stored across three lines (plus a distinct response and an
id-less line) and assert its usage is counted once.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
$ uvx ruff@0.15.17 check headroom/subscription/session_tracking.py tests/test_subscription_session_tracking.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/subscription/session_tracking.py
Success: no issues found in 1 source file
# session_tracking is import-light, so I ran the exact logic against the real
# module in the project venv (uv sync): a response stored on 3 lines yields
# input=106 (100 once + 5 + 1 id-less), not 306.
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12, project venv (`uv sync --extra
proxy`), `uvx ruff@0.15.17` / `uvx mypy@1.20.2`.
- Exact command / steps: wrote a transcript with `msg_dup` repeated on
three lines (same usage, 100/10), one distinct `msg_other` (5/2), and
one id-less line (1/1); called the real `compute_window_tokens` over the
window.
- Observed result: `input == 106` and `output == 13` (the duplicated
response counted once, the id-less line still counted); the pre-fix code
would report `input == 306`. Because `session_tracking` has no heavy
imports, this ran against the actual module.
- Not tested: a live Claude Code transcript with real multi-block
responses; the added unit test reproduces the multi-line-per-response
shape.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

`session_tracking` is a light module, so I verified against the real
code in the venv (output above) in addition to the added test. This is
deliberately scoped to the usage-double-count sub-part of #2340; the
retry-amplification/SSE side is separate and untouched. Keyed on
`message.id` so it is safe by construction: no id or a unique id behaves
exactly as before.
This commit is contained in:
Abhay Singh 2026-08-12 10:35:04 +05:30 committed by GitHub
parent 1db6d88ab4
commit 74275b7c3e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 55 additions and 0 deletions

View file

@ -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")

View file

@ -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