From 406904c64eabb2c7dabb7a1fc9bfef9ff14cc176 Mon Sep 17 00:00:00 2001 From: Victor Gutierrez Calderon Date: Thu, 23 Jul 2026 13:09:05 +0200 Subject: [PATCH 1/2] =?UTF-8?q?perf(codex):=20incremental=20turn-end=20tok?= =?UTF-8?q?en=20calc=20to=20fix=20O(N=C2=B2)=20reparse?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Entire-Checkpoint: 01KY7AMC577ZGBN3J6QYQHK2P0 --- CLAUDE.md | 1 + cmd/entire/cli/agent/capabilities.go | 8 + .../cli/agent/codex/incremental_test.go | 189 ++++++++++++ .../cli/agent/codex/perf_validation_test.go | 275 ++++++++++++++++++ cmd/entire/cli/agent/codex/transcript.go | 105 ++++++- cmd/entire/cli/agent/token_snapshot.go | 46 +++ cmd/entire/cli/lifecycle.go | 45 ++- cmd/entire/cli/session/state.go | 8 + 8 files changed, 665 insertions(+), 12 deletions(-) create mode 100644 cmd/entire/cli/agent/codex/incremental_test.go create mode 100644 cmd/entire/cli/agent/codex/perf_validation_test.go create mode 100644 cmd/entire/cli/agent/token_snapshot.go diff --git a/CLAUDE.md b/CLAUDE.md index b3b9408297..174807bda3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -595,6 +595,7 @@ The manual-commit strategy (`manual_commit*.go`) does not modify the active bran - Rewind restores files from shadow branch commit tree (does not use `git reset`) - **Location-independent transcript resolution** - transcript paths are always computed dynamically from the current repo location (via `agent.GetSessionDir` + `agent.ResolveSessionFile`), never stored in checkpoint metadata. This ensures restore/rewind works after repo relocation or across machines. - **Token usage scoping** - `SessionState.TokenUsage` is the session-wide total used by `entire status`; `SessionState.CheckpointTokenUsage` is the pending checkpoint delta since the last condensation. Checkpoint metadata must stay scoped to `CheckpointTranscriptStart` or the pending checkpoint delta. Cursor tokens come only from stop-hook payloads, while Copilot CLI can also backfill full-session totals from `session.shutdown`. +- **Incremental token calc for cumulative-count agents** - agents whose transcript reports *cumulative* token totals (Codex: `event_msg`→`token_count`→`total_token_usage`) implement `agent.IncrementalTokenCalculator`. At turn-end the lifecycle scans only the transcript lines added since the last checkpoint, baselining off `SessionState.CumulativeTokenBaseline` (the last cumulative snapshot, persisted after `SaveStep`) instead of re-parsing the whole rollout every hook (the old O(N) per-hook / O(N²) per-session cost). Cold start or a shrunk/reset transcript (`snapshot.LineCount > fromOffset`) falls back to a one-time full scan; the delta produced always equals the full-scan `CalculateTokenUsage`. Codex also implements `SubagentAwareExtractor` (no subagents) so turn-end modified-file extraction runs on the in-memory transcript bytes rather than a second disk read. - Tracks session state in `.git/entire-sessions/` (shared across worktrees) - **Shadow branch migration** - if user does stash/pull/rebase (HEAD changes without commit), shadow branch is automatically moved to new base commit - **Orphaned branch cleanup** - if a shadow branch exists without a corresponding session state file, it is automatically reset when a new session starts diff --git a/cmd/entire/cli/agent/capabilities.go b/cmd/entire/cli/agent/capabilities.go index 77084156bb..50a15f6994 100644 --- a/cmd/entire/cli/agent/capabilities.go +++ b/cmd/entire/cli/agent/capabilities.go @@ -127,6 +127,14 @@ func AsSubagentAwareExtractor(ag Agent) (SubagentAwareExtractor, bool) { return declaredCapability[SubagentAwareExtractor](ag, func(c DeclaredCaps) bool { return c.SubagentAwareExtractor }) } +// AsIncrementalTokenCalculator returns the agent as IncrementalTokenCalculator +// if it implements the interface. Incremental token calculation is a built-in +// optimization (for agents with cumulative-count transcripts like Codex) with +// no external-protocol equivalent, so it resolves by type assertion alone. +func AsIncrementalTokenCalculator(ag Agent) (IncrementalTokenCalculator, bool) { + return builtinCapability[IncrementalTokenCalculator](ag) +} + // AsSessionBaseDirProvider returns the agent as SessionBaseDirProvider if it implements // the interface. No capability declaration is needed since this is a built-in-only feature // (external agents use the agent binary's own session resolution). diff --git a/cmd/entire/cli/agent/codex/incremental_test.go b/cmd/entire/cli/agent/codex/incremental_test.go new file mode 100644 index 0000000000..c8f324db63 --- /dev/null +++ b/cmd/entire/cli/agent/codex/incremental_test.go @@ -0,0 +1,189 @@ +package codex + +import ( + "strings" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/transcript" + "github.com/stretchr/testify/require" +) + +// lineCount counts non-empty JSONL lines the same way GetTranscriptPosition does +// for well-formed rollouts (no blank lines). +func lineCount(data []byte) int { return len(splitJSONL(data)) } + +// TestIncremental_MatchesFullScan replays every turn-end hook of a growing +// session and asserts the incremental calc produces byte-identical numbers to +// the full-scan CalculateTokenUsage — carrying the persisted snapshot forward +// like the lifecycle does. +func TestIncremental_MatchesFullScan(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + + const turns = 8 + var prior *agent.CumulativeTokenSnapshot + fromOffset := 0 // turn 1 starts at line 0 + + for turn := 1; turn <= turns; turn++ { + data, _ := buildRollout(turn, false, 0) + + full, err := ag.CalculateTokenUsage(data, fromOffset) + require.NoError(t, err) + + inc, next, err := ag.CalculateTokenUsageIncremental(data, fromOffset, prior) + require.NoError(t, err) + require.Equal(t, full, inc, "turn %d: incremental delta must equal full scan", turn) + + // Next turn starts where this transcript ended. + prior = next + fromOffset = lineCount(data) + } +} + +// TestIncremental_MatchesFullScan_WithCachedTokens uses the shipped fixture that +// includes cached-input tokens and a mid-session offset. +func TestIncremental_MatchesFullScan_WithCachedTokens(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + data := []byte(sampleRollout) + + for _, offset := range []int{0, 4, 8, 11} { + full, err := ag.CalculateTokenUsage(data, offset) + require.NoError(t, err) + + // Cold start (no prior): must fall back to a full scan and match exactly. + inc, next, err := ag.CalculateTokenUsageIncremental(data, offset, nil) + require.NoError(t, err) + require.Equal(t, full, inc, "offset %d: cold-start incremental must equal full scan", offset) + require.NotNil(t, next) + require.Equal(t, lineCount(data), next.LineCount, "offset %d: snapshot line count", offset) + } +} + +// TestIncremental_ColdStartProducesReusableBaseline verifies a cold-start call +// yields a snapshot that a subsequent incremental call can baseline off to get +// the correct per-checkpoint delta. +func TestIncremental_ColdStartProducesReusableBaseline(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + + // Turn 1: 2 token_count lines, cold start. + turn1, _ := buildRollout(2, false, 0) + _, snap, err := ag.CalculateTokenUsageIncremental(turn1, 0, nil) + require.NoError(t, err) + require.NotNil(t, snap) + require.Equal(t, lineCount(turn1), snap.LineCount) + + // Turn 2: one more token_count line appended. + turn2, _ := buildRollout(3, false, 0) + fromOffset := lineCount(turn1) + + inc, _, err := ag.CalculateTokenUsageIncremental(turn2, fromOffset, snap) + require.NoError(t, err) + full, err := ag.CalculateTokenUsage(turn2, fromOffset) + require.NoError(t, err) + require.Equal(t, full, inc, "warm incremental delta must equal full scan") +} + +// TestIncremental_ResumedOffsetFallsBackToFullScan covers a stale baseline +// (LineCount past the current fromOffset, as after a transcript reset/resume): +// the calc must ignore it and full-scan, still producing correct numbers. +func TestIncremental_ResumedOffsetFallsBackToFullScan(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + data, _ := buildRollout(3, false, 0) + fromOffset := 3 + + stale := &agent.CumulativeTokenSnapshot{LineCount: 9999, InputTokens: 1, CachedInputTokens: 1, OutputTokens: 1} + inc, _, err := ag.CalculateTokenUsageIncremental(data, fromOffset, stale) + require.NoError(t, err) + full, err := ag.CalculateTokenUsage(data, fromOffset) + require.NoError(t, err) + require.Equal(t, full, inc, "stale baseline must be ignored (full-scan fallback)") +} + +// TestIncremental_DoesNotParseHistory is the acceptance criterion: parse work is +// bounded by lines added since the last checkpoint, not the whole session. +// We corrupt every pre-offset line with bogus token_counts. A full scan of the +// corrupted transcript changes the numbers (it reads the bogus baseline); the +// incremental call, given a valid prior baseline, slices the history off and is +// UNAFFECTED — proving it never parsed those lines. +func TestIncremental_DoesNotParseHistory(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + + clean, _ := buildRollout(5, false, 0) + fromOffset := lineCount(clean) + + // Correct baseline as of fromOffset = the last cumulative in `clean`. + _, prior, err := ag.CalculateTokenUsageIncremental(clean, fromOffset, nil) + require.NoError(t, err) + require.NotNil(t, prior) + + // Append one more turn to `clean` (the new delta the next hook sees). + cleanNext, _ := buildRollout(6, false, 0) + tail := transcript.SliceFromLine(cleanNext, fromOffset) + + // Build a corrupted transcript: same line count for lines <= fromOffset, but + // every pre-offset token_count carries absurd values, followed by the real tail. + var corrupt strings.Builder + corrupt.WriteString(`{"timestamp":"t","type":"session_meta","payload":{"id":"x"}}`) + corrupt.WriteByte('\n') + for i := 2; i <= fromOffset; i++ { + corrupt.WriteString(tokenCountLine(9_000_000, 8_000_000, 900_000)) + corrupt.WriteByte('\n') + } + corrupt.Write(tail) + corrupted := []byte(corrupt.String()) + require.Equal(t, fromOffset, lineCount(corrupted)-lineCount(tail), + "corrupted history must have the same pre-offset line count") + + // A FULL scan of the corrupted transcript reads the bogus baseline → wrong. + fullCorrupt, err := ag.CalculateTokenUsage(corrupted, fromOffset) + require.NoError(t, err) + + // The INCREMENTAL call with the valid prior slices off history → correct, + // identical to the clean full scan. + incCorrupt, _, err := ag.CalculateTokenUsageIncremental(corrupted, fromOffset, prior) + require.NoError(t, err) + cleanFull, err := ag.CalculateTokenUsage(cleanNext, fromOffset) + require.NoError(t, err) + + require.Equal(t, cleanFull, incCorrupt, + "incremental result must ignore corrupted history (only delta lines parsed)") + require.NotEqual(t, fullCorrupt, incCorrupt, + "a full scan WOULD have been corrupted — proving incremental skipped history") +} + +// TestExtractAllModifiedFiles_FromBytes verifies the bytes-based extractor +// returns the same files as the disk-based ExtractModifiedFilesFromOffset, so +// the turn-end pipeline avoids the second full-file disk read. +func TestExtractAllModifiedFiles_FromBytes(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + data := []byte(sampleRollout) + path := writeSampleRollout(t) + + for _, offset := range []int{0, 9} { + fromBytes, err := ag.ExtractAllModifiedFiles(data, offset, "") + require.NoError(t, err) + fromDisk, _, err := ag.ExtractModifiedFilesFromOffset(path, offset) + require.NoError(t, err) + require.ElementsMatch(t, fromDisk, fromBytes, "offset %d: bytes extractor must match disk extractor", offset) + } +} + +// TestCodex_DeclaresBytesCapabilities locks in that Codex now routes both token +// calc and file extraction through in-memory bytes (no disk reread, incremental +// tokens) — the opposite of the pre-fix validation guard H3. +func TestCodex_DeclaresBytesCapabilities(t *testing.T) { + t.Parallel() + ag := NewCodexAgent() + + _, incOK := agent.AsIncrementalTokenCalculator(ag) + require.True(t, incOK, "Codex must be an IncrementalTokenCalculator") + + _, subOK := agent.AsSubagentAwareExtractor(ag) + require.True(t, subOK, "Codex must expose a bytes-based (SubagentAware) file extractor") +} diff --git a/cmd/entire/cli/agent/codex/perf_validation_test.go b/cmd/entire/cli/agent/codex/perf_validation_test.go new file mode 100644 index 0000000000..ac56aaecd3 --- /dev/null +++ b/cmd/entire/cli/agent/codex/perf_validation_test.go @@ -0,0 +1,275 @@ +package codex + +// perf_validation_test.go — SCENARIOS validating the hypotheses in issue #1836 +// (Codex checkpoint hooks O(N²) transcript reparse). +// +// These tests/benchmarks are the empirical contract that the root-cause +// analysis is correct BEFORE any fix lands. They are written to PASS against +// the current (slow) implementation: each asserts the buggy behavior exists. +// After the fix, the behavioral guards (H1, H1b) should be UPDATED to assert +// the incremental behavior; the benchmarks let you measure the win. +// +// Hypotheses: +// H1 — CalculateTokenUsage parses the whole file every hook, regardless of +// fromOffset (per-hook O(N), whole-session O(N²)). +// H1b — A naive Claude-style SliceFromLine(offset) breaks Codex numbers, +// because Codex counts are CUMULATIVE and the delta needs the +// pre-offset baseline. So the fix must persist the baseline, not just +// slice. +// H3 — Codex does NOT implement SubagentAwareExtractor, so turn-end file +// extraction falls back to the disk-reread path. +// H4 — encrypted_content reasoning blobs inflate the per-line byte cost of +// every size-linear step. + +import ( + "fmt" + "strings" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/transcript" + "github.com/stretchr/testify/require" +) + +// tokenCountLine builds an event_msg/token_count line with the given cumulative totals. +func tokenCountLine(input, cached, output int) string { + return fmt.Sprintf( + `{"timestamp":"t","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":%d,"cached_input_tokens":%d,"output_tokens":%d,"reasoning_output_tokens":0,"total_tokens":%d}}}}`, + input, cached, output, input+output) +} + +// fatReasoningLine is a response_item carrying a large encrypted_content blob, +// mirroring real Codex rollouts. It is not a token_count line, so it never +// contributes to the token delta — but the token calc still splits and +// top-level-unmarshals it on every hook. +func fatReasoningLine(blobBytes int) string { + blob := strings.Repeat("A", blobBytes) + return `{"timestamp":"t","type":"response_item","payload":{"type":"reasoning","summary":[{"text":"x"}],"encrypted_content":"` + blob + `"}}` +} + +// buildRollout returns a synthetic rollout of nTurns turns. Each turn is one +// token_count line with monotonically increasing cumulative counts, optionally +// preceded by a fat reasoning line. Returns the bytes and the total line count. +func buildRollout(nTurns int, withBlobs bool, blobBytes int) ([]byte, int) { + var b strings.Builder + b.WriteString(`{"timestamp":"t","type":"session_meta","payload":{"id":"x"}}`) + b.WriteByte('\n') + lines := 1 + for i := 1; i <= nTurns; i++ { + if withBlobs { + b.WriteString(fatReasoningLine(blobBytes)) + b.WriteByte('\n') + lines++ + } + // cumulative: grows every turn + b.WriteString(tokenCountLine(1000*i, 800*i, 100*i)) + b.WriteByte('\n') + lines++ + } + return []byte(b.String()), lines +} + +// --------------------------------------------------------------------------- +// H1 — the token calc consults pre-offset lines, proving whole-file reparse. +// --------------------------------------------------------------------------- + +// TestValidate_H1_TokenCalcReadsPreOffsetBaseline proves the calc scans lines +// BEFORE fromOffset. We set fromOffset so only the final token_count line is +// "after" it; a correct delta can only be produced by having read the earlier +// (pre-offset) cumulative baseline. If the implementation sliced at the offset +// (touched only new lines), baseline would be nil and the delta would equal the +// full final cumulative value instead of the per-turn delta. +func TestValidate_H1_TokenCalcReadsPreOffsetBaseline(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + + // 3 token_count lines at file lines 2, 3, 4 (line 1 is session_meta). + data, total := buildRollout(3, false, 0) + require.Equal(t, 4, total) + + // fromOffset=3 → only the line-4 token_count (cumulative {3000,2400,300}) + // is "after"; the line-3 baseline is {2000,1600,200}. + usage, err := ag.CalculateTokenUsage(data, 3) + require.NoError(t, err) + require.NotNil(t, usage) + + // baseline line3 {2000,1600,200}, last line4 {3000,2400,300}. + // delta = 1000 input, 800 cache, 100 output. fresh input = 1000-800 = 200. + require.Equal(t, 100, usage.OutputTokens, "output delta uses pre-offset baseline") + require.Equal(t, 800, usage.CacheReadTokens, "cache delta uses pre-offset baseline") + require.Equal(t, 200, usage.InputTokens, "fresh input delta uses pre-offset baseline") + require.Equal(t, 1, usage.APICallCount, "only one token_count line is after the offset") + + // The smoking gun: baseline came from a line BEFORE the offset. A sliced + // (delta-only) implementation cannot produce these numbers — see H1b. +} + +// --------------------------------------------------------------------------- +// H1b — naive Claude-style slicing breaks Codex numbers (baseline lost). +// This validates the fix CONSTRAINT: persist the baseline, don't just slice. +// --------------------------------------------------------------------------- + +func TestValidate_H1b_NaiveSliceBreaksCumulativeDelta(t *testing.T) { + t.Parallel() + ag := &CodexAgent{} + + data, _ := buildRollout(3, false, 0) + fromOffset := 3 + + correct, err := ag.CalculateTokenUsage(data, fromOffset) + require.NoError(t, err) + require.NotNil(t, correct) + + // Simulate the Claude-style approach: slice to the checkpoint window, then + // run the calc from offset 0 on the slice. This is what a naive port would + // do — and it is WRONG for Codex because the baseline lives before the slice. + sliced := transcript.SliceFromLine(data, fromOffset) + naive, err := ag.CalculateTokenUsage(sliced, 0) + require.NoError(t, err) + require.NotNil(t, naive) + + // The naive slice has no baseline, so it reports the full cumulative value + // as this checkpoint's delta — inflated vs correct. + require.NotEqual(t, correct.OutputTokens, naive.OutputTokens, + "naive slicing must diverge — proves baseline persistence is required") + require.Greater(t, naive.OutputTokens, correct.OutputTokens, + "naive slice over-counts because it lost the cumulative baseline") + t.Logf("correct delta output=%d, naive-slice output=%d (inflated)", + correct.OutputTokens, naive.OutputTokens) +} + +// --------------------------------------------------------------------------- +// H3 — Codex has no bytes-based subagent-aware extractor → disk-reread path. +// --------------------------------------------------------------------------- + +// H3 (post-fix): Codex now exposes a bytes-based file extractor, so lifecycle no +// longer falls back to the disk-reread path. Pre-fix this asserted the opposite; +// see TestCodex_DeclaresBytesCapabilities in incremental_test.go for the guard. +func TestValidate_H3_CodexNowHasBytesExtractor(t *testing.T) { + t.Parallel() + + _, ok := agent.AsSubagentAwareExtractor(NewCodexAgent()) + require.True(t, ok, + "Codex must be a SubagentAwareExtractor after the fix — turn-end file "+ + "extraction runs on in-memory bytes instead of a 2nd full-file disk read") +} + +// --------------------------------------------------------------------------- +// Benchmarks — empirical proof of the O(N)/O(N²) cost and the blob amplifier. +// Run: go test ./cmd/entire/cli/agent/codex/ -run '^$' -bench BenchmarkValidate -benchmem +// --------------------------------------------------------------------------- + +// benchSink keeps benchmark results reachable so the compiler can't elide the +// call under test. +var benchSink *agent.TokenUsage + +// H1 per-hook cost: fromOffset pinned to the LAST turn (delta = 1 line). If the +// calc were incremental this would be ~flat across N; today it grows with N. +func BenchmarkValidate_H1_PerHook_LastTurnOnly(b *testing.B) { + for _, n := range []int{100, 1000, 10000} { + data, total := buildRollout(n, false, 0) + offset := total - 1 // only the final token_count is "after" + b.Run(fmt.Sprintf("turns=%d", n), func(b *testing.B) { + ag := &CodexAgent{} + b.ReportAllocs() + b.ResetTimer() + for range b.N { + usage, err := ag.CalculateTokenUsage(data, offset) + if err != nil { + b.Fatal(err) + } + benchSink = usage + } + }) + } +} + +// H1 whole-session cost: replay every turn-end hook over a growing prefix, as +// the real pipeline does. Total work is the sum of per-hook O(prefix) → O(N²). +func BenchmarkValidate_H1_WholeSession(b *testing.B) { + for _, n := range []int{200, 1000} { + data, _ := buildRollout(n, false, 0) + // Precompute per-turn byte prefixes (line boundaries). + prefixes := make([][]byte, 0, n) + nl := 0 + for i, c := range data { + if c == '\n' { + nl++ + if nl >= 2 { // after session_meta, snapshot each turn boundary + prefixes = append(prefixes, data[:i+1]) + } + } + } + b.Run(fmt.Sprintf("turns=%d", n), func(b *testing.B) { + ag := &CodexAgent{} + b.ReportAllocs() + b.ResetTimer() + for range b.N { + off := 0 + for _, p := range prefixes { + usage, err := ag.CalculateTokenUsage(p, off) + if err != nil { + b.Fatal(err) + } + benchSink = usage + off = strings.Count(string(p), "\n") - 1 + } + } + }) + } +} + +// H4 constant-factor amplifier: same number of token_count lines, but one +// variant interleaves fat encrypted_content reasoning blobs. Compare ns/op and +// B/op — the blob variant is markedly heavier per hook. +func BenchmarkValidate_H4_EncryptedContentAmplifier(b *testing.B) { + const turns = 1000 + const blob = 4096 // ~4KB per reasoning line, conservative vs real rollouts + + plain, _ := buildRollout(turns, false, 0) + fat, _ := buildRollout(turns, true, blob) + + run := func(name string, data []byte) { + b.Run(name, func(b *testing.B) { + ag := &CodexAgent{} + b.ReportAllocs() + b.SetBytes(int64(len(data))) + b.ResetTimer() + for range b.N { + usage, err := ag.CalculateTokenUsage(data, 0) + if err != nil { + b.Fatal(err) + } + benchSink = usage + } + }) + } + run("no_blobs", plain) + run("with_blobs", fat) +} + +// After-fix proof: per-hook incremental cost with a warm baseline stays flat as +// the session grows, vs the full-scan growth in BenchmarkValidate_H1_PerHook. +func BenchmarkValidate_Fixed_IncrementalPerHook(b *testing.B) { + ag := &CodexAgent{} + for _, n := range []int{100, 1000, 10000} { + data, _ := buildRollout(n, false, 0) + prev, _ := buildRollout(n-1, false, 0) + fromOffset := lineCount(prev) + _, prior, err := ag.CalculateTokenUsageIncremental(prev, lineCount(prev)-1, nil) + if err != nil { + b.Fatal(err) + } + b.Run(fmt.Sprintf("turns=%d", n), func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for range b.N { + usage, _, incErr := ag.CalculateTokenUsageIncremental(data, fromOffset, prior) + if incErr != nil { + b.Fatal(incErr) + } + benchSink = usage + } + }) + } +} diff --git a/cmd/entire/cli/agent/codex/transcript.go b/cmd/entire/cli/agent/codex/transcript.go index 16b04f557a..b94a91353c 100644 --- a/cmd/entire/cli/agent/codex/transcript.go +++ b/cmd/entire/cli/agent/codex/transcript.go @@ -14,12 +14,15 @@ import ( "time" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/transcript" ) // Compile-time interface assertions. var ( _ agent.TranscriptAnalyzer = (*CodexAgent)(nil) _ agent.TokenCalculator = (*CodexAgent)(nil) + _ agent.IncrementalTokenCalculator = (*CodexAgent)(nil) + _ agent.SubagentAwareExtractor = (*CodexAgent)(nil) _ agent.PromptExtractor = (*CodexAgent)(nil) _ agent.RestoredSessionPathResolver = (*CodexAgent)(nil) ) @@ -266,13 +269,63 @@ func classifyApplyPatchPaths(input string) (added, modified, deleted []string) { // CalculateTokenUsage computes token usage from the transcript starting at the given line offset. // Codex reports cumulative total_token_usage, so we compute the delta between the last // token_count at/before the offset and the last token_count after the offset. +// +// This full-scan form parses every line from 0; prefer CalculateTokenUsageIncremental +// on the turn-end hot path, which parses only lines added since the last checkpoint. 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 +} + +// CalculateTokenUsageIncremental computes the same delta as CalculateTokenUsage but, +// given the cumulative baseline persisted at the previous checkpoint, scans only the +// transcript lines added since then. Codex rollouts are append-only with no token_count +// events between turns, so the persisted snapshot remains a valid pre-offset baseline +// and history need not be re-parsed. Falls back to a full scan on cold start (prior nil) +// or when the transcript shrank/reset (prior.LineCount > fromOffset). +func (c *CodexAgent) CalculateTokenUsageIncremental(transcriptData []byte, fromOffset int, prior *agent.CumulativeTokenSnapshot) (*agent.TokenUsage, *agent.CumulativeTokenSnapshot, error) { + startLine := 0 + var seed *tokenUsageData + if prior != nil && prior.LineCount <= fromOffset { + startLine = prior.LineCount + seed = &tokenUsageData{ + InputTokens: prior.InputTokens, + CachedInputTokens: prior.CachedInputTokens, + OutputTokens: prior.OutputTokens, + } + } + + scanData := transcriptData + if startLine > 0 { + scanData = transcript.SliceFromLine(transcriptData, startLine) + } + + usage, latest := scanTokenDelta(scanData, startLine, fromOffset, seed) + + // Persist the final cumulative snapshot at the current transcript length so + // the next checkpoint can baseline off it without rescanning history. + next := &agent.CumulativeTokenSnapshot{LineCount: startLine + len(splitJSONL(scanData))} + if latest != nil { + next.InputTokens = latest.InputTokens + next.CachedInputTokens = latest.CachedInputTokens + next.OutputTokens = latest.OutputTokens + } + return usage, next, nil +} + +// scanTokenDelta scans token_count events in data, numbering lines starting at +// startLineNum+1 (the absolute line number of data's first line). seed is the +// cumulative baseline as of startLineNum (nil for a from-scratch scan). It returns +// the checkpoint delta (nil when no token_count exists after fromOffset, matching the +// original behavior) and the last cumulative snapshot observed (seed if none seen). +func scanTokenDelta(data []byte, startLineNum, fromOffset int, seed *tokenUsageData) (usage *agent.TokenUsage, latest *tokenUsageData) { + baselineUsage := seed // last token_count at or before offset + latest = seed // last token_count seen overall (for the next baseline) + var lastUsage *tokenUsageData // last token_count after offset apiCalls := 0 - lineNum := 0 + lineNum := startLineNum - for _, lineData := range splitJSONL(transcriptData) { + for _, lineData := range splitJSONL(data) { lineNum++ var line rolloutLine @@ -294,6 +347,7 @@ func (c *CodexAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) continue } + latest = info.TotalTokenUsage if lineNum <= fromOffset { baselineUsage = info.TotalTokenUsage } else { @@ -303,10 +357,10 @@ func (c *CodexAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) } if lastUsage == nil { - return nil, nil //nolint:nilnil // no usage data found + return nil, latest } - // Subtract baseline to get the delta for this checkpoint range + // Subtract baseline to get the delta for this checkpoint range. inputTokens := lastUsage.InputTokens cacheReadTokens := lastUsage.CachedInputTokens outputTokens := lastUsage.OutputTokens @@ -316,17 +370,46 @@ func (c *CodexAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) outputTokens -= baselineUsage.OutputTokens } - freshInputTokens := inputTokens - cacheReadTokens - if freshInputTokens < 0 { - freshInputTokens = 0 - } + freshInputTokens := max(inputTokens-cacheReadTokens, 0) return &agent.TokenUsage{ InputTokens: freshInputTokens, CacheReadTokens: cacheReadTokens, OutputTokens: outputTokens, APICallCount: apiCalls, - }, nil + }, latest +} + +// ExtractAllModifiedFiles implements agent.SubagentAwareExtractor for Codex. Codex +// has no subagents, so subagentsDir is ignored; the value of implementing this +// interface is that the turn-end pipeline extracts modified files from the in-memory +// transcript bytes it already holds, instead of falling back to a second full-file +// disk read via ExtractModifiedFilesFromOffset. +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 +} + +// CalculateTotalTokenUsage implements agent.SubagentAwareExtractor. Codex has no +// subagents, so this is just the main-transcript delta. The turn-end pipeline uses +// CalculateTokenUsageIncremental directly (it needs session state for the baseline); +// this exists to satisfy the interface and as a stateless fallback. +func (c *CodexAgent) CalculateTotalTokenUsage(transcriptData []byte, fromOffset int, _ string) (*agent.TokenUsage, error) { + return c.CalculateTokenUsage(transcriptData, fromOffset) } // ExtractPrompts returns user prompts from the transcript starting at the given offset. diff --git a/cmd/entire/cli/agent/token_snapshot.go b/cmd/entire/cli/agent/token_snapshot.go new file mode 100644 index 0000000000..03a3fe670a --- /dev/null +++ b/cmd/entire/cli/agent/token_snapshot.go @@ -0,0 +1,46 @@ +package agent + +// CumulativeTokenSnapshot caches an agent's running cumulative token totals at a +// known transcript position. Agents that report CUMULATIVE counts (e.g. Codex, +// whose rollout emits token_count events carrying total_token_usage since +// session start) can use it to compute a checkpoint delta by scanning only the +// transcript lines added since this snapshot, instead of re-parsing the whole +// rollout on every turn-end hook. +// +// It is persisted in session state and round-trips as JSON. +type CumulativeTokenSnapshot struct { + // LineCount is the transcript line count when this snapshot was captured + // (same counting convention as TranscriptAnalyzer.GetTranscriptPosition). + // The next incremental calc scans from here and uses it as the fast-path + // guard: the snapshot is a valid baseline only when LineCount <= the next + // fromOffset (otherwise the transcript was reset/rewound and we full-scan). + LineCount int `json:"line_count"` + // InputTokens is the cumulative total input tokens reported at LineCount. + InputTokens int `json:"input_tokens"` + // CachedInputTokens is the cumulative cached-input tokens reported at LineCount. + CachedInputTokens int `json:"cached_input_tokens"` + // OutputTokens is the cumulative output tokens reported at LineCount. + OutputTokens int `json:"output_tokens"` +} + +// IncrementalTokenCalculator is a TokenCalculator whose transcript reports +// CUMULATIVE token counts. It computes the checkpoint delta from a persisted +// baseline snapshot, scanning only the lines added since that snapshot rather +// than the whole transcript — the whole-file rescan is O(session) per hook and +// O(session^2) across a session. +type IncrementalTokenCalculator interface { + TokenCalculator + + // CalculateTokenUsageIncremental returns the checkpoint token delta for + // transcript lines after fromOffset, plus a fresh snapshot to persist for the + // next checkpoint. prior is the snapshot persisted at the previous checkpoint, + // or nil on cold start (first checkpoint, resumed/imported session, or after a + // transcript reset); when prior is nil or stale the implementation falls back + // to a full scan for that one call. + // + // CONTRACT: for any given inputs, the returned usage MUST equal + // CalculateTokenUsage(transcriptData, fromOffset) — only the amount of parsing + // differs. next is nil only when the transcript carries no cumulative token + // data at all. + CalculateTokenUsageIncremental(transcriptData []byte, fromOffset int, prior *CumulativeTokenSnapshot) (usage *TokenUsage, next *CumulativeTokenSnapshot, err error) +} diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index e89a6b0b63..0ed598c2d3 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -922,8 +922,17 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev // back to transcript-based computation, preferring SubagentAwareExtractor // to include subagent tokens. tokenUsage := event.TokenUsage + var nextTokenBaseline *agent.CumulativeTokenSnapshot if tokenUsage == nil { - tokenUsage = agent.CalculateTokenUsage(ctx, ag, transcriptData, transcriptLinesAtStart, subagentsDir) + 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) + } } // Build fully-populated step context and delegate to strategy @@ -948,6 +957,20 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev return fmt.Errorf("failed to save step: %w", err) } + // Persist the cumulative token baseline after SaveStep (SaveStep may + // reinitialize session state, which would drop an earlier update). The next + // checkpoint reads this to scan only newly added transcript lines. + if nextTokenBaseline != nil { + mutErr := strategy.MutateSessionState(ctx, sessionID, func(state *strategy.SessionState) error { + state.CumulativeTokenBaseline = nextTokenBaseline + return nil + }) + if mutErr != nil && !errors.Is(mutErr, strategy.ErrStateNotFound) { + logging.Warn(logCtx, "failed to persist cumulative token baseline", + slog.String("error", mutErr.Error())) + } + } + // Update session state with backfilled prompt after SaveStep. // Done after SaveStep because SaveStep may reinitialize session state, // which would overwrite an earlier LastPrompt update. @@ -975,6 +998,26 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev return nil } +// computeIncrementalTokenUsage computes the checkpoint token delta for a +// cumulative-count agent (Codex) by scanning only the transcript lines added +// since the last checkpoint. It reads the persisted cumulative baseline from +// session state and returns both the delta and the fresh baseline to persist +// after SaveStep. On any failure it returns (nil, nil); a nil prior baseline +// makes the calculator fall back to a one-time full scan. +func computeIncrementalTokenUsage(ctx, logCtx context.Context, inc agent.IncrementalTokenCalculator, sessionID string, transcriptData []byte, fromOffset int) (*agent.TokenUsage, *agent.CumulativeTokenSnapshot) { + var prior *agent.CumulativeTokenSnapshot + if st, err := strategy.LoadSessionState(ctx, sessionID); err == nil && st != nil { + prior = st.CumulativeTokenBaseline + } + usage, next, err := inc.CalculateTokenUsageIncremental(transcriptData, fromOffset, prior) + if err != nil { + logging.Debug(logCtx, "incremental token calculation failed", + slog.String("error", err.Error())) + return nil, nil + } + return usage, next +} + // handleLifecycleCompaction handles context compaction: saves current progress // but stays in ACTIVE phase (unlike TurnEnd which transitions to IDLE). // Also resets the transcript offset since the transcript may be truncated. diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index 1eed68ebfb..281c3ed523 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -309,6 +309,14 @@ type State struct { // cumulative total on every checkpoint. SubagentTokensBaseline *agent.TokenUsage `json:"subagent_tokens_baseline,omitempty"` + // CumulativeTokenBaseline caches the last cumulative token snapshot for agents + // that report CUMULATIVE token counts (Codex), so turn-end token calc scans + // only the transcript lines added since the last checkpoint instead of the + // whole rollout (O(session) per hook → O(session^2) across a session). nil + // until the first checkpoint that observes token data, or after a transcript + // reset — both fall back to a one-time full scan. + CumulativeTokenBaseline *agent.CumulativeTokenSnapshot `json:"cumulative_token_baseline,omitempty"` + // SkillEvents records explicit native skill signals observed during this session. // Stored as sidecar metadata so consumers can collapse skill-related transcript events // without mutating the raw agent transcript. From 1ca43abb97aa768d0e3e9cddf7f4bb719fb99c57 Mon Sep 17 00:00:00 2001 From: Victor Gutierrez Calderon Date: Fri, 24 Jul 2026 09:54:09 +0200 Subject: [PATCH 2/2] perf(codex): fall back on incremental token error; slice before file extract Address review feedback: - computeIncrementalTokenUsage now falls back to the full-scan token calc on incremental error instead of silently dropping the checkpoint's token usage. - ExtractAllModifiedFiles slices to the post-offset lines before splitting, so per-turn work scales with the delta rather than the whole session. Co-Authored-By: Claude Opus 4.8 (1M context) Entire-Checkpoint: 01KY9HW4BJJRTQ7T5TMK7KGPR6 --- cmd/entire/cli/agent/codex/transcript.go | 13 +++++++------ cmd/entire/cli/lifecycle.go | 13 +++++++++++-- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/cmd/entire/cli/agent/codex/transcript.go b/cmd/entire/cli/agent/codex/transcript.go index b94a91353c..1c9a6dcb1c 100644 --- a/cmd/entire/cli/agent/codex/transcript.go +++ b/cmd/entire/cli/agent/codex/transcript.go @@ -386,14 +386,15 @@ func scanTokenDelta(data []byte, startLineNum, fromOffset int, seed *tokenUsageD // transcript bytes it already holds, instead of falling back to a second full-file // disk read via ExtractModifiedFilesFromOffset. func (c *CodexAgent) ExtractAllModifiedFiles(transcriptData []byte, fromOffset int, _ string) ([]string, error) { + // Slice to the lines added since fromOffset before splitting, so per-turn work + // scales with the delta rather than the whole session. + data := transcriptData + if fromOffset > 0 { + data = transcript.SliceFromLine(transcriptData, fromOffset) + } seen := make(map[string]struct{}) var files []string - lineNum := 0 - for _, lineData := range splitJSONL(transcriptData) { - lineNum++ - if lineNum <= fromOffset { - continue - } + for _, lineData := range splitJSONL(data) { for _, f := range extractFilesFromLine(lineData) { if _, ok := seen[f]; !ok { seen[f] = struct{}{} diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index 0ed598c2d3..0328f79491 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1011,9 +1011,18 @@ func computeIncrementalTokenUsage(ctx, logCtx context.Context, inc agent.Increme } usage, next, err := inc.CalculateTokenUsageIncremental(transcriptData, fromOffset, prior) if err != nil { - logging.Debug(logCtx, "incremental token calculation failed", + // Fall back to the full-scan path so a checkpoint's token usage is not + // silently dropped; skip persisting a baseline so the next hook re-derives + // one from a clean full scan. + logging.Debug(logCtx, "incremental token calculation failed, falling back to full scan", slog.String("error", err.Error())) - return nil, nil + fallback, fbErr := inc.CalculateTokenUsage(transcriptData, fromOffset) + if fbErr != nil { + logging.Debug(logCtx, "full-scan token calculation also failed", + slog.String("error", fbErr.Error())) + return nil, nil + } + return fallback, nil } return usage, next }