diff --git a/cmd/entire/cli/integration_test/cursor_token_condensation_test.go b/cmd/entire/cli/integration_test/cursor_token_condensation_test.go index 89c038115e..a772d9ba9b 100644 --- a/cmd/entire/cli/integration_test/cursor_token_condensation_test.go +++ b/cmd/entire/cli/integration_test/cursor_token_condensation_test.go @@ -6,6 +6,7 @@ import ( "bytes" "context" "encoding/json" + "fmt" "os" "path/filepath" "testing" @@ -21,9 +22,18 @@ import ( func runCursorHook(t *testing.T, env *TestEnv, cursorProjectDir, hookName string, input map[string]any) { t.Helper() + require.NoError(t, tryCursorHook(t, env, cursorProjectDir, hookName, input)) +} + +// tryCursorHook runs a Cursor hook and returns its failure instead of failing +// the test, so it can run off the test goroutine. +func tryCursorHook(t *testing.T, env *TestEnv, cursorProjectDir, hookName string, input map[string]any) error { + t.Helper() inputJSON, err := json.Marshal(input) - require.NoError(t, err) + if err != nil { + return fmt.Errorf("marshal %s input: %w", hookName, err) + } cmd := execx.NonInteractive(context.Background(), getTestBinary(), "hooks", "cursor", hookName) cmd.Dir = env.RepoDir @@ -31,8 +41,11 @@ func runCursorHook(t *testing.T, env *TestEnv, cursorProjectDir, hookName string cmd.Env = append(env.cliEnv(), "ENTIRE_TEST_CURSOR_PROJECT_DIR="+cursorProjectDir) out, err := cmd.CombinedOutput() - require.NoErrorf(t, err, "cursor %s hook failed\ninput: %s\noutput: %s", hookName, inputJSON, out) t.Logf("cursor %s output: %s", hookName, out) + if err != nil { + return fmt.Errorf("cursor %s hook failed\ninput: %s\noutput: %s: %w", hookName, inputJSON, out, err) + } + return nil } func TestCursorTokenUsage_SurvivesCondensation(t *testing.T) { @@ -203,6 +216,164 @@ func TestCursorTokenUsage_PerCheckpointScoping(t *testing.T) { "checkpoint 2 OutputTokens must be turn 2 only (30), not the cumulative session total (80)") } +// TestCursorTokenUsage_TurnWithoutFileChanges covers the checkpoint-less +// turn-end path: a Stop that changes no files skips SaveStep, so the stop +// payload's tokens have to be recorded there or they are lost from both the +// session total and the next checkpoint (Cursor's transcript carries no usage +// for condensation to recompute from). +func TestCursorTokenUsage_TurnWithoutFileChanges(t *testing.T) { + t.Parallel() + + env := NewFeatureBranchEnv(t) + env.InitEntireWithAgent(agent.AgentNameCursor) + + cursorProjectDir := t.TempDir() + if resolved, err := filepath.EvalSymlinks(cursorProjectDir); err == nil { + cursorProjectDir = resolved + } + + const conversationID = "cursor-nochange-session" + + transcriptDir := filepath.Join(cursorProjectDir, conversationID) + require.NoError(t, os.MkdirAll(transcriptDir, 0o755)) + transcriptPath := filepath.Join(transcriptDir, conversationID+".jsonl") + + hook := func(name string, extra map[string]any) { + t.Helper() + input := map[string]any{"conversation_id": conversationID, "transcript_path": transcriptPath} + for k, v := range extra { + input[k] = v + } + runCursorHook(t, env, cursorProjectDir, name, input) + } + // turn appends a prompt/answer pair, runs before-submit-prompt, writes + // files (none for a read-only turn), and stops with the given counts. + // Cursor's input_tokens includes cache reads and writes. + turn := func(prompt string, files []string, input, output, cacheRead, cacheWrite int) { + t.Helper() + f, err := os.OpenFile(transcriptPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600) + require.NoError(t, err) + _, werr := f.WriteString(`{"type":"user","text":"` + prompt + `"}` + "\n") + require.NoError(t, f.Close()) + require.NoError(t, werr) + hook("before-submit-prompt", map[string]any{"prompt": prompt}) + for _, name := range files { + env.WriteFile(name, "package main\n// "+prompt+"\n") + } + hook("stop", map[string]any{ + "model": "cursor-default", + "loop_count": 1, + "input_tokens": input, + "output_tokens": output, + "cache_read_tokens": cacheRead, + "cache_write_tokens": cacheWrite, + }) + } + sessionTokens := func() *agent.TokenUsage { + t.Helper() + state, err := env.GetSessionState(conversationID) + require.NoError(t, err) + require.NotNil(t, state, "session state must exist") + require.NotNil(t, state.TokenUsage, "session state must carry token usage") + return state.TokenUsage + } + + hook("session-start", map[string]any{"model": "cursor-default"}) + + // Turn 1 writes a file: fresh input 200, output 50, cache 4000/800. + turn("turn one", []string{"turn1.go"}, 5000, 50, 4000, 800) + // Turn 2 changes nothing: fresh input 100, output 7, cache 300/20. + turn("turn two", nil, 420, 7, 300, 20) + + total := sessionTokens() + require.Equal(t, 300, total.InputTokens, "session total must include the no-change turn (200+100)") + require.Equal(t, 57, total.OutputTokens, "session total must include the no-change turn (50+7)") + require.Equal(t, 4300, total.CacheReadTokens) + require.Equal(t, 820, total.CacheCreationTokens) + + env.GitCommitWithHooks("Turn 1", "turn1.go") + checkpoint1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, checkpoint1, "expected a checkpoint after the commit") + + cp1 := readCommittedTokenUsage(t, env, checkpoint1) + require.NotNil(t, cp1, "checkpoint must carry token usage") + require.Equal(t, 300, cp1.InputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 57, cp1.OutputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 4300, cp1.CacheReadTokens) + require.Equal(t, 820, cp1.CacheCreationTokens) + require.Equal(t, 300, sessionTokens().InputTokens, "condensation must not change the session total") + + // Turn 3 writes a file: the next checkpoint must not repeat turn 2. + turn("turn three", []string{"turn3.go"}, 1000, 9, 0, 0) + env.GitCommitWithHooks("Turn 3", "turn3.go") + checkpoint2 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, checkpoint2) + require.NotEqual(t, checkpoint1, checkpoint2, "turn 3 must produce a distinct checkpoint") + + cp2 := readCommittedTokenUsage(t, env, checkpoint2) + require.NotNil(t, cp2, "checkpoint 2 must carry turn 3 token usage") + require.Equal(t, 1000, cp2.InputTokens, "checkpoint 2 must hold turn 3 only") + require.Equal(t, 9, cp2.OutputTokens, "checkpoint 2 must hold turn 3 only") + require.Equal(t, 1300, sessionTokens().InputTokens, "session total = every turn once") +} + +// TestCursorTokenUsage_RepeatedStopCountsNoChangeTurnOnce pins that a +// no-change turn's stop-payload tokens are recorded once per turn: Cursor's +// payload is the turn's own usage, so a repeated or concurrent stop for the +// same turn would otherwise add it again. +func TestCursorTokenUsage_RepeatedStopCountsNoChangeTurnOnce(t *testing.T) { + t.Parallel() + + env := NewFeatureBranchEnv(t) + env.InitEntireWithAgent(agent.AgentNameCursor) + + cursorProjectDir := t.TempDir() + if resolved, err := filepath.EvalSymlinks(cursorProjectDir); err == nil { + cursorProjectDir = resolved + } + + const conversationID = "cursor-repeated-stop-session" + transcriptDir := filepath.Join(cursorProjectDir, conversationID) + require.NoError(t, os.MkdirAll(transcriptDir, 0o755)) + transcriptPath := filepath.Join(transcriptDir, conversationID+".jsonl") + + input := func(extra map[string]any) map[string]any { + in := map[string]any{"conversation_id": conversationID, "transcript_path": transcriptPath} + for k, v := range extra { + in[k] = v + } + return in + } + startTurn := func(prompt string) { + t.Helper() + appendLines(t, transcriptPath, `{"type":"user","text":"`+prompt+`"}`) + runCursorHook(t, env, cursorProjectDir, "before-submit-prompt", input(map[string]any{"prompt": prompt})) + } + // Cursor's input_tokens includes cache reads and writes; none here. + stop := func(in, out int) func() error { + return func() error { + return tryCursorHook(t, env, cursorProjectDir, "stop", input(map[string]any{ + "model": "cursor-default", "loop_count": 1, "input_tokens": in, "output_tokens": out, + })) + } + } + + runCursorHook(t, env, cursorProjectDir, "session-start", input(map[string]any{"model": "cursor-default"})) + + startTurn("turn one") + require.NoError(t, stop(100, 10)()) + requireTurnTokens(t, env, conversationID, 100, 10, "after a no-change turn") + + require.NoError(t, stop(100, 10)(), "a repeated stop must still succeed") + requireTurnTokens(t, env, conversationID, 100, 10, "a repeated stop must add nothing") + + startTurn("turn two") + for i, err := range runConcurrently(stop(40, 4), stop(40, 4)) { + require.NoError(t, err, "concurrent stop %d", i) + } + requireTurnTokens(t, env, conversationID, 140, 14, "concurrent stops must record the turn once") +} + func readCommittedTokenUsage(t *testing.T, env *TestEnv, checkpointID string) *agent.TokenUsage { t.Helper() content, found := env.ReadFileFromBranch(paths.MetadataBranchName, SessionMetadataPath(checkpointID)) diff --git a/cmd/entire/cli/integration_test/factoryai_worker_session_test.go b/cmd/entire/cli/integration_test/factoryai_worker_session_test.go index a3ec8cbac0..94182f5dd2 100644 --- a/cmd/entire/cli/integration_test/factoryai_worker_session_test.go +++ b/cmd/entire/cli/integration_test/factoryai_worker_session_test.go @@ -3,7 +3,12 @@ package integration import ( + "encoding/json" "testing" + + "github.com/entireio/cli/cmd/entire/cli/checkpoint" + "github.com/entireio/cli/cmd/entire/cli/paths" + "github.com/stretchr/testify/require" ) // Factory AI Droid dispatches Workers as detached sessions: the parent's @@ -128,6 +133,113 @@ func TestFactoryDroidWorkerSessionBecomesTaskCheckpoint(t *testing.T) { }) } +// TestFactoryDroidReadOnlyWorkerReachesParentCheckpoint covers a Worker turn +// that changes no files: it must still become a completed task record on the +// parent, so its transcript and tokens reach the parent's next checkpoint, +// with the tokens on the task once and never on the Worker's own session. +func TestFactoryDroidReadOnlyWorkerReachesParentCheckpoint(t *testing.T) { + t.Parallel() + + env := NewTestEnv(t) + env.InitRepo() + env.InitEntire() + env.WriteFile(".gitignore", ".entire/\n") + env.WriteFile("README.md", "# Test Repository") + env.GitAdd(".gitignore") + env.GitAdd("README.md") + env.GitCommit("Initial commit") + + parent := env.NewFactoryDroidSession() + require.NoError(t, env.SimulateFactoryDroidUserPromptSubmit(parent.ID)) + + // A read-only Worker: it reads and reports, consuming tokens, and writes + // nothing. + const toolUseID = "toolu_reader_01" + worker := env.NewFactoryDroidSession() + require.NoError(t, env.SimulateFactoryDroidUserPromptSubmit(worker.ID)) + appendLines(t, worker.TranscriptPath, + `{"type":"message","id":"w1","message":{"role":"user","content":[{"type":"text","text":"Summarize the repo"}]}}`, + `{"type":"message","id":"w2","message":{"role":"assistant","id":"msg_w1","content":[{"type":"text","text":"It is a test repository."}],"usage":{"input_tokens":70,"output_tokens":7}}}`) + worker.MarkAsWorkerSession(parent.ID, toolUseID, "worker: Summarize the repo") + + // Twice: a repeated Stop must not count the Worker's usage again. + require.NoError(t, env.SimulateFactoryDroidStop(worker.ID, worker.TranscriptPath)) + require.NoError(t, env.SimulateFactoryDroidStop(worker.ID, worker.TranscriptPath)) + + state, err := env.GetSessionState(parent.ID) + require.NoError(t, err) + require.NotNil(t, state) + rec := state.FindTaskRecord(toolUseID) + require.NotNil(t, rec, "a read-only Worker turn must still record a task on the parent") + require.False(t, rec.CompletedAt.IsZero(), "the record must be completed") + require.Empty(t, rec.Files, "a read-only Worker touches no files") + require.Equal(t, worker.TranscriptPath, rec.DeclaredTranscriptPath) + require.NotNil(t, rec.TokenUsage, "the record must carry the Worker's usage") + require.Equal(t, 70, rec.TokenUsage.InputTokens, "the Worker's usage, once") + require.Equal(t, 7, rec.TokenUsage.OutputTokens, "the Worker's usage, once") + + if workerState, err := env.GetSessionState(worker.ID); err != nil { + t.Fatalf("GetSessionState(worker) failed: %v", err) + } else if workerState != nil { + require.Zero(t, workerState.StepCount, "a Worker's turn must not step its own session") + if workerState.TokenUsage != nil { + require.Zero(t, workerState.TokenUsage.InputTokens, "the Worker's usage belongs to the parent's task, not its own session") + } + } + + // The parent's own turn changes a file; committing it condenses the + // parent, Worker task included. + env.WriteFile("feature.go", "package main\n") + parent.CreateDroidTranscript("Add a feature", []FileChange{{Path: "feature.go", Content: "package main\n"}}) + require.NoError(t, env.SimulateFactoryDroidStop(parent.ID, parent.TranscriptPath)) + env.GitCommitWithHooks("Add feature", "feature.go") + checkpointID := env.TryGetLatestCheckpointID() + require.NotEmpty(t, checkpointID, "expected a checkpoint for the parent's commit") + + raw, ok := env.ReadFileFromBranch(paths.MetadataBranchName, CheckpointTaskFilePath(checkpointID, toolUseID, "task.json")) + require.True(t, ok, "the read-only Worker's task must be materialized in the parent's checkpoint") + var task checkpoint.TaskRecord + require.NoError(t, json.Unmarshal([]byte(raw), &task)) + require.Equal(t, worker.ID, task.AgentID) + require.NotNil(t, task.TokenUsage, "the stored task must carry the Worker's usage: %s", raw) + require.Equal(t, 70, task.TokenUsage.InputTokens, "the stored task holds the Worker's usage once: %s", raw) + require.Equal(t, 7, task.TokenUsage.OutputTokens, "the stored task holds the Worker's usage once: %s", raw) + + transcript, ok := env.ReadFileFromBranch(paths.MetadataBranchName, + CheckpointTaskFilePath(checkpointID, toolUseID, paths.AgentTranscriptFileName(worker.ID))) + require.True(t, ok, "the Worker's transcript must be stored with its task") + require.Contains(t, transcript, "It is a test repository.") + + if usage := readCommittedTokenUsage(t, env, checkpointID); usage != nil { + require.NotEqual(t, 70, usage.InputTokens, "the parent session's own usage must not absorb the Worker's") + } + + // A second Worker turn after that commit: the next checkpoint's task must + // hold only the second turn, not the first again. + require.NoError(t, env.SimulateFactoryDroidUserPromptSubmit(parent.ID)) + require.NoError(t, env.SimulateFactoryDroidUserPromptSubmitWithTranscript(worker.ID, worker.TranscriptPath)) + appendLines(t, worker.TranscriptPath, + `{"type":"message","id":"w3","message":{"role":"user","content":[{"type":"text","text":"And the tests?"}]}}`, + `{"type":"message","id":"w4","message":{"role":"assistant","id":"msg_w2","content":[{"type":"text","text":"There are none."}],"usage":{"input_tokens":30,"output_tokens":3}}}`) + require.NoError(t, env.SimulateFactoryDroidStop(worker.ID, worker.TranscriptPath)) + + env.WriteFile("feature2.go", "package main\n") + parent.CreateDroidTranscript("Add another feature", []FileChange{{Path: "feature2.go", Content: "package main\n"}}) + require.NoError(t, env.SimulateFactoryDroidStop(parent.ID, parent.TranscriptPath)) + env.GitCommitWithHooks("Add feature 2", "feature2.go") + checkpoint2 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, checkpoint2) + require.NotEqual(t, checkpointID, checkpoint2, "expected a second checkpoint") + + raw, ok = env.ReadFileFromBranch(paths.MetadataBranchName, CheckpointTaskFilePath(checkpoint2, toolUseID, "task.json")) + require.True(t, ok, "the second Worker turn must be materialized in the second checkpoint") + var task2 checkpoint.TaskRecord + require.NoError(t, json.Unmarshal([]byte(raw), &task2)) + require.NotNil(t, task2.TokenUsage, "the second stored task must carry the Worker's usage: %s", raw) + require.Equal(t, 30, task2.TokenUsage.InputTokens, "the first Worker turn is already in the first checkpoint: %s", raw) + require.Equal(t, 3, task2.TokenUsage.OutputTokens, "the first Worker turn is already in the first checkpoint: %s", raw) +} + // TestFactoryDroidTopLevelSessionStillCheckpoints guards the fallback: an // ordinary Droid session has no calling IDs and must keep its own checkpoint. func TestFactoryDroidTopLevelSessionStillCheckpoints(t *testing.T) { diff --git a/cmd/entire/cli/integration_test/hooks.go b/cmd/entire/cli/integration_test/hooks.go index 595af0fc4b..8a4e9323dd 100644 --- a/cmd/entire/cli/integration_test/hooks.go +++ b/cmd/entire/cli/integration_test/hooks.go @@ -826,10 +826,17 @@ func (r *FactoryDroidHookRunner) runDroidHookInRepoDir(hookName string, inputJSO // SimulateUserPromptSubmit simulates the UserPromptSubmit hook for Factory Droid. func (r *FactoryDroidHookRunner) SimulateUserPromptSubmit(sessionID string) error { r.T.Helper() + return r.SimulateUserPromptSubmitWithTranscript(sessionID, "") +} + +// SimulateUserPromptSubmitWithTranscript simulates UserPromptSubmit with the +// transcript path Droid reports, so turn start records the transcript offset. +func (r *FactoryDroidHookRunner) SimulateUserPromptSubmitWithTranscript(sessionID, transcriptPath string) error { + r.T.Helper() input := map[string]string{ "session_id": sessionID, - "transcript_path": "", + "transcript_path": transcriptPath, "prompt": "test prompt", } @@ -999,6 +1006,13 @@ func (env *TestEnv) SimulateFactoryDroidUserPromptSubmit(sessionID string) error return runner.SimulateUserPromptSubmit(sessionID) } +// SimulateFactoryDroidUserPromptSubmitWithTranscript is a convenience method on TestEnv. +func (env *TestEnv) SimulateFactoryDroidUserPromptSubmitWithTranscript(sessionID, transcriptPath string) error { + env.T.Helper() + runner := NewFactoryDroidHookRunner(env.RepoDir, env.T) + return runner.SimulateUserPromptSubmitWithTranscript(sessionID, transcriptPath) +} + // SimulateFactoryDroidStop is a convenience method on TestEnv. func (env *TestEnv) SimulateFactoryDroidStop(sessionID, transcriptPath string) error { env.T.Helper() diff --git a/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go b/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go new file mode 100644 index 0000000000..c3b12f454b --- /dev/null +++ b/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go @@ -0,0 +1,230 @@ +//go:build integration + +package integration + +import ( + "fmt" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/stretchr/testify/require" +) + +// Transcript-token agents keep their usage in the transcript, so a turn that +// changes no files is still covered by the next checkpoint's transcript slice. +// What it used to miss is the session total: the checkpoint-less turn-end path +// skipped SaveStep and recorded nothing. These tests drive a file-changing +// turn, a no-change turn, a commit, then one more turn and commit, and assert +// each turn lands once in the session total and once in a checkpoint. + +// appendLines appends JSONL lines to path, creating it if needed. +func appendLines(t *testing.T, path string, lines ...string) { + t.Helper() + require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o755)) + f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600) + require.NoError(t, err) + for _, line := range lines { + if _, err := f.WriteString(line + "\n"); err != nil { + _ = f.Close() + require.NoError(t, err) + } + } + require.NoError(t, f.Close()) +} + +// sessionTokenTotal returns the session-wide token total from session state. +func sessionTokenTotal(t *testing.T, env *TestEnv, sessionID string) *agent.TokenUsage { + t.Helper() + state, err := env.GetSessionState(sessionID) + require.NoError(t, err) + require.NotNil(t, state, "session state must exist") + require.NotNil(t, state.TokenUsage, "session state must carry token usage") + return state.TokenUsage +} + +// commitCheckpointTokens commits files with hooks and returns the new +// checkpoint's committed token usage, requiring it differ from previous. +func commitCheckpointTokens(t *testing.T, env *TestEnv, previous, message string, files ...string) (string, *agent.TokenUsage) { + t.Helper() + env.GitCommitWithHooks(message, files...) + checkpointID := env.TryGetLatestCheckpointID() + require.NotEmpty(t, checkpointID, "%s: expected a checkpoint", message) + require.NotEqual(t, previous, checkpointID, "%s: expected a new checkpoint", message) + usage := readCommittedTokenUsage(t, env, checkpointID) + require.NotNil(t, usage, "%s: checkpoint must carry token usage", message) + return checkpointID, usage +} + +func TestClaudeCodeTokenUsage_TurnWithoutFileChanges(t *testing.T) { + t.Parallel() + + env := NewFeatureBranchEnv(t) + session := env.NewSession() + transcript := session.TranscriptPath + + // turn runs one Claude Code turn: prompt hook, then the turn's transcript + // lines (one assistant message with the given usage), its files, and Stop. + n := 0 + turn := func(files []string, input, output int) { + t.Helper() + n++ + prompt := fmt.Sprintf("turn %d", n) + require.NoError(t, env.SimulateUserPromptSubmitWithPromptAndTranscriptPath(session.ID, prompt, transcript)) + appendLines(t, transcript, + fmt.Sprintf(`{"uuid":"user-%d","type":"user","message":{"content":%q}}`, n, prompt), + fmt.Sprintf(`{"uuid":"asst-%d","type":"assistant","message":{"id":"msg_%d","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":%d,"output_tokens":%d}}}`, n, n, input, output)) + for _, name := range files { + env.WriteFile(name, prompt+"\n") + } + require.NoError(t, env.SimulateStop(session.ID, transcript)) + } + + turn([]string{"turn1.txt"}, 100, 10) + turn(nil, 40, 4) // changes nothing + + total := sessionTokenTotal(t, env, session.ID) + require.Equal(t, 140, total.InputTokens, "session total must include the no-change turn (100+40)") + require.Equal(t, 14, total.OutputTokens, "session total must include the no-change turn (10+4)") + + cp1, usage1 := commitCheckpointTokens(t, env, "", "Turn 1", "turn1.txt") + require.Equal(t, 140, usage1.InputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 14, usage1.OutputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 140, sessionTokenTotal(t, env, session.ID).InputTokens, "condensation must not change the session total") + + turn([]string{"turn3.txt"}, 7, 1) + _, usage2 := commitCheckpointTokens(t, env, cp1, "Turn 3", "turn3.txt") + require.Equal(t, 7, usage2.InputTokens, "checkpoint 2 must hold turn 3 only") + require.Equal(t, 1, usage2.OutputTokens, "checkpoint 2 must hold turn 3 only") + total = sessionTokenTotal(t, env, session.ID) + require.Equal(t, 147, total.InputTokens, "session total = every turn once") + require.Equal(t, 15, total.OutputTokens, "session total = every turn once") +} + +func TestCodexTokenUsage_TurnWithoutFileChanges(t *testing.T) { + t.Parallel() + + env := NewFeatureBranchEnv(t) + const sessionID = "codex-nochange-session" + rollout := filepath.Join(env.RepoDir, ".entire", "tmp", "codex-nochange-rollout.jsonl") + appendLines(t, rollout, + `{"timestamp":"2026-01-01T00:00:00Z","type":"session_meta","payload":{"id":"`+sessionID+`","thread_source":"user","cwd":"`+env.RepoDir+`"}}`) + hook := codexHooker(t, env.RepoDir, sessionID, rollout) + + // Codex reports cumulative totals; each turn appends the new cumulative + // snapshot, so the per-turn delta is the difference from the previous one. + var cumInput, cumCached, cumOutput int + n := 0 + turn := func(files []string, input, cached, output int) { + t.Helper() + n++ + prompt := fmt.Sprintf("turn %d", n) + hook("user-prompt-submit", map[string]any{"prompt": prompt, "hook_event_name": "UserPromptSubmit"}) + cumInput += input + cumCached += cached + cumOutput += output + appendLines(t, rollout, + fmt.Sprintf(`{"timestamp":"2026-01-01T00:0%d:01Z","type":"response_item","payload":{"type":"message","role":"user","content":[{"type":"input_text","text":%q}]}}`, n, prompt), + fmt.Sprintf(`{"timestamp":"2026-01-01T00:0%d:02Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":%d,"cached_input_tokens":%d,"output_tokens":%d}}}}`, n, cumInput, cumCached, cumOutput), + fmt.Sprintf(`{"timestamp":"2026-01-01T00:0%d:03Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}}`, n)) + for i, name := range files { + env.WriteFile(name, prompt+"\n") + applyPatchHook(hook, fmt.Sprintf("call_%d_%d", n, i), "*** Begin Patch\n*** Add File: "+name+"\n+"+prompt+"\n*** End Patch\n") + } + hook("stop", map[string]any{"hook_event_name": "Stop"}) + } + + // Fresh input is input minus cached: turn 1 = 80, turn 2 = 30, turn 3 = 5. + turn([]string{"turn1.txt"}, 100, 20, 10) + turn(nil, 40, 10, 4) // changes nothing + + total := sessionTokenTotal(t, env, sessionID) + require.Equal(t, 110, total.InputTokens, "session total must include the no-change turn (80+30)") + require.Equal(t, 30, total.CacheReadTokens, "session total must include the no-change turn (20+10)") + require.Equal(t, 14, total.OutputTokens, "session total must include the no-change turn (10+4)") + + cp1, usage1 := commitCheckpointTokens(t, env, "", "Turn 1", "turn1.txt") + require.Equal(t, 110, usage1.InputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 14, usage1.OutputTokens, "checkpoint must hold both turns, the no-change turn once") + require.Equal(t, 110, sessionTokenTotal(t, env, sessionID).InputTokens, "condensation must not change the session total") + + turn([]string{"turn3.txt"}, 5, 0, 1) + _, usage2 := commitCheckpointTokens(t, env, cp1, "Turn 3", "turn3.txt") + require.Equal(t, 5, usage2.InputTokens, "checkpoint 2 must hold turn 3 only") + require.Equal(t, 1, usage2.OutputTokens, "checkpoint 2 must hold turn 3 only") + total = sessionTokenTotal(t, env, sessionID) + require.Equal(t, 115, total.InputTokens, "session total = every turn once") + require.Equal(t, 15, total.OutputTokens, "session total = every turn once") +} + +// requireTurnTokens asserts the session total and the pending checkpoint +// window both hold exactly the given usage. +func requireTurnTokens(t *testing.T, env *TestEnv, sessionID string, input, output int, msg string) { + t.Helper() + state, err := env.GetSessionState(sessionID) + require.NoError(t, err) + require.NotNil(t, state, "session state must exist") + require.NotNil(t, state.TokenUsage, "%s: session total", msg) + require.Equal(t, input, state.TokenUsage.InputTokens, "%s: session total input", msg) + require.Equal(t, output, state.TokenUsage.OutputTokens, "%s: session total output", msg) + require.NotNil(t, state.CheckpointTokenUsage, "%s: pending checkpoint usage", msg) + require.Equal(t, input, state.CheckpointTokenUsage.InputTokens, "%s: pending checkpoint input", msg) + require.Equal(t, output, state.CheckpointTokenUsage.OutputTokens, "%s: pending checkpoint output", msg) +} + +// runConcurrently runs each fn at once and returns their errors. +func runConcurrently(fns ...func() error) []error { + errs := make([]error, len(fns)) + var wg sync.WaitGroup + for i, fn := range fns { + wg.Add(1) + go func() { + defer wg.Done() + errs[i] = fn() + }() + } + wg.Wait() + return errs +} + +// TestClaudeCodeTokenUsage_RepeatedStopCountsNoChangeTurnOnce pins that a +// no-change turn's tokens are recorded by the Stop that ends the turn and by +// no other. A retried Stop runs after the pre-prompt state is gone, so it +// resolves its transcript offset from the session instead and would re-add +// every earlier turn; concurrent Stops would each add the turn. +func TestClaudeCodeTokenUsage_RepeatedStopCountsNoChangeTurnOnce(t *testing.T) { + t.Parallel() + + env := NewFeatureBranchEnv(t) + session := env.NewSession() + transcript := session.TranscriptPath + + n := 0 + startTurn := func(input, output int) { + t.Helper() + n++ + prompt := fmt.Sprintf("turn %d", n) + require.NoError(t, env.SimulateUserPromptSubmitWithPromptAndTranscriptPath(session.ID, prompt, transcript)) + appendLines(t, transcript, + fmt.Sprintf(`{"uuid":"user-%d","type":"user","message":{"content":%q}}`, n, prompt), + fmt.Sprintf(`{"uuid":"asst-%d","type":"assistant","message":{"id":"msg_%d","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":%d,"output_tokens":%d}}}`, n, n, input, output)) + } + stop := func() error { return env.SimulateStop(session.ID, transcript) } + + startTurn(100, 10) + require.NoError(t, stop()) + startTurn(40, 4) + require.NoError(t, stop()) + requireTurnTokens(t, env, session.ID, 140, 14, "after two no-change turns") + + require.NoError(t, stop(), "a repeated Stop must still succeed") + requireTurnTokens(t, env, session.ID, 140, 14, "a repeated Stop must add nothing") + + startTurn(7, 1) + for i, err := range runConcurrently(stop, stop) { + require.NoError(t, err, "concurrent Stop %d", i) + } + requireTurnTokens(t, env, session.ID, 147, 15, "concurrent Stops must record the turn once") +} diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index 118e0bdd62..a1eca7f59f 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1014,6 +1014,30 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev codexInventoryUsage, codexLedgerVersion = refreshCodexInventory(ctx, ag, sessionID, transcriptData, inventoryOffset) } + // Agents that run subagents as sessions of their own (Factory AI Droid's + // Workers) reach turn-end here for the subagent's own turn. Attribute that + // work to the parent task invocation instead of minting a top-level session + // checkpoint unrelated to the session the user is actually driving. This + // runs before the no-changes gate: a read-only Worker turn is still the + // parent's task, and a zero-file completed record is itself a condensation + // trigger (State.HasTaskContent), so its transcript and tokens reach the + // parent's next checkpoint. Its usage goes on that record only, never on + // the Worker's own session total, so it is counted once. + if link, isSubagent := resolveSubagentSessionLink(ctx, ag, transcriptRef); isSubagent { + return saveSubagentSessionTaskStep(ctx, subagentSessionStep{ + link: link, + sessionID: sessionID, + event: event, + transcriptRef: transcriptRef, + modifiedFiles: relModifiedFiles, + newFiles: relNewFiles, + deletedFiles: relDeletedFiles, + turnUsage: workerTurnTokenUsage(ctx, ag, event, preState, transcriptData), + ag: ag, + strat: GetStrategy(ctx), + }) + } + // Check if there are any changes totalChanges := len(relModifiedFiles) + len(relNewFiles) + len(relDeletedFiles) if totalChanges == 0 { @@ -1025,33 +1049,16 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev slog.String("error", recErr.Error())) } recordCaptureDegraded(ctx, sessionID, captureDegraded) - // SaveStep is skipped, but out-of-band token usage must still be - // recorded: an Antigravity turn that commits ALL its work mid-turn - // (its normal flow) ends with a clean tree, and the mid-turn - // condensation ran with a zero delta (the baseline only re-snapshots - // at TurnStart). Without this, CleanupPrePromptState deletes the - // baseline and the turn's tokens are lost permanently. - if oobUsage := computeOutOfBandTokenUsage(ctx, ag, sessionID, preState); oobUsage != nil { - if accErr := strategy.AccumulateSessionTokenUsage(ctx, sessionID, oobUsage); accErr != nil { - // This is the only path that records a checkpoint-less turn's - // tokens, so a swallowed failure here is a permanent loss. Name - // the two causes apart: a session whose state was removed - // between the turn-end transition and this accumulate (nothing - // left to attribute to) versus an I/O or lock failure on state - // that still exists. - if errors.Is(accErr, strategy.ErrStateNotFound) { - logging.Warn(logCtx, "session state already removed; out-of-band token usage for checkpoint-less turn not recorded", - slog.String("session_id", sessionID), - slog.Int("input_tokens", oobUsage.InputTokens), - slog.Int("output_tokens", oobUsage.OutputTokens)) - } else { - logging.Warn(logCtx, "failed to record out-of-band token usage for checkpoint-less turn", - slog.String("session_id", sessionID), - slog.String("error", accErr.Error())) - } - } - } - transitionSessionTurnEnd(ctx, sessionID, event) + // SaveStep is skipped, but the turn's tokens must still be recorded, + // resolved and accounted exactly as SaveStep would. Otherwise the + // session total misses the turn, and so does the next checkpoint for + // agents whose transcript carries no usage (Cursor's stop payload, + // Antigravity's out-of-band pipe). An Antigravity turn that commits ALL + // its work mid-turn (its normal flow) lands here with a clean tree. + // The tokens are applied by the turn-end transition itself, so a + // repeated or concurrent Stop for the same turn adds nothing. + turnUsage := resolveTurnTokenUsage(ctx, ag, event, sessionID, preState, transcriptData, subagentsDir, codexInventoryUsage) + transitionSessionTurnEnd(ctx, sessionID, event, &checkpointlessTurnUsage{usage: turnUsage, ledgerVersion: codexLedgerVersion}) if cleanupErr := CleanupPrePromptState(ctx, sessionID); cleanupErr != nil { logging.Warn(logCtx, "failed to cleanup pre-prompt state", slog.String("error", cleanupErr.Error())) @@ -1072,24 +1079,6 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev strat := GetStrategy(ctx) agentType := ag.Type() - // Agents that run subagents as sessions of their own (Factory AI Droid's - // Workers) reach turn-end here for the subagent's own turn. Attribute that - // work to the parent task invocation instead of minting a top-level session - // checkpoint unrelated to the session the user is actually driving. - if link, isSubagent := resolveSubagentSessionLink(ctx, ag, transcriptRef); isSubagent { - return saveSubagentSessionTaskStep(ctx, subagentSessionStep{ - link: link, - sessionID: sessionID, - event: event, - transcriptRef: transcriptRef, - modifiedFiles: relModifiedFiles, - newFiles: relNewFiles, - deletedFiles: relDeletedFiles, - agentType: agentType, - strat: strat, - }) - } - // Get transcript position/identifier from pre-prompt state var transcriptIdentifierAtStart string var transcriptLinesAtStart int @@ -1098,27 +1087,7 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev transcriptLinesAtStart = preState.TranscriptOffset } - // Resolve token usage. Hook-provided counts (e.g., Cursor's stop hook, - // which is the only authoritative source for Cursor sessions because the - // JSONL transcript has no usage fields) take precedence; otherwise fall - // back to transcript-based computation, preferring SubagentAwareExtractor - // to include subagent tokens. - tokenUsage := event.TokenUsage - if tokenUsage == nil { - if codexInventoryUsage != nil { - tokenUsage = codexInventoryUsage - } else { - tokenUsage = agent.CalculateTokenUsage(ctx, ag, transcriptData, transcriptLinesAtStart, subagentsDir) - } - } - - // Out-of-band fallback: Antigravity exposes token usage only via its - // title/statusline pipe (captured by the title-tee shim), never in the - // transcript. Delta = current cumulative totals minus the TurnStart - // baseline stored in PrePromptState. - if tokenUsage == nil { - tokenUsage = computeOutOfBandTokenUsage(ctx, ag, sessionID, preState) - } + tokenUsage := resolveTurnTokenUsage(ctx, ag, event, sessionID, preState, transcriptData, subagentsDir, codexInventoryUsage) // Build fully-populated step context and delegate to strategy stepCtx := strategy.StepContext{ @@ -1159,7 +1128,7 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev } } recordCaptureDegraded(ctx, sessionID, captureDegraded) - transitionSessionTurnEnd(ctx, sessionID, event) + transitionSessionTurnEnd(ctx, sessionID, event, nil) if cleanupErr := CleanupPrePromptState(ctx, sessionID); cleanupErr != nil { logging.Warn(logCtx, "failed to cleanup pre-prompt state", slog.String("error", cleanupErr.Error())) @@ -1167,6 +1136,49 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev return nil } +// workerTurnTokenUsage returns a detached subagent session's usage for the +// turn ending now: the hook-reported counts, or its transcript from the turn's +// start offset. Its own subagents are not counted, since a task record holds +// only the subagent's own usage. +func workerTurnTokenUsage(ctx context.Context, ag agent.Agent, event *agent.Event, preState *PrePromptState, transcriptData []byte) *agent.TokenUsage { + if event.TokenUsage != nil { + return event.TokenUsage + } + turnStart := 0 + if preState != nil { + turnStart = preState.TranscriptOffset + } + return agent.CalculateTokenUsage(ctx, ag, transcriptData, turnStart, "") +} + +// resolveTurnTokenUsage returns the token usage of the turn ending now. +// Hook-provided counts (e.g., Cursor's stop hook, which is the only +// authoritative source for Cursor sessions because the JSONL transcript has no +// usage fields) take precedence; otherwise fall back to Codex's child ledger, +// then to transcript-based computation from the turn's start offset (preferring +// SubagentAwareExtractor to include subagent tokens), and last to the +// out-of-band source: Antigravity exposes token usage only via its +// title/statusline pipe (captured by the title-tee shim), never in the +// transcript, so its delta is the current cumulative totals minus the +// TurnStart baseline stored in PrePromptState. Only one source is used, so a +// turn is never counted twice. +func resolveTurnTokenUsage(ctx context.Context, ag agent.Agent, event *agent.Event, sessionID string, preState *PrePromptState, transcriptData []byte, subagentsDir string, codexInventoryUsage *agent.TokenUsage) *agent.TokenUsage { + if event.TokenUsage != nil { + return event.TokenUsage + } + if codexInventoryUsage != nil { + return codexInventoryUsage + } + turnStart := 0 + if preState != nil { + turnStart = preState.TranscriptOffset + } + if usage := agent.CalculateTokenUsage(ctx, ag, transcriptData, turnStart, subagentsDir); usage != nil { + return usage + } + return computeOutOfBandTokenUsage(ctx, ag, sessionID, preState) +} + // 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. @@ -2374,7 +2386,8 @@ type subagentSessionStep struct { modifiedFiles []string newFiles []string deletedFiles []string - agentType types.AgentType + turnUsage *agent.TokenUsage + ag agent.Agent strat *strategy.ManualCommitStrategy } @@ -2389,6 +2402,12 @@ type subagentSessionStep struct { // subagent's own session ID doubles as the agent ID; the declared transcript // path is the Worker session's transcript, known here at attribution time — // if it is gone by condensation, the materializer records it unavailable. +// +// The record also carries the Worker's token usage. Each turn adds its own +// usage, scoped to the turn's transcript slice, and only the Stop that ends +// the Worker's active turn adds it, so a turn condensed into one checkpoint is +// never counted again by a later one, and none of it lands on the Worker's +// own session total. func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) error { logCtx := logging.WithComponent(ctx, "lifecycle") logging.Info(logCtx, "recording subagent session as task record", @@ -2401,12 +2420,21 @@ func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) // The parent may not have any session state yet when its Worker finishes // first. - if err := step.strat.EnsureSessionExists(ctx, step.link.ParentSessionID, step.agentType); err != nil { + if err := step.strat.EnsureSessionExists(ctx, step.link.ParentSessionID, step.ag.Type()); err != nil { return fmt.Errorf("failed to ensure parent session state: %w", err) } + // Retire the subagent's own session the same way an ordinary turn-end does, + // so its phase and pre-prompt state do not linger as an active session. + // Ending it first is what claims the turn's usage: a repeated Stop finds + // the Worker already IDLE and adds none. + var turnUsage *agent.TokenUsage + if transitionSessionTurnEnd(ctx, step.sessionID, step.event, nil) { + turnUsage = step.turnUsage + } + files := mergeUnique(mergeUnique(step.modifiedFiles, step.newFiles), step.deletedFiles) - if err := strategy.UpsertCompletedTaskRecord(ctx, step.link.ParentSessionID, session.TaskRecord{ + if err := strategy.UpsertCompletedTaskRecordTurn(ctx, step.link.ParentSessionID, session.TaskRecord{ ToolUseID: step.link.ToolUseID, AgentID: step.sessionID, StartedAt: time.Now(), @@ -2414,13 +2442,11 @@ func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) TaskDescription: step.link.TaskDescription, DeclaredTranscriptPath: step.transcriptRef, Files: files, + TokenUsage: turnUsage, }); err != nil { return fmt.Errorf("failed to record subagent session task: %w", err) } - // Retire the subagent's own session the same way an ordinary turn-end does, - // so its phase and pre-prompt state do not linger as an active session. - transitionSessionTurnEnd(ctx, step.sessionID, step.event) if cleanupErr := CleanupPrePromptState(ctx, step.sessionID); cleanupErr != nil { logging.Warn(logCtx, "failed to cleanup pre-prompt state", slog.String("error", cleanupErr.Error())) @@ -2494,11 +2520,20 @@ func recordCaptureDegraded(ctx context.Context, sessionID string, degraded bool) } // transitionSessionTurnEnd transitions the session phase to IDLE and dispatches turn-end actions. -func transitionSessionTurnEnd(ctx context.Context, sessionID string, event *agent.Event) { +func transitionSessionTurnEnd(ctx context.Context, sessionID string, event *agent.Event, turnUsage *checkpointlessTurnUsage) (endedActiveTurn bool) { logCtx := logging.WithComponent(ctx, "lifecycle") var appendedSkillEvents []agent.SkillEvent + wasActive := false mutErr := strategy.MutateSessionStateOnSaved(ctx, sessionID, func(state *strategy.SessionState) error { appendedSkillEvents = persistEventMetadataToState(event, state) + // Only the Stop that ends the ACTIVE turn records its tokens. A + // repeated, retried, or concurrent Stop for the same turn finds the + // session already IDLE (this closure runs under the session lock), so + // it adds nothing. + wasActive = state.Phase.IsActive() + if turnUsage != nil && wasActive { + strategy.ApplyTurnTokenUsage(state, turnUsage.usage, turnUsage.ledgerVersion) + } if err := strategy.TransitionAndLog(ctx, state, session.EventTurnEnd, session.TransitionContext{}, session.NoOpActionHandler{}); err != nil { logging.Warn(logCtx, "turn-end transition failed", slog.String("error", err.Error())) @@ -2530,6 +2565,41 @@ func transitionSessionTurnEnd(ctx context.Context, sessionID string, event *agen logging.Warn(logCtx, "failed to update session phase on turn end", slog.String("error", mutErr.Error())) } + // The closure's changes count only once they are saved. + endedActiveTurn = wasActive && mutErr == nil + if turnUsage != nil && turnUsage.usage != nil && !endedActiveTurn { + // This is the only path that records a checkpoint-less turn's tokens, + // so name why they were not: no state to attribute them to, a failed + // write, or no active turn to end. + switch { + case errors.Is(mutErr, strategy.ErrStateNotFound): + logging.Warn(logCtx, "session state already removed; token usage for checkpoint-less turn not recorded", + slog.String("session_id", sessionID), + slog.Int("input_tokens", turnUsage.usage.InputTokens), + slog.Int("output_tokens", turnUsage.usage.OutputTokens)) + case mutErr != nil: + logging.Warn(logCtx, "failed to record token usage for checkpoint-less turn", + slog.String("session_id", sessionID), + slog.String("error", mutErr.Error())) + default: + // Either an earlier Stop already ended this turn, or no TurnStart + // was recorded for it. The two look alike here, and with no turn + // start there is no offset to scope the usage to, so neither + // records it. + logging.Info(logCtx, "no active turn to end; token usage for checkpoint-less turn not recorded", + slog.String("session_id", sessionID)) + } + } + return endedActiveTurn +} + +// checkpointlessTurnUsage is the token usage of a turn that ends without a +// step, for transitionSessionTurnEnd to record when it ends the turn. +// ledgerVersion is the subagent ledger version the usage was computed +// against (Codex), or nil. +type checkpointlessTurnUsage struct { + usage *agent.TokenUsage + ledgerVersion *uint64 } // sessionEndedAt resolves the EndedAt stamp for a session being finalized under @@ -2635,7 +2705,7 @@ func computeOutOfBandTokenUsage(ctx context.Context, ag agent.Agent, sessionID s // the session's first tracked turn (no snapshot existed yet). Mid-session // — after earlier turns already accumulated deltas — it means the // PrePromptState was lost or corrupt: counting from zero would return - // session-cumulative totals and AccumulateSessionTokenUsage would re-add + // session-cumulative totals and ApplyTurnTokenUsage would re-add // tokens the earlier turns already recorded. Degrade to no-data instead. if len(baseline) == 0 { if state, stateErr := strategy.LoadSessionState(ctx, sessionID); stateErr == nil && diff --git a/cmd/entire/cli/lifecycle_test.go b/cmd/entire/cli/lifecycle_test.go index b9e73a7341..f9c9e36014 100644 --- a/cmd/entire/cli/lifecycle_test.go +++ b/cmd/entire/cli/lifecycle_test.go @@ -3309,7 +3309,7 @@ func TestSaveSubagentSessionTaskStep_SecondTurn_MergesAndKeepsDeclaredPath(t *te event: &agent.Event{Type: agent.TurnEnd, SessionID: "droid-worker", Timestamp: time.Now()}, transcriptRef: "/tmp/worker.jsonl", modifiedFiles: []string{"a.txt"}, - agentType: agent.AgentTypeClaudeCode, + ag: &mockLifecycleAgent{name: "mock", agentType: agent.AgentTypeClaudeCode}, strat: GetStrategy(ctx), } require.NoError(t, saveSubagentSessionTaskStep(ctx, step)) diff --git a/cmd/entire/cli/strategy/accumulate_token_usage_test.go b/cmd/entire/cli/strategy/accumulate_token_usage_test.go index d20c1d05b5..39636fc714 100644 --- a/cmd/entire/cli/strategy/accumulate_token_usage_test.go +++ b/cmd/entire/cli/strategy/accumulate_token_usage_test.go @@ -1,52 +1,75 @@ package strategy import ( - "context" "testing" - "time" "github.com/entireio/cli/cmd/entire/cli/agent" - agenttypes "github.com/entireio/cli/cmd/entire/cli/agent/types" - "github.com/entireio/cli/cmd/entire/cli/testutil" "github.com/stretchr/testify/require" ) -// TestAccumulateSessionTokenUsage verifies the helper mirrors SaveStep's token +// TestApplyTurnTokenUsage verifies the helper mirrors SaveStep's token // accounting (accumulate into BOTH the session-cumulative TokenUsage and the // checkpoint-scoped CheckpointTokenUsage). It exists for turns that end with // no uncommitted changes (e.g. Antigravity committing all its work mid-turn): -// SaveStep is skipped, but the turn's out-of-band token delta must still be -// recorded so the next condensation attributes it. -func TestAccumulateSessionTokenUsage(t *testing.T) { - dir := t.TempDir() - testutil.InitRepo(t, dir) - t.Chdir(dir) - - sessionID := "agy-accumulate-tokens" - now := time.Now() +// SaveStep is skipped, but the turn's token delta must still be recorded so +// the next condensation attributes it. +func TestApplyTurnTokenUsage(t *testing.T) { + t.Parallel() + + state := &SessionState{ + TokenUsage: &agent.TokenUsage{InputTokens: 100, OutputTokens: 10, APICallCount: 1}, + } + + ApplyTurnTokenUsage(state, &agent.TokenUsage{InputTokens: 50, OutputTokens: 5, APICallCount: 1}, nil) + + require.NotNil(t, state.TokenUsage) + require.Equal(t, 150, state.TokenUsage.InputTokens, "session-cumulative input") + require.Equal(t, 15, state.TokenUsage.OutputTokens, "session-cumulative output") + require.Equal(t, 2, state.TokenUsage.APICallCount, "session-cumulative calls") + require.NotNil(t, state.CheckpointTokenUsage, "checkpoint-scoped accumulator must be populated") + require.Equal(t, 50, state.CheckpointTokenUsage.InputTokens, "checkpoint-scoped input") + + // Nil delta is a no-op. + ApplyTurnTokenUsage(state, nil, nil) + require.Equal(t, 150, state.TokenUsage.InputTokens) + require.Equal(t, 50, state.CheckpointTokenUsage.InputTokens) +} + +// TestApplyTurnTokenUsage_SubagentTokensMatchSaveStep pins that a +// checkpoint-less turn's cumulative subagent snapshot is handled as SaveStep +// handles it: it replaces the session's cumulative child total, the pending +// checkpoint window gets it rescoped against the baseline (so children already +// condensed are not counted again), and a snapshot computed against a stale +// subagent ledger version is dropped. +func TestApplyTurnTokenUsage_SubagentTokensMatchSaveStep(t *testing.T) { + t.Parallel() + state := &SessionState{ - SessionID: sessionID, - AgentType: agenttypes.AgentType("Antigravity"), - StartedAt: now, // zero StartedAt would be auto-deleted as stale on load - LastInteractionTime: &now, - TokenUsage: &agent.TokenUsage{ - InputTokens: 100, OutputTokens: 10, APICallCount: 1, - }, + AgentType: agent.AgentTypeClaudeCode, + SubagentLedgerVersion: 2, + TokenUsage: &agent.TokenUsage{InputTokens: 100, SubagentTokens: &agent.TokenUsage{InputTokens: 30}}, + SubagentTokensBaseline: &agent.TokenUsage{InputTokens: 30}, + } + + ApplyTurnTokenUsage(state, &agent.TokenUsage{ + InputTokens: 10, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 45}, + }, nil) + + require.Equal(t, 110, state.TokenUsage.InputTokens) + require.NotNil(t, state.TokenUsage.SubagentTokens) + require.Equal(t, 45, state.TokenUsage.SubagentTokens.InputTokens, "session child total is the latest cumulative snapshot") + require.NotNil(t, state.CheckpointTokenUsage) + require.Equal(t, 10, state.CheckpointTokenUsage.InputTokens) + require.NotNil(t, state.CheckpointTokenUsage.SubagentTokens) + require.Equal(t, 15, state.CheckpointTokenUsage.SubagentTokens.InputTokens, "window child total is rescoped against the baseline (45-30)") + + stale := uint64(1) + ApplyTurnTokenUsage(state, &agent.TokenUsage{ + InputTokens: 5, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 999}, + }, &stale) + + require.Equal(t, 115, state.TokenUsage.InputTokens, "main-agent tokens survive a stale ledger version") + if state.TokenUsage.SubagentTokens != nil { + require.NotEqual(t, 999, state.TokenUsage.SubagentTokens.InputTokens, "a stale-ledger child snapshot must not be recorded") } - require.NoError(t, SaveSessionState(context.Background(), state)) - - delta := &agent.TokenUsage{InputTokens: 50, OutputTokens: 5, APICallCount: 1} - require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, delta)) - - got, err := LoadSessionState(context.Background(), sessionID) - require.NoError(t, err) - require.NotNil(t, got.TokenUsage) - require.Equal(t, 150, got.TokenUsage.InputTokens, "session-cumulative input") - require.Equal(t, 15, got.TokenUsage.OutputTokens, "session-cumulative output") - require.Equal(t, 2, got.TokenUsage.APICallCount, "session-cumulative calls") - require.NotNil(t, got.CheckpointTokenUsage, "checkpoint-scoped accumulator must be populated") - require.Equal(t, 50, got.CheckpointTokenUsage.InputTokens, "checkpoint-scoped input") - - // Nil delta is a no-op, not an error. - require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, nil)) } diff --git a/cmd/entire/cli/strategy/manual_commit_git.go b/cmd/entire/cli/strategy/manual_commit_git.go index 4fdc2b9ef1..f71e6d145b 100644 --- a/cmd/entire/cli/strategy/manual_commit_git.go +++ b/cmd/entire/cli/strategy/manual_commit_git.go @@ -97,52 +97,7 @@ func (s *ManualCommitStrategy) SaveStep(ctx context.Context, step StepContext) e if state.StepCount == 1 { state.TranscriptIdentifierAtStart = step.StepTranscriptIdentifier } - if step.TokenUsage != nil { - state.TokenUsage = accumulateTokenUsage(state.TokenUsage, step.TokenUsage) - state.CheckpointTokenUsage = accumulateTokenUsage(state.CheckpointTokenUsage, step.TokenUsage) - // step.TokenUsage.SubagentTokens is a cumulative-since-session-start - // snapshot (agent IDs are discovered from the full transcript and each - // subagent's own transcript is re-read from its start on every call — - // see CalculateTotalTokenUsage in the claudecode/factoryaidroid - // packages), not a per-step delta like the rest of TokenUsage. - // accumulateTokenUsage already replaces (rather than adds) the - // SubagentTokens field for that reason, so state.TokenUsage ends up - // correctly holding the latest cumulative total. CheckpointTokenUsage - // additionally needs rescoping to "since last condensation" by - // subtracting the baseline captured at the last reset, otherwise the - // full cumulative subagent total would be reported again at every - // checkpoint instead of just this checkpoint's share. - // - // Derive the checkpoint delta FRESH each call from the session-wide - // cumulative (state.TokenUsage.SubagentTokens) minus the baseline — - // do NOT mutate CheckpointTokenUsage.SubagentTokens in place. A later - // step in the same window can carry step.TokenUsage != nil but - // SubagentTokens == nil (the subagent transcript was cleaned up, so - // CalculateTotalTokenUsage returned APICallCount==0 and left it nil); - // accumulateTokenUsage then leaves CheckpointTokenUsage.SubagentTokens - // at its already-rescoped value, and re-subtracting the baseline from - // that would double-subtract and (via clampSubtract) shrink or zero a - // real subagent total. Recomputing from the session-wide cumulative - // is idempotent regardless of whether this step carried a snapshot. - if state.CheckpointTokenUsage != nil && state.TokenUsage != nil { - complete := state.TokenUsage.SubagentTokensComplete - switch { - case complete != nil && !*complete: - state.CheckpointTokenUsage = types.WithClearedSubagentTokens(state.CheckpointTokenUsage, false) - case state.SubagentTokensBaselineComplete != nil && !*state.SubagentTokensBaselineComplete: - // A known-incomplete baseline cannot yield an exact delta, even - // when the current inventory has become complete again. - state.CheckpointTokenUsage = types.WithClearedSubagentTokens(state.CheckpointTokenUsage, false) - default: - state.CheckpointTokenUsage.SubagentTokens = types.SubtractTokenUsage( - state.TokenUsage.SubagentTokens, state.SubagentTokensBaseline) - if complete != nil { - value := *complete - state.CheckpointTokenUsage.SubagentTokensComplete = &value - } - } - } - } + applyTurnTokenUsage(state, step.TokenUsage) logCtx := logging.WithComponent(ctx, "checkpoint") logging.Info(logCtx, "checkpoint saved", @@ -163,13 +118,75 @@ func (s *ManualCommitStrategy) SaveStep(ctx context.Context, step StepContext) e } func invalidateStaleSubagentSnapshot(step *StepContext, state *SessionState) { - if step.SubagentLedgerVersion == nil || step.TokenUsage == nil || - state.SubagentLedgerVersion == *step.SubagentLedgerVersion { - return + step.TokenUsage = withoutStaleSubagentSnapshot(step.TokenUsage, step.SubagentLedgerVersion, state) +} + +// withoutStaleSubagentSnapshot drops usage's child aggregate when it was +// computed against a subagent ledger version other than the one state holds. +func withoutStaleSubagentSnapshot(usage *agent.TokenUsage, ledgerVersion *uint64, state *SessionState) *agent.TokenUsage { + if ledgerVersion == nil || usage == nil || state.SubagentLedgerVersion == *ledgerVersion { + return usage } // Keep valid main-agent deltas but never persist a child aggregate // computed against an older authoritative inventory. - step.TokenUsage = types.WithClearedSubagentTokens(step.TokenUsage, false) + return types.WithClearedSubagentTokens(usage, false) +} + +// applyTurnTokenUsage records one turn's token usage on state: it adds the +// per-turn counters to both the session total (TokenUsage) and the pending +// checkpoint window (CheckpointTokenUsage), and rescopes the window's +// cumulative SubagentTokens against the baseline. SaveStep and +// ApplyTurnTokenUsage (turns that change no files) both go through +// here, so a turn is accounted the same way whichever path records it. A nil +// usage is a no-op. +func applyTurnTokenUsage(state *SessionState, usage *agent.TokenUsage) { + if usage == nil { + return + } + state.TokenUsage = accumulateTokenUsage(state.TokenUsage, usage) + state.CheckpointTokenUsage = accumulateTokenUsage(state.CheckpointTokenUsage, usage) + // usage.SubagentTokens is a cumulative-since-session-start + // snapshot (agent IDs are discovered from the full transcript and each + // subagent's own transcript is re-read from its start on every call — + // see CalculateTotalTokenUsage in the claudecode/factoryaidroid + // packages), not a per-step delta like the rest of TokenUsage. + // accumulateTokenUsage already replaces (rather than adds) the + // SubagentTokens field for that reason, so state.TokenUsage ends up + // correctly holding the latest cumulative total. CheckpointTokenUsage + // additionally needs rescoping to "since last condensation" by + // subtracting the baseline captured at the last reset, otherwise the + // full cumulative subagent total would be reported again at every + // checkpoint instead of just this checkpoint's share. + // + // Derive the checkpoint delta FRESH each call from the session-wide + // cumulative (state.TokenUsage.SubagentTokens) minus the baseline — + // do NOT mutate CheckpointTokenUsage.SubagentTokens in place. A later + // step in the same window can carry usage != nil but + // SubagentTokens == nil (the subagent transcript was cleaned up, so + // CalculateTotalTokenUsage returned APICallCount==0 and left it nil); + // accumulateTokenUsage then leaves CheckpointTokenUsage.SubagentTokens + // at its already-rescoped value, and re-subtracting the baseline from + // that would double-subtract and (via clampSubtract) shrink or zero a + // real subagent total. Recomputing from the session-wide cumulative + // is idempotent regardless of whether this step carried a snapshot. + if state.CheckpointTokenUsage != nil && state.TokenUsage != nil { + complete := state.TokenUsage.SubagentTokensComplete + switch { + case complete != nil && !*complete: + state.CheckpointTokenUsage = types.WithClearedSubagentTokens(state.CheckpointTokenUsage, false) + case state.SubagentTokensBaselineComplete != nil && !*state.SubagentTokensBaselineComplete: + // A known-incomplete baseline cannot yield an exact delta, even + // when the current inventory has become complete again. + state.CheckpointTokenUsage = types.WithClearedSubagentTokens(state.CheckpointTokenUsage, false) + default: + state.CheckpointTokenUsage.SubagentTokens = types.SubtractTokenUsage( + state.TokenUsage.SubagentTokens, state.SubagentTokensBaseline) + if complete != nil { + value := *complete + state.CheckpointTokenUsage.SubagentTokensComplete = &value + } + } + } } // ensureSessionInitialized creates the session state file if it doesn't yet @@ -310,9 +327,25 @@ func CompleteTaskRecord(ctx context.Context, sessionID string, rec session.TaskR // (Factory Droid Workers reach turn-end once per Worker turn, all attributed // to one parent tool use), unlike CompleteTaskRecord's exactly-once claim. func UpsertCompletedTaskRecord(ctx context.Context, sessionID string, rec session.TaskRecord) error { + return upsertCompletedTaskRecord(ctx, sessionID, rec, false) +} + +// UpsertCompletedTaskRecordTurn is UpsertCompletedTaskRecord for a producer +// that reports usage per turn (Factory Droid Workers): rec.TokenUsage is this +// turn's delta and is added to the record's usage rather than replacing it. +// Condensation removes a completed record once stored, so a later turn starts +// a fresh record holding only usage no checkpoint has yet. +func UpsertCompletedTaskRecordTurn(ctx context.Context, sessionID string, rec session.TaskRecord) error { + return upsertCompletedTaskRecord(ctx, sessionID, rec, true) +} + +func upsertCompletedTaskRecord(ctx context.Context, sessionID string, rec session.TaskRecord, addUsage bool) error { return MutateSessionState(ctx, sessionID, func(state *SessionState) error { if existing := state.FindTaskRecord(rec.ToolUseID); existing != nil { rec.Files = mergeFilesTouched(existing.Files, rec.Files) + if addUsage && rec.TokenUsage != nil { + rec.TokenUsage = types.AddTokenUsage(existing.TokenUsage, rec.TokenUsage) + } } else { state.AddTaskRecord(launchStubTaskRecord(rec)) } diff --git a/cmd/entire/cli/strategy/session_state.go b/cmd/entire/cli/strategy/session_state.go index 6064c3496e..5516f7ec26 100644 --- a/cmd/entire/cli/strategy/session_state.go +++ b/cmd/entire/cli/strategy/session_state.go @@ -1022,21 +1022,16 @@ func ClearSessionState(ctx context.Context, sessionID string) error { return nil } -// AccumulateSessionTokenUsage adds a token-usage delta to both the -// session-cumulative total and the checkpoint-scoped accumulator, mirroring -// SaveStep's accounting (manual_commit_git.go). It exists for turns that end -// with no uncommitted changes — e.g. Antigravity committing all its work -// mid-turn, its normal flow — where the TurnEnd handler skips SaveStep -// entirely but the turn's out-of-band token delta must still be recorded so -// the next condensation (or `entire status`) attributes it. A nil or empty -// delta is a no-op. -func AccumulateSessionTokenUsage(ctx context.Context, sessionID string, delta *agent.TokenUsage) error { - if delta == nil { - return nil - } - return MutateSessionState(ctx, sessionID, func(state *SessionState) error { - state.TokenUsage = accumulateTokenUsage(state.TokenUsage, delta) - state.CheckpointTokenUsage = accumulateTokenUsage(state.CheckpointTokenUsage, delta) - return nil - }) +// ApplyTurnTokenUsage records a turn's token usage on state exactly as SaveStep +// does (applyTurnTokenUsage): into the session total and the pending +// checkpoint window, with SubagentTokens rescoped against the baseline. It +// exists for turns that end with no uncommitted changes, where the TurnEnd +// handler skips SaveStep but the turn's tokens must still be recorded so +// `entire status` and the next condensation attribute them. The caller applies +// it inside the MutateSessionState that ends the turn, so the turn's phase +// decides whether it is still unaccounted. ledgerVersion is the subagent +// ledger version the usage was computed against (Codex), or nil; a stale one +// drops the child aggregate, as it does for a step. A nil delta is a no-op. +func ApplyTurnTokenUsage(state *SessionState, delta *agent.TokenUsage, ledgerVersion *uint64) { + applyTurnTokenUsage(state, withoutStaleSubagentSnapshot(delta, ledgerVersion, state)) }