akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
2630 commits 410 branches 169 tags 142 MiB
Commit graph

2 commits

Author SHA1 Message Date
Abhay Singh
74275b7c3e
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.
 2026-08-12 00:05:04 -05:00
Gautam Sharma
793d20fb2a
fix(subscription): read newest transcript tail (#2310) 
## Description

Large Claude Code transcript files were capped by reading the first 10
MB of each append-only JSONL file. Because recent entries are appended
at the end, current-window and weighted token usage could silently omit
the newest activity.

Oversized transcripts are now read from EOF. If the capped tail begins
within a JSONL record, only that partial record is discarded. A
  complete record beginning exactly at the boundary remains included.

  ## Type of Change

  - [x] Bug fix
  - [ ] New feature
  - [ ] Breaking change
  - [ ] Documentation-only change
  - [ ] Refactoring

  ## Changes Made

- Read the newest capped transcript bytes instead of the oldest prefix.
  - Determine the tail offset using the opened file handle.
- Inspect the preceding byte to distinguish a partial record from an
exact line boundary.
  - Remove partial bytes before UTF-8 decoding.
  - Preserve existing behavior for transcripts below the 10 MB cap.
- Add direct session-tracking tests for aggregation and boundary
behavior.
  - Add an Unreleased changelog entry.

  ## Testing

  - [x] Added regression tests
  - [x] Focused tests pass
  - [x] Subscription test suite passes
  - [x] Ruff checks pass
  - [x] Mypy passes
  - [x] Changed files pass formatting checks
  - [ ] Entire repository test suite passes without baseline failures

  Commands and results:

- `uv run --extra dev --frozen pytest
tests/test_subscription_session_tracking.py -q`
    - `4 passed`
  - Subscription-focused suite
    - `53 passed`
  - `uv run --extra dev --frozen ruff check .`
    - Passed
  - `uv run --extra dev --frozen mypy headroom --ignore-missing-imports`
    - Success across 504 source files
  - Changed-file Ruff formatting
    - Passed
  - `uv run --extra dev --frozen pytest -q`
    - `9364 passed, 565 skipped, 4 failed`
    - The four existing, unrelated failures are:
      - `test_l2_appends_transform_label`
      - `test_recovery_records_sockets_and_secures_both_backups`
      - `test_dashboard_uses_cached_stats_and_lazy_history_feed_polling`
      - `test_smart_crusher_log_fallback_runs_for_valid_json`

Repository-wide `ruff format --check .` identifies pre-existing
formatting drift only in the untouched
`headroom/proxy/handlers/anthropic.py`.

  ## Real Behavior Proof

A focused reproduction created a 10,485,787-byte transcript with a
marker entry appended after the 10 MB boundary.

  Before the fix:

  ```text
{'file_bytes': 10485787, 'line_count': 1, 'newest_entry_present': False}

  After the fix:

{'file_bytes': 10485787, 'line_count': 1, 'newest_entry_present': True}

  The regression tests additionally verify that:

1. Recent token usage beyond the cap contributes to raw and weighted
totals.
  2. A partial initial JSONL record is discarded.
  3. A complete record exactly at the tail boundary is preserved.
  4. Small transcripts retain their existing behavior.

  Environment: macOS arm64, CPython 3.12.13.

Not tested: mutation of the transcript during the individual file read
by a live Claude Code process. Reads remain bounded to a single recent
  snapshot.

  ## Review Readiness

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

  ## Checklist

- [x] The change follows existing project style and error-handling
conventions.
  - [x] Tests cover the reported failure and relevant boundary cases.
  - [x] The 10 MB memory/read cap remains enforced.
  - [x] No unrelated files or formatting changes are included.
  - [x] No temporary logging or debug code remains.
  - [x] The changelog has been updated.

  ## Additional Notes

The four full-suite failures listed above occur outside the modified
subscription code and are unrelated to this PR. All tests covering
  transcript reading and subscription tracking pass.

---------

Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
 2026-07-16 14:39:10 -07:00