perf(codex): stop unmarshalling the whole rollout to compute token usage

[ecgang]

Fixes #1836.

Per @gtrrz-victor's direction on the issue, this reevaluates the Codex token-usage path rather than porting Claude's slice-first pattern. Scope held to Codex-only, token-usage calculation only, same numbers out.

Where the cost was

Not the whole-file split — splitJSONL's bytes.Split/TrimSpace return subslices and copy nothing. It's json.RawMessage.UnmarshalJSON, which is *m = append((*m)[0:0], data...): it copies. So every response_item line's encrypted_content reasoning blob was scanned, validated and heap-copied on every hook, then thrown away unread. The old path measured ~280 MB/s — unmarshal-bound.

CalculateTokenUsage needs three things from all that work: the last cumulative total_token_usage at or before the boundary, the last one after it, and how many token_count events followed. Everything else parsed was waste.

What changed

Walk the rollout with a bytes.IndexByte cursor and unmarshal only lines containing the literal token_count.

The filter cannot miss an event: that string is the value of the event's own type field, so any line the envelope check would accept contains it. Every admitted line still goes through the full three-step envelope check, so per-line decisions — and therefore the numbers — are unchanged. Lines that merely mention token_count in their text are rejected there, exactly as before; nothing is trusted on the strength of the byte match alone.

Two tempting shortcuts were tried and rejected because they silently corrupt token counts:

• Defer the baseline parse (remember the last line containing token_count, unmarshal once after the walk). ~3× faster again, and wrong: a plain event_msg of another type whose text mentions token_count — an agent discussing this issue emits one — overwrites the real baseline candidate, the deferred parse fails, the baseline goes nil, and nothing is subtracted. On a fixture, truth {fresh 1000, cache 3000, out 200} reported as {2000, 7000, 300}. Both decoy shapes are now regression tests.
• Persist a cumulative baseline in SessionState + reverse-scan from EOF. fromOffset is a line number, so locating its byte boundary needs the forward walk anyway, and apiCalls still needs the delta span. What's left is invalidation surface across resume/adopt/import/compaction that fails silently when stale.

Measured

BenchmarkCalculateTokenUsage, synthetic rollout with 4 reasoning blobs per turn at 8 KB each, boundary one turn back from EOF:

session length	before	after
50 turns	6.03 ms/hook	0.216 ms
200 turns	23.9 ms/hook	0.843 ms
500 turns	59.4 ms/hook	2.11 ms

Allocation at 500 turns: 20.7 MB / 33.5k allocs → 0.64 MB / 13k allocs.

This is a constant-factor fix, not a complexity-class one, and it does not satisfy the issue's acceptance criterion as literally worded. Being explicit since the issue asked for parses "bounded by the lines added that turn (plus O(1) baseline lookup)":

• The byte scan is still linear in transcript size, so a session still costs O(N²) bytes scanned — what's removed is unmarshalling and copying them.
• Parses per hook are one per token_count event, i.e. O(turns since session start), not O(delta). Total parse work over a session drops from O(N²) to O(N).
• The only designs that hit the literal wording are the two rejected above. Happy to build the persisted-baseline version instead if you'd rather have the bound and accept its invalidation surface — that's your call, not mine.

Verification

Tests were written against the pre-change implementation and pass unchanged against both, which is what makes "same numbers out" a measurement rather than a claim:

• full offset sweep over the shared fixture, including both ends and past EOF (the nil, nil contract)
• blank- and whitespace-only-line numbering parity — offsets are in splitJSONL coordinates, which count only non-empty lines
• both decoy shapes: an event_msg of another type mentioning token_count, and the substring inside a base64 encrypted_content blob (URL-safe base64 includes _)
• malformed and truncated lines
• baseline > last (session reset/compaction), pinning that CacheReadTokens/OutputTokens stay unclamped while only InputTokens is floored at zero — deliberate existing asymmetry this change must not alter
• a CI-gated assertion that parse count tracks token_count events rather than transcript size, counting unmarshals rather than wall-clock so it's deterministic

Additionally, a throwaway differential harness ran the old and new implementations over every real Codex rollout on my machine — 481 files, 609 MB, 80,410 lines, 19,742 real token_count events, 7,696 comparisons at both ends, past EOF and 11 interior offsets per file: 0 mismatches and 0 false negatives. Every real event contained the byte marker, which is the empirical form of the safety argument above. That harness read local session data and is not part of this PR.

mise run check green (fmt, lint 0 issues, unit + integration), e2e canary 59/59 and 4/4.

Out of scope

Left alone deliberately, per your comment narrowing scope to token-usage only:

• The modified-files second full-file read. The issue's AC list includes a bytes-based SubagentAwareExtractor for Codex so turn-end extraction stops re-reading the rollout from disk. It's a real cost worth fixing, but it's a different interface with a different risk surface. Say the word and I'll file it as a follow-up.
• A pre-existing line-numbering skew, unrelated to this change but worth knowing: GetTranscriptPosition — which produces the stored offset — counts every line including blanks, while the token calculation counts only non-empty ones. Blank-free rollouts are unaffected, so nothing is broken today. Happy to fix separately. It's also why porting the Claude slice-first pattern verbatim would have misfired: transcript.SliceFromLine counts physical newlines, so it would shift the baseline.

Also unchanged: the shared pipeline's whole-file transcript copy, encrypted_content in stored checkpoints, Claude and other agents, the shared transcript package, and accounting semantics or reported fields.

Commits

• refactor(codex): one reader for the token_count envelope — review_tokens.go already had this three-step unmarshal for the entire review stream tailer, but dropped cached_input_tokens. Extracted parseRolloutTokenUsage; the tailer's parseRolloutTokenCount is now a wrapper over it. Separate commit because it touches a working, unrelated feature.
• perf(codex): stop unmarshalling the whole rollout to read two numbers — the walk, the prefilter, and the tests.

---

🤖 Generated with Claude Code