Skip to content

Commit 406904c

Browse files
gtrrz-victorclaude
andcommitted
perf(codex): incremental turn-end token calc to fix O(N²) reparse
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
1 parent 3b483c3 commit 406904c

8 files changed

Lines changed: 665 additions & 12 deletions

File tree

CLAUDE.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,7 @@ The manual-commit strategy (`manual_commit*.go`) does not modify the active bran
595595
- Rewind restores files from shadow branch commit tree (does not use `git reset`)
596596
- **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.
597597
- **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`.
598+
- **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.
598599
- Tracks session state in `.git/entire-sessions/` (shared across worktrees)
599600
- **Shadow branch migration** - if user does stash/pull/rebase (HEAD changes without commit), shadow branch is automatically moved to new base commit
600601
- **Orphaned branch cleanup** - if a shadow branch exists without a corresponding session state file, it is automatically reset when a new session starts

cmd/entire/cli/agent/capabilities.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,14 @@ func AsSubagentAwareExtractor(ag Agent) (SubagentAwareExtractor, bool) {
127127
return declaredCapability[SubagentAwareExtractor](ag, func(c DeclaredCaps) bool { return c.SubagentAwareExtractor })
128128
}
129129

130+
// AsIncrementalTokenCalculator returns the agent as IncrementalTokenCalculator
131+
// if it implements the interface. Incremental token calculation is a built-in
132+
// optimization (for agents with cumulative-count transcripts like Codex) with
133+
// no external-protocol equivalent, so it resolves by type assertion alone.
134+
func AsIncrementalTokenCalculator(ag Agent) (IncrementalTokenCalculator, bool) {
135+
return builtinCapability[IncrementalTokenCalculator](ag)
136+
}
137+
130138
// AsSessionBaseDirProvider returns the agent as SessionBaseDirProvider if it implements
131139
// the interface. No capability declaration is needed since this is a built-in-only feature
132140
// (external agents use the agent binary's own session resolution).
Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
package codex
2+
3+
import (
4+
"strings"
5+
"testing"
6+
7+
"github.com/entireio/cli/cmd/entire/cli/agent"
8+
"github.com/entireio/cli/cmd/entire/cli/transcript"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
// lineCount counts non-empty JSONL lines the same way GetTranscriptPosition does
13+
// for well-formed rollouts (no blank lines).
14+
func lineCount(data []byte) int { return len(splitJSONL(data)) }
15+
16+
// TestIncremental_MatchesFullScan replays every turn-end hook of a growing
17+
// session and asserts the incremental calc produces byte-identical numbers to
18+
// the full-scan CalculateTokenUsage — carrying the persisted snapshot forward
19+
// like the lifecycle does.
20+
func TestIncremental_MatchesFullScan(t *testing.T) {
21+
t.Parallel()
22+
ag := &CodexAgent{}
23+
24+
const turns = 8
25+
var prior *agent.CumulativeTokenSnapshot
26+
fromOffset := 0 // turn 1 starts at line 0
27+
28+
for turn := 1; turn <= turns; turn++ {
29+
data, _ := buildRollout(turn, false, 0)
30+
31+
full, err := ag.CalculateTokenUsage(data, fromOffset)
32+
require.NoError(t, err)
33+
34+
inc, next, err := ag.CalculateTokenUsageIncremental(data, fromOffset, prior)
35+
require.NoError(t, err)
36+
require.Equal(t, full, inc, "turn %d: incremental delta must equal full scan", turn)
37+
38+
// Next turn starts where this transcript ended.
39+
prior = next
40+
fromOffset = lineCount(data)
41+
}
42+
}
43+
44+
// TestIncremental_MatchesFullScan_WithCachedTokens uses the shipped fixture that
45+
// includes cached-input tokens and a mid-session offset.
46+
func TestIncremental_MatchesFullScan_WithCachedTokens(t *testing.T) {
47+
t.Parallel()
48+
ag := &CodexAgent{}
49+
data := []byte(sampleRollout)
50+
51+
for _, offset := range []int{0, 4, 8, 11} {
52+
full, err := ag.CalculateTokenUsage(data, offset)
53+
require.NoError(t, err)
54+
55+
// Cold start (no prior): must fall back to a full scan and match exactly.
56+
inc, next, err := ag.CalculateTokenUsageIncremental(data, offset, nil)
57+
require.NoError(t, err)
58+
require.Equal(t, full, inc, "offset %d: cold-start incremental must equal full scan", offset)
59+
require.NotNil(t, next)
60+
require.Equal(t, lineCount(data), next.LineCount, "offset %d: snapshot line count", offset)
61+
}
62+
}
63+
64+
// TestIncremental_ColdStartProducesReusableBaseline verifies a cold-start call
65+
// yields a snapshot that a subsequent incremental call can baseline off to get
66+
// the correct per-checkpoint delta.
67+
func TestIncremental_ColdStartProducesReusableBaseline(t *testing.T) {
68+
t.Parallel()
69+
ag := &CodexAgent{}
70+
71+
// Turn 1: 2 token_count lines, cold start.
72+
turn1, _ := buildRollout(2, false, 0)
73+
_, snap, err := ag.CalculateTokenUsageIncremental(turn1, 0, nil)
74+
require.NoError(t, err)
75+
require.NotNil(t, snap)
76+
require.Equal(t, lineCount(turn1), snap.LineCount)
77+
78+
// Turn 2: one more token_count line appended.
79+
turn2, _ := buildRollout(3, false, 0)
80+
fromOffset := lineCount(turn1)
81+
82+
inc, _, err := ag.CalculateTokenUsageIncremental(turn2, fromOffset, snap)
83+
require.NoError(t, err)
84+
full, err := ag.CalculateTokenUsage(turn2, fromOffset)
85+
require.NoError(t, err)
86+
require.Equal(t, full, inc, "warm incremental delta must equal full scan")
87+
}
88+
89+
// TestIncremental_ResumedOffsetFallsBackToFullScan covers a stale baseline
90+
// (LineCount past the current fromOffset, as after a transcript reset/resume):
91+
// the calc must ignore it and full-scan, still producing correct numbers.
92+
func TestIncremental_ResumedOffsetFallsBackToFullScan(t *testing.T) {
93+
t.Parallel()
94+
ag := &CodexAgent{}
95+
data, _ := buildRollout(3, false, 0)
96+
fromOffset := 3
97+
98+
stale := &agent.CumulativeTokenSnapshot{LineCount: 9999, InputTokens: 1, CachedInputTokens: 1, OutputTokens: 1}
99+
inc, _, err := ag.CalculateTokenUsageIncremental(data, fromOffset, stale)
100+
require.NoError(t, err)
101+
full, err := ag.CalculateTokenUsage(data, fromOffset)
102+
require.NoError(t, err)
103+
require.Equal(t, full, inc, "stale baseline must be ignored (full-scan fallback)")
104+
}
105+
106+
// TestIncremental_DoesNotParseHistory is the acceptance criterion: parse work is
107+
// bounded by lines added since the last checkpoint, not the whole session.
108+
// We corrupt every pre-offset line with bogus token_counts. A full scan of the
109+
// corrupted transcript changes the numbers (it reads the bogus baseline); the
110+
// incremental call, given a valid prior baseline, slices the history off and is
111+
// UNAFFECTED — proving it never parsed those lines.
112+
func TestIncremental_DoesNotParseHistory(t *testing.T) {
113+
t.Parallel()
114+
ag := &CodexAgent{}
115+
116+
clean, _ := buildRollout(5, false, 0)
117+
fromOffset := lineCount(clean)
118+
119+
// Correct baseline as of fromOffset = the last cumulative in `clean`.
120+
_, prior, err := ag.CalculateTokenUsageIncremental(clean, fromOffset, nil)
121+
require.NoError(t, err)
122+
require.NotNil(t, prior)
123+
124+
// Append one more turn to `clean` (the new delta the next hook sees).
125+
cleanNext, _ := buildRollout(6, false, 0)
126+
tail := transcript.SliceFromLine(cleanNext, fromOffset)
127+
128+
// Build a corrupted transcript: same line count for lines <= fromOffset, but
129+
// every pre-offset token_count carries absurd values, followed by the real tail.
130+
var corrupt strings.Builder
131+
corrupt.WriteString(`{"timestamp":"t","type":"session_meta","payload":{"id":"x"}}`)
132+
corrupt.WriteByte('\n')
133+
for i := 2; i <= fromOffset; i++ {
134+
corrupt.WriteString(tokenCountLine(9_000_000, 8_000_000, 900_000))
135+
corrupt.WriteByte('\n')
136+
}
137+
corrupt.Write(tail)
138+
corrupted := []byte(corrupt.String())
139+
require.Equal(t, fromOffset, lineCount(corrupted)-lineCount(tail),
140+
"corrupted history must have the same pre-offset line count")
141+
142+
// A FULL scan of the corrupted transcript reads the bogus baseline → wrong.
143+
fullCorrupt, err := ag.CalculateTokenUsage(corrupted, fromOffset)
144+
require.NoError(t, err)
145+
146+
// The INCREMENTAL call with the valid prior slices off history → correct,
147+
// identical to the clean full scan.
148+
incCorrupt, _, err := ag.CalculateTokenUsageIncremental(corrupted, fromOffset, prior)
149+
require.NoError(t, err)
150+
cleanFull, err := ag.CalculateTokenUsage(cleanNext, fromOffset)
151+
require.NoError(t, err)
152+
153+
require.Equal(t, cleanFull, incCorrupt,
154+
"incremental result must ignore corrupted history (only delta lines parsed)")
155+
require.NotEqual(t, fullCorrupt, incCorrupt,
156+
"a full scan WOULD have been corrupted — proving incremental skipped history")
157+
}
158+
159+
// TestExtractAllModifiedFiles_FromBytes verifies the bytes-based extractor
160+
// returns the same files as the disk-based ExtractModifiedFilesFromOffset, so
161+
// the turn-end pipeline avoids the second full-file disk read.
162+
func TestExtractAllModifiedFiles_FromBytes(t *testing.T) {
163+
t.Parallel()
164+
ag := &CodexAgent{}
165+
data := []byte(sampleRollout)
166+
path := writeSampleRollout(t)
167+
168+
for _, offset := range []int{0, 9} {
169+
fromBytes, err := ag.ExtractAllModifiedFiles(data, offset, "")
170+
require.NoError(t, err)
171+
fromDisk, _, err := ag.ExtractModifiedFilesFromOffset(path, offset)
172+
require.NoError(t, err)
173+
require.ElementsMatch(t, fromDisk, fromBytes, "offset %d: bytes extractor must match disk extractor", offset)
174+
}
175+
}
176+
177+
// TestCodex_DeclaresBytesCapabilities locks in that Codex now routes both token
178+
// calc and file extraction through in-memory bytes (no disk reread, incremental
179+
// tokens) — the opposite of the pre-fix validation guard H3.
180+
func TestCodex_DeclaresBytesCapabilities(t *testing.T) {
181+
t.Parallel()
182+
ag := NewCodexAgent()
183+
184+
_, incOK := agent.AsIncrementalTokenCalculator(ag)
185+
require.True(t, incOK, "Codex must be an IncrementalTokenCalculator")
186+
187+
_, subOK := agent.AsSubagentAwareExtractor(ag)
188+
require.True(t, subOK, "Codex must expose a bytes-based (SubagentAware) file extractor")
189+
}

0 commit comments

Comments
 (0)