From 1ea72520d81953f9d7794e54f1c01f4e67440284 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 18:29:35 -0700 Subject: [PATCH 1/4] lifecycle: keep hook-reported tokens for turns that change no files A turn-end that changes no files skips SaveStep, and that branch only recorded out-of-band usage (Antigravity). The hook-reported usage on the event (Cursor's stop payload) was dropped. Cursor's transcript has no usage fields, so condensation could not recover it either: the turn was lost from the session total and from the next checkpoint. Record event.TokenUsage on that branch, falling back to out-of-band usage only when the hook reported none, the same precedence the SaveStep path uses, so no turn is counted twice. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4HPW4188RCPJJK89HX2XPNM --- .../cursor_token_condensation_test.go | 101 ++++++++++++++++++ cmd/entire/cli/lifecycle.go | 35 +++--- 2 files changed, 124 insertions(+), 12 deletions(-) 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..993dfc9e8a 100644 --- a/cmd/entire/cli/integration_test/cursor_token_condensation_test.go +++ b/cmd/entire/cli/integration_test/cursor_token_condensation_test.go @@ -203,6 +203,107 @@ 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") +} + 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/lifecycle.go b/cmd/entire/cli/lifecycle.go index ffcee3e285..0ed50b8c82 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1021,14 +1021,25 @@ 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 { + // SaveStep is skipped, but the turn's token usage must still be + // recorded when no transcript can reproduce it later. Two sources, + // in the same precedence as the SaveStep path below: + // - Hook-reported usage (Cursor's stop payload). Cursor's transcript + // has no usage fields, so condensation falls back to + // CheckpointTokenUsage; dropping the event here loses the turn + // from both the session total and the next checkpoint. + // - Out-of-band usage (Antigravity). A 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. + // Only one is used, so a turn is never counted twice. + turnUsage := event.TokenUsage + if turnUsage == nil { + turnUsage = computeOutOfBandTokenUsage(ctx, ag, sessionID, preState) + } + if turnUsage != nil { + if accErr := strategy.AccumulateSessionTokenUsage(ctx, sessionID, turnUsage); 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 @@ -1036,12 +1047,12 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev // 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", + logging.Warn(logCtx, "session state already removed; 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)) + slog.Int("input_tokens", turnUsage.InputTokens), + slog.Int("output_tokens", turnUsage.OutputTokens)) } else { - logging.Warn(logCtx, "failed to record out-of-band token usage for checkpoint-less turn", + logging.Warn(logCtx, "failed to record token usage for checkpoint-less turn", slog.String("session_id", sessionID), slog.String("error", accErr.Error())) } From db6ebd3893fed5c575b9611d7f488a82c868e0fd Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 19:14:38 -0700 Subject: [PATCH 2/4] lifecycle: record transcript-derived tokens for turns that change no files The checkpoint-less turn-end path recorded only hook-reported and out-of-band usage, so for Claude Code, Codex and the other agents whose usage lives in the transcript, a turn that changed no files never reached the session total that `entire status` shows. (The next checkpoint was unaffected: condensation recomputes it from the transcript slice.) Resolve the turn's usage on that path exactly as SaveStep does, with one helper (resolveTurnTokenUsage: hook, Codex child ledger, transcript from the turn's start offset, out-of-band), and record it through the same accounting SaveStep uses, now extracted as applyTurnTokenUsage: the session total, the pending checkpoint window, the subagent baseline rescoping, and dropping a child snapshot computed against a stale Codex ledger version. Subagent sessions of their own (Droid Workers) are skipped, as SaveStep never runs for them either. Co-Authored-By: Claude Opus 5.5 --- .../turn_without_changes_tokens_test.go | 159 ++++++++++++++++++ cmd/entire/cli/lifecycle.go | 80 ++++----- .../strategy/accumulate_token_usage_test.go | 55 +++++- cmd/entire/cli/strategy/manual_commit_git.go | 117 +++++++------ cmd/entire/cli/strategy/session_state.go | 22 +-- 5 files changed, 332 insertions(+), 101 deletions(-) create mode 100644 cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go 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..79f363d5ca --- /dev/null +++ b/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go @@ -0,0 +1,159 @@ +//go:build integration + +package integration + +import ( + "fmt" + "os" + "path/filepath" + "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") +} diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index 0ed50b8c82..512d203170 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1021,25 +1021,21 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev slog.String("error", recErr.Error())) } recordCaptureDegraded(ctx, sessionID, captureDegraded) - // SaveStep is skipped, but the turn's token usage must still be - // recorded when no transcript can reproduce it later. Two sources, - // in the same precedence as the SaveStep path below: - // - Hook-reported usage (Cursor's stop payload). Cursor's transcript - // has no usage fields, so condensation falls back to - // CheckpointTokenUsage; dropping the event here loses the turn - // from both the session total and the next checkpoint. - // - Out-of-band usage (Antigravity). A 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. - // Only one is used, so a turn is never counted twice. - turnUsage := event.TokenUsage - if turnUsage == nil { - turnUsage = computeOutOfBandTokenUsage(ctx, ag, sessionID, preState) + // 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. + // A subagent's own session (Factory AI Droid's Workers) is skipped: + // with changes it never reaches SaveStep either, because its work is + // the parent's task, so recording here would count it twice. + var turnUsage *agent.TokenUsage + if _, isSubagent := resolveSubagentSessionLink(ctx, ag, transcriptRef); !isSubagent { + turnUsage = resolveTurnTokenUsage(ctx, ag, event, sessionID, preState, transcriptData, subagentsDir, codexInventoryUsage) } if turnUsage != nil { - if accErr := strategy.AccumulateSessionTokenUsage(ctx, sessionID, turnUsage); accErr != nil { + if accErr := strategy.AccumulateSessionTokenUsage(ctx, sessionID, turnUsage, codexLedgerVersion); 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 @@ -1105,27 +1101,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{ @@ -1174,6 +1150,34 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev return nil } +// 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. diff --git a/cmd/entire/cli/strategy/accumulate_token_usage_test.go b/cmd/entire/cli/strategy/accumulate_token_usage_test.go index d20c1d05b5..12feec3e71 100644 --- a/cmd/entire/cli/strategy/accumulate_token_usage_test.go +++ b/cmd/entire/cli/strategy/accumulate_token_usage_test.go @@ -36,7 +36,7 @@ func TestAccumulateSessionTokenUsage(t *testing.T) { require.NoError(t, SaveSessionState(context.Background(), state)) delta := &agent.TokenUsage{InputTokens: 50, OutputTokens: 5, APICallCount: 1} - require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, delta)) + require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, delta, nil)) got, err := LoadSessionState(context.Background(), sessionID) require.NoError(t, err) @@ -48,5 +48,56 @@ func TestAccumulateSessionTokenUsage(t *testing.T) { 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)) + require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, nil, nil)) +} + +// TestAccumulateSessionTokenUsage_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 TestAccumulateSessionTokenUsage_SubagentTokensMatchSaveStep(t *testing.T) { + dir := t.TempDir() + testutil.InitRepo(t, dir) + t.Chdir(dir) + ctx := context.Background() + + sessionID := "subagent-accumulate-tokens" + now := time.Now() + require.NoError(t, SaveSessionState(ctx, &SessionState{ + SessionID: sessionID, + AgentType: agent.AgentTypeClaudeCode, + StartedAt: now, + LastInteractionTime: &now, + SubagentLedgerVersion: 2, + TokenUsage: &agent.TokenUsage{InputTokens: 100, SubagentTokens: &agent.TokenUsage{InputTokens: 30}}, + SubagentTokensBaseline: &agent.TokenUsage{InputTokens: 30}, + })) + + require.NoError(t, AccumulateSessionTokenUsage(ctx, sessionID, &agent.TokenUsage{ + InputTokens: 10, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 45}, + }, nil)) + + got, err := LoadSessionState(ctx, sessionID) + require.NoError(t, err) + require.Equal(t, 110, got.TokenUsage.InputTokens) + require.NotNil(t, got.TokenUsage.SubagentTokens) + require.Equal(t, 45, got.TokenUsage.SubagentTokens.InputTokens, "session child total is the latest cumulative snapshot") + require.NotNil(t, got.CheckpointTokenUsage) + require.Equal(t, 10, got.CheckpointTokenUsage.InputTokens) + require.NotNil(t, got.CheckpointTokenUsage.SubagentTokens) + require.Equal(t, 15, got.CheckpointTokenUsage.SubagentTokens.InputTokens, "window child total is rescoped against the baseline (45-30)") + + stale := uint64(1) + require.NoError(t, AccumulateSessionTokenUsage(ctx, sessionID, &agent.TokenUsage{ + InputTokens: 5, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 999}, + }, &stale)) + + got, err = LoadSessionState(ctx, sessionID) + require.NoError(t, err) + require.Equal(t, 115, got.TokenUsage.InputTokens, "main-agent tokens survive a stale ledger version") + if got.TokenUsage.SubagentTokens != nil { + require.NotEqual(t, 999, got.TokenUsage.SubagentTokens.InputTokens, "a stale-ledger child snapshot must not be recorded") + } } diff --git a/cmd/entire/cli/strategy/manual_commit_git.go b/cmd/entire/cli/strategy/manual_commit_git.go index dd6768e32e..d82dba0972 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 +// AccumulateSessionTokenUsage (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 diff --git a/cmd/entire/cli/strategy/session_state.go b/cmd/entire/cli/strategy/session_state.go index 6064c3496e..94cd59d9d2 100644 --- a/cmd/entire/cli/strategy/session_state.go +++ b/cmd/entire/cli/strategy/session_state.go @@ -1022,21 +1022,21 @@ 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 { +// AccumulateSessionTokenUsage records a turn's token usage 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. 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 AccumulateSessionTokenUsage(ctx context.Context, sessionID string, delta *agent.TokenUsage, ledgerVersion *uint64) 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) + applyTurnTokenUsage(state, withoutStaleSubagentSnapshot(delta, ledgerVersion, state)) return nil }) } From e8cf4320a69af0b7e1376307c82b9f64673af7d4 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Wed, 14 Oct 2026 10:36:08 -0700 Subject: [PATCH 3/4] turn-end: record a no-change turn's tokens once, and route read-only Workers A no-change turn's tokens were accumulated before the turn-end transition, with no check that the turn was still open, so a repeated, retried, or concurrent Stop added them again. A retried Claude Code Stop was worse: with the pre-prompt state already cleaned up it resolved its transcript offset from the session and re-added every earlier turn. The tokens are now applied inside the MutateSessionState that performs the ACTIVE->IDLE transition, and only while the session is ACTIVE, so the second Stop finds it IDLE and adds nothing. A Factory AI Droid Worker turn that changed no files returned through the no-changes gate before the Worker routing, so it left no task record on the parent and its transcript and tokens never reached the parent checkpoint. Workers are now resolved before that gate, and their record carries the Worker's usage computed from its whole transcript (recounted at condensation), so it lands on the parent's task once and never on the Worker's own session. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4XQRQ3HNAPCABZXM4XG28D6 --- .../cursor_token_condensation_test.go | 74 +++++++++- .../factoryai_worker_session_test.go | 87 ++++++++++++ .../turn_without_changes_tokens_test.go | 71 ++++++++++ cmd/entire/cli/lifecycle.go | 134 +++++++++++------- cmd/entire/cli/lifecycle_test.go | 2 +- .../strategy/accumulate_token_usage_test.go | 100 +++++-------- cmd/entire/cli/strategy/manual_commit_git.go | 2 +- cmd/entire/cli/strategy/session_state.go | 21 ++- 8 files changed, 358 insertions(+), 133 deletions(-) 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 993dfc9e8a..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) { @@ -304,6 +317,63 @@ func TestCursorTokenUsage_TurnWithoutFileChanges(t *testing.T) { 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..251657c6d6 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,88 @@ 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") + } +} + // 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/turn_without_changes_tokens_test.go b/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go index 79f363d5ca..c3b12f454b 100644 --- a/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go +++ b/cmd/entire/cli/integration_test/turn_without_changes_tokens_test.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "sync" "testing" "github.com/entireio/cli/cmd/entire/cli/agent" @@ -157,3 +158,73 @@ func TestCodexTokenUsage_TurnWithoutFileChanges(t *testing.T) { 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 b3ee6971e5..444e03b721 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1014,6 +1014,29 @@ 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, + ag: ag, + strat: GetStrategy(ctx), + }) + } + // Check if there are any changes totalChanges := len(relModifiedFiles) + len(relNewFiles) + len(relDeletedFiles) if totalChanges == 0 { @@ -1031,34 +1054,10 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev // 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. - // A subagent's own session (Factory AI Droid's Workers) is skipped: - // with changes it never reaches SaveStep either, because its work is - // the parent's task, so recording here would count it twice. - var turnUsage *agent.TokenUsage - if _, isSubagent := resolveSubagentSessionLink(ctx, ag, transcriptRef); !isSubagent { - turnUsage = resolveTurnTokenUsage(ctx, ag, event, sessionID, preState, transcriptData, subagentsDir, codexInventoryUsage) - } - if turnUsage != nil { - if accErr := strategy.AccumulateSessionTokenUsage(ctx, sessionID, turnUsage, codexLedgerVersion); 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; token usage for checkpoint-less turn not recorded", - slog.String("session_id", sessionID), - slog.Int("input_tokens", turnUsage.InputTokens), - slog.Int("output_tokens", turnUsage.OutputTokens)) - } else { - logging.Warn(logCtx, "failed to record token usage for checkpoint-less turn", - slog.String("session_id", sessionID), - slog.String("error", accErr.Error())) - } - } - } - transitionSessionTurnEnd(ctx, sessionID, event) + // 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())) @@ -1079,24 +1078,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 @@ -1146,7 +1127,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())) @@ -2389,7 +2370,7 @@ type subagentSessionStep struct { modifiedFiles []string newFiles []string deletedFiles []string - agentType types.AgentType + ag agent.Agent strat *strategy.ManualCommitStrategy } @@ -2404,6 +2385,11 @@ 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, computed from its whole +// transcript (recounted at condensation, like a foreground task's), so every +// Worker turn replaces it rather than adding to it, 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", @@ -2416,11 +2402,15 @@ 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) } files := mergeUnique(mergeUnique(step.modifiedFiles, step.newFiles), step.deletedFiles) + tokenUsage := step.event.TokenUsage + if tokenUsage == nil { + tokenUsage = subagentTokenUsage(logCtx, step.ag, step.event, step.transcriptRef) + } if err := strategy.UpsertCompletedTaskRecord(ctx, step.link.ParentSessionID, session.TaskRecord{ ToolUseID: step.link.ToolUseID, AgentID: step.sessionID, @@ -2429,13 +2419,16 @@ func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) TaskDescription: step.link.TaskDescription, DeclaredTranscriptPath: step.transcriptRef, Files: files, + TokenUsage: tokenUsage, + // The usage is the whole transcript's, so condensation may recount it. + TokenUsageFromTranscript: step.event.TokenUsage == nil, }); 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) + transitionSessionTurnEnd(ctx, step.sessionID, step.event, nil) if cleanupErr := CleanupPrePromptState(ctx, step.sessionID); cleanupErr != nil { logging.Warn(logCtx, "failed to cleanup pre-prompt state", slog.String("error", cleanupErr.Error())) @@ -2509,11 +2502,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) { logCtx := logging.WithComponent(ctx, "lifecycle") var appendedSkillEvents []agent.SkillEvent + usageApplied := 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. + if turnUsage != nil && turnUsage.usage != nil && state.Phase.IsActive() { + strategy.ApplyTurnTokenUsage(state, turnUsage.usage, turnUsage.ledgerVersion) + usageApplied = true + } 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())) @@ -2545,6 +2547,34 @@ 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())) } + if turnUsage != nil && turnUsage.usage != nil && !usageApplied { + // 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 a turn already ended by an earlier Stop. + 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: + logging.Info(logCtx, "turn already ended; not recording its token usage again", + slog.String("session_id", sessionID)) + } + } +} + +// 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 @@ -2650,7 +2680,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 12feec3e71..39636fc714 100644 --- a/cmd/entire/cli/strategy/accumulate_token_usage_test.go +++ b/cmd/entire/cli/strategy/accumulate_token_usage_test.go @@ -1,103 +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) +// 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() - sessionID := "agy-accumulate-tokens" - now := time.Now() 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, - }, + TokenUsage: &agent.TokenUsage{InputTokens: 100, OutputTokens: 10, APICallCount: 1}, } - require.NoError(t, SaveSessionState(context.Background(), state)) - delta := &agent.TokenUsage{InputTokens: 50, OutputTokens: 5, APICallCount: 1} - require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, delta, nil)) + ApplyTurnTokenUsage(state, &agent.TokenUsage{InputTokens: 50, OutputTokens: 5, APICallCount: 1}, nil) - 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") + 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, not an error. - require.NoError(t, AccumulateSessionTokenUsage(context.Background(), sessionID, nil, nil)) + // Nil delta is a no-op. + ApplyTurnTokenUsage(state, nil, nil) + require.Equal(t, 150, state.TokenUsage.InputTokens) + require.Equal(t, 50, state.CheckpointTokenUsage.InputTokens) } -// TestAccumulateSessionTokenUsage_SubagentTokensMatchSaveStep pins that a +// 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 TestAccumulateSessionTokenUsage_SubagentTokensMatchSaveStep(t *testing.T) { - dir := t.TempDir() - testutil.InitRepo(t, dir) - t.Chdir(dir) - ctx := context.Background() +func TestApplyTurnTokenUsage_SubagentTokensMatchSaveStep(t *testing.T) { + t.Parallel() - sessionID := "subagent-accumulate-tokens" - now := time.Now() - require.NoError(t, SaveSessionState(ctx, &SessionState{ - SessionID: sessionID, + state := &SessionState{ AgentType: agent.AgentTypeClaudeCode, - StartedAt: now, - LastInteractionTime: &now, SubagentLedgerVersion: 2, TokenUsage: &agent.TokenUsage{InputTokens: 100, SubagentTokens: &agent.TokenUsage{InputTokens: 30}}, SubagentTokensBaseline: &agent.TokenUsage{InputTokens: 30}, - })) + } - require.NoError(t, AccumulateSessionTokenUsage(ctx, sessionID, &agent.TokenUsage{ + ApplyTurnTokenUsage(state, &agent.TokenUsage{ InputTokens: 10, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 45}, - }, nil)) + }, nil) - got, err := LoadSessionState(ctx, sessionID) - require.NoError(t, err) - require.Equal(t, 110, got.TokenUsage.InputTokens) - require.NotNil(t, got.TokenUsage.SubagentTokens) - require.Equal(t, 45, got.TokenUsage.SubagentTokens.InputTokens, "session child total is the latest cumulative snapshot") - require.NotNil(t, got.CheckpointTokenUsage) - require.Equal(t, 10, got.CheckpointTokenUsage.InputTokens) - require.NotNil(t, got.CheckpointTokenUsage.SubagentTokens) - require.Equal(t, 15, got.CheckpointTokenUsage.SubagentTokens.InputTokens, "window child total is rescoped against the baseline (45-30)") + 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) - require.NoError(t, AccumulateSessionTokenUsage(ctx, sessionID, &agent.TokenUsage{ + ApplyTurnTokenUsage(state, &agent.TokenUsage{ InputTokens: 5, APICallCount: 1, SubagentTokens: &agent.TokenUsage{InputTokens: 999}, - }, &stale)) + }, &stale) - got, err = LoadSessionState(ctx, sessionID) - require.NoError(t, err) - require.Equal(t, 115, got.TokenUsage.InputTokens, "main-agent tokens survive a stale ledger version") - if got.TokenUsage.SubagentTokens != nil { - require.NotEqual(t, 999, got.TokenUsage.SubagentTokens.InputTokens, "a stale-ledger child snapshot must not be recorded") + 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") } } diff --git a/cmd/entire/cli/strategy/manual_commit_git.go b/cmd/entire/cli/strategy/manual_commit_git.go index dcacd0719a..3562bbd340 100644 --- a/cmd/entire/cli/strategy/manual_commit_git.go +++ b/cmd/entire/cli/strategy/manual_commit_git.go @@ -136,7 +136,7 @@ func withoutStaleSubagentSnapshot(usage *agent.TokenUsage, ledgerVersion *uint64 // 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 -// AccumulateSessionTokenUsage (turns that change no files) both go through +// 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) { diff --git a/cmd/entire/cli/strategy/session_state.go b/cmd/entire/cli/strategy/session_state.go index 94cd59d9d2..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 records a turn's token usage exactly as SaveStep +// 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. 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 AccumulateSessionTokenUsage(ctx context.Context, sessionID string, delta *agent.TokenUsage, ledgerVersion *uint64) error { - if delta == nil { - return nil - } - return MutateSessionState(ctx, sessionID, func(state *SessionState) error { - applyTurnTokenUsage(state, withoutStaleSubagentSnapshot(delta, ledgerVersion, state)) - return nil - }) +// `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)) } From 179dd3dbd1cd601dc29483917145d9d0849e29ca Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Wed, 14 Oct 2026 13:47:00 -0700 Subject: [PATCH 4/4] turn-end: scope Worker task usage per turn and fix the applied check A multi-turn Droid Worker's record carried whole-transcript usage that condensation recounted from line 0, and condensation drops a completed record once stored, so a turn already in one checkpoint was counted again by the next. Each Worker turn now adds only its own slice (from the turn's start offset) to the record, and only when ending the Worker's own turn claims it, so a repeated Stop adds nothing; the record is no longer marked for a whole-transcript recount. Usage counted as applied only after the turn-end mutation saved, so a failed save is reported. The log for a Stop with no active turn no longer claims it was a duplicate: with no turn start there is no offset to scope the usage to, so it is not recorded in either case. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4Y2P7H5M6JSVCYPKHR2MXPS --- .../factoryai_worker_session_test.go | 25 +++++++ cmd/entire/cli/integration_test/hooks.go | 16 ++++- cmd/entire/cli/lifecycle.go | 69 +++++++++++++------ cmd/entire/cli/strategy/manual_commit_git.go | 16 +++++ 4 files changed, 103 insertions(+), 23 deletions(-) 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 251657c6d6..94182f5dd2 100644 --- a/cmd/entire/cli/integration_test/factoryai_worker_session_test.go +++ b/cmd/entire/cli/integration_test/factoryai_worker_session_test.go @@ -213,6 +213,31 @@ func TestFactoryDroidReadOnlyWorkerReachesParentCheckpoint(t *testing.T) { 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 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/lifecycle.go b/cmd/entire/cli/lifecycle.go index 444e03b721..a1eca7f59f 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -1032,6 +1032,7 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev modifiedFiles: relModifiedFiles, newFiles: relNewFiles, deletedFiles: relDeletedFiles, + turnUsage: workerTurnTokenUsage(ctx, ag, event, preState, transcriptData), ag: ag, strat: GetStrategy(ctx), }) @@ -1135,6 +1136,21 @@ 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 @@ -2370,6 +2386,7 @@ type subagentSessionStep struct { modifiedFiles []string newFiles []string deletedFiles []string + turnUsage *agent.TokenUsage ag agent.Agent strat *strategy.ManualCommitStrategy } @@ -2386,10 +2403,11 @@ type subagentSessionStep struct { // 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, computed from its whole -// transcript (recounted at condensation, like a foreground task's), so every -// Worker turn replaces it rather than adding to it, and none of it lands on -// the Worker's own session total. +// 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", @@ -2406,12 +2424,17 @@ func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) return fmt.Errorf("failed to ensure parent session state: %w", err) } - files := mergeUnique(mergeUnique(step.modifiedFiles, step.newFiles), step.deletedFiles) - tokenUsage := step.event.TokenUsage - if tokenUsage == nil { - tokenUsage = subagentTokenUsage(logCtx, step.ag, step.event, step.transcriptRef) + // 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 } - if err := strategy.UpsertCompletedTaskRecord(ctx, step.link.ParentSessionID, session.TaskRecord{ + + files := mergeUnique(mergeUnique(step.modifiedFiles, step.newFiles), step.deletedFiles) + if err := strategy.UpsertCompletedTaskRecordTurn(ctx, step.link.ParentSessionID, session.TaskRecord{ ToolUseID: step.link.ToolUseID, AgentID: step.sessionID, StartedAt: time.Now(), @@ -2419,16 +2442,11 @@ func saveSubagentSessionTaskStep(ctx context.Context, step subagentSessionStep) TaskDescription: step.link.TaskDescription, DeclaredTranscriptPath: step.transcriptRef, Files: files, - TokenUsage: tokenUsage, - // The usage is the whole transcript's, so condensation may recount it. - TokenUsageFromTranscript: step.event.TokenUsage == nil, + 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, nil) if cleanupErr := CleanupPrePromptState(ctx, step.sessionID); cleanupErr != nil { logging.Warn(logCtx, "failed to cleanup pre-prompt state", slog.String("error", cleanupErr.Error())) @@ -2502,19 +2520,19 @@ 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, turnUsage *checkpointlessTurnUsage) { +func transitionSessionTurnEnd(ctx context.Context, sessionID string, event *agent.Event, turnUsage *checkpointlessTurnUsage) (endedActiveTurn bool) { logCtx := logging.WithComponent(ctx, "lifecycle") var appendedSkillEvents []agent.SkillEvent - usageApplied := false + 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. - if turnUsage != nil && turnUsage.usage != nil && state.Phase.IsActive() { + wasActive = state.Phase.IsActive() + if turnUsage != nil && wasActive { strategy.ApplyTurnTokenUsage(state, turnUsage.usage, turnUsage.ledgerVersion) - usageApplied = true } if err := strategy.TransitionAndLog(ctx, state, session.EventTurnEnd, session.TransitionContext{}, session.NoOpActionHandler{}); err != nil { logging.Warn(logCtx, "turn-end transition failed", @@ -2547,10 +2565,12 @@ 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())) } - if turnUsage != nil && turnUsage.usage != nil && !usageApplied { + // 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 a turn already ended by an earlier Stop. + // 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", @@ -2562,10 +2582,15 @@ func transitionSessionTurnEnd(ctx context.Context, sessionID string, event *agen slog.String("session_id", sessionID), slog.String("error", mutErr.Error())) default: - logging.Info(logCtx, "turn already ended; not recording its token usage again", + // 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 diff --git a/cmd/entire/cli/strategy/manual_commit_git.go b/cmd/entire/cli/strategy/manual_commit_git.go index 3562bbd340..f71e6d145b 100644 --- a/cmd/entire/cli/strategy/manual_commit_git.go +++ b/cmd/entire/cli/strategy/manual_commit_git.go @@ -327,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)) }