perf(codex): incremental turn-end token calc to fix O(N²) reparse#1840
Draft
gtrrz-victor wants to merge 1 commit into
Draft
perf(codex): incremental turn-end token calc to fix O(N²) reparse#1840gtrrz-victor wants to merge 1 commit into
gtrrz-victor wants to merge 1 commit into
Conversation
Codex reports cumulative token totals, so the old turn-end token calc re-parsed the entire rollout from line 0 every hook — O(N) per hook, O(N²) per session, amplified by large encrypted_content blobs. Long sessions visibly stalled at turn end. Add agent.IncrementalTokenCalculator: scan only the lines added since the last checkpoint, baselining off a cumulative snapshot persisted in SessionState.CumulativeTokenBaseline. Codex is append-only with a per-rollout session id and fires no compaction hook, so the snapshot always indexes stable content; cold start or a shrunk transcript falls back to a one-time full scan. Deltas are provably identical to the full scan. Also implement SubagentAwareExtractor so turn-end file extraction runs on the in-memory transcript bytes instead of a second disk read. Per-hook cost with a warm baseline drops from O(N) to O(delta): allocs flat at 32/op across 100→10000 turns (was 260k). Fixes #1836. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Entire-Checkpoint: 01KY7AMC577ZGBN3J6QYQHK2P0
Contributor
There was a problem hiding this comment.
Pull request overview
This PR improves turn-end performance for the Codex agent by introducing an incremental token-usage calculation path that avoids re-parsing the full cumulative rollout transcript on every checkpoint, and by enabling in-memory transcript parsing for modified-file extraction.
Changes:
- Added a persisted
CumulativeTokenBaselinesnapshot in session state and a newagent.IncrementalTokenCalculatorinterface to support incremental token deltas for cumulative-count transcripts (Codex). - Updated lifecycle turn-end handling to prefer incremental token calculation when available and persist the updated baseline after
SaveStep. - Updated Codex to implement
IncrementalTokenCalculatorandSubagentAwareExtractor, and added/updated unit & perf validation tests to lock in correctness and measure the win.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| cmd/entire/cli/session/state.go | Persists a cumulative token baseline snapshot in session state for incremental calc. |
| cmd/entire/cli/lifecycle.go | Uses incremental token calc when supported and persists the next baseline post-SaveStep. |
| cmd/entire/cli/agent/token_snapshot.go | Introduces CumulativeTokenSnapshot and the IncrementalTokenCalculator interface. |
| cmd/entire/cli/agent/codex/transcript.go | Implements incremental token calculation and a bytes-based modified-file extractor for Codex. |
| cmd/entire/cli/agent/codex/perf_validation_test.go | Adds perf/scenario tests documenting the original O(N²) behavior and the post-fix benchmarks/guards. |
| cmd/entire/cli/agent/codex/incremental_test.go | Adds unit tests asserting incremental == full scan and bounded parsing behavior. |
| cmd/entire/cli/agent/capabilities.go | Adds AsIncrementalTokenCalculator helper for built-in-only capability detection. |
| CLAUDE.md | Documents the new incremental token calc behavior in the architecture notes. |
Comment on lines
+927
to
+935
| if inc, ok := agent.AsIncrementalTokenCalculator(ag); ok { | ||
| // Cumulative-count agents (Codex): parse only lines added since the | ||
| // last checkpoint, baselining off the persisted cumulative snapshot. | ||
| // nextTokenBaseline is persisted AFTER SaveStep (below), since SaveStep | ||
| // may reinitialize session state. | ||
| tokenUsage, nextTokenBaseline = computeIncrementalTokenUsage(ctx, logCtx, inc, sessionID, transcriptData, transcriptLinesAtStart) | ||
| } else { | ||
| tokenUsage = agent.CalculateTokenUsage(ctx, ag, transcriptData, transcriptLinesAtStart, subagentsDir) | ||
| } |
Comment on lines
275
to
+278
| func (c *CodexAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) { | ||
| var baselineUsage *tokenUsageData // last token_count at or before offset | ||
| var lastUsage *tokenUsageData // last token_count after offset | ||
| usage, _ := scanTokenDelta(transcriptData, 0, fromOffset, nil) | ||
| return usage, nil | ||
| } |
Comment on lines
+388
to
+405
| func (c *CodexAgent) ExtractAllModifiedFiles(transcriptData []byte, fromOffset int, _ string) ([]string, error) { | ||
| seen := make(map[string]struct{}) | ||
| var files []string | ||
| lineNum := 0 | ||
| for _, lineData := range splitJSONL(transcriptData) { | ||
| lineNum++ | ||
| if lineNum <= fromOffset { | ||
| continue | ||
| } | ||
| for _, f := range extractFilesFromLine(lineData) { | ||
| if _, ok := seen[f]; !ok { | ||
| seen[f] = struct{}{} | ||
| files = append(files, f) | ||
| } | ||
| } | ||
| } | ||
| return files, nil | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
https://entire.io/gh/entireio/cli/trails/929
Fixes #1836.
Problem
Codex reports cumulative token totals (
event_msg→token_count→total_token_usage), so the turn-end token calc re-parsed the entire rollout from line 0 on every hook — O(N) per hook, O(N²) per session — amplified by largeencrypted_contentreasoning blobs. Long sessions visibly stalled at each turn end. Claude Code avoids this by slicing (its counts are per-message, not cumulative).Fix
agent.IncrementalTokenCalculator(new, built-in-only capability): scan only the transcript lines added since the last checkpoint, baselining off a cumulative snapshot persisted inSessionState.CumulativeTokenBaseline.snapshot.LineCount > fromOffset) falls back to a one-time full scan.CalculateTokenUsage— only the amount of parsing changes.SubagentAwareExtractor(no subagents) so turn-end modified-file extraction runs on the in-memory transcript bytes instead of a second full-file disk read.CalculateTokenUsage(full-scan) kept byte-for-byte identical for back-compat and the fallback path.Measured win (per-hook, warm baseline)
Allocs flat at 32/op regardless of session length — the definitive O(delta) proof (residual ns/op is
SliceFromLine's cheap byte-walk, not JSON parsing).Safety (no regressions)
fromOffsetas the old disk path and produces identical results for JSONL (no blank lines); the skip-checkpoint decision also relies on git-status new/deleted + merged modified, independent of transcript parsing. No missed changes → no spurious carry-over branches.Test plan
incremental_test.go): incremental == full scan across turns/cached/mid-offset; cold-start reusable baseline; resumed/stale-offset fallback;DoesNotParseHistory(corrupted pre-offset lines ignored — proves bounded parsing); bytes extractor == disk extractor.perf_validation_test.go) documenting the root cause + benchmarks.mise run fmt/ lint (v2.11.3): clean · unit (codex/agent/session/cli): pass · integration: 437 pass · Vogon canary: 4/4.🤖 Generated with Claude Code
Note
Cursor Bugbot is generating a summary for commit 406904c. Configure here.