diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index c2176b1c77..a6a3f1e90b 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -9,6 +9,7 @@ import ( "log/slog" "os" "os/exec" + "reflect" "slices" "strings" "time" @@ -366,11 +367,7 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa return fmt.Errorf("failed to get git author: %w", err) } - tokenUsage := agent.CalculateTokenUsage(logCtx, ag, transcriptData, window.start, "") - sessionTokens := tokenUsage - if window.start > 0 { - sessionTokens = agent.CalculateTokenUsage(logCtx, ag, transcriptData, 0, "") - } + tokenUsage, sessionUsage, tokenPos := attachTokens(logCtx, ag, existingState, transcriptData, transcriptPath) // attach writes checkpoints and historically never configured // redaction; a scanner-config failure must fail the attach. @@ -403,7 +400,7 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa // the new HEAD. Seeding BaseCommit makes the session link future commits // on HEAD; an attach to an older commit is about that commit only. seedBase := target.Hash.Equal(headCommit.Hash) - if err := saveAttachSessionState(logCtx, repo, existingState, sessionID, ag.Type(), transcriptPath, agentHome, checkpointID, meta, sessionTokens, window.end, opts, reviewSkills, seedBase); err != nil { + if err := saveAttachSessionState(logCtx, repo, existingState, sessionID, ag.Type(), transcriptPath, agentHome, checkpointID, meta, sessionUsage, tokenPos, window.end, opts, reviewSkills, seedBase); err != nil { logging.Warn(logCtx, "failed to save session state", "error", err) } else if activeHome { // The active home was resolved from the user's environment, so it may @@ -415,6 +412,26 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa return linkErr } +// attachTokens returns the tokens the attach checkpoint stores, the session's +// whole-transcript total for its state, and the token position to consume. +// The checkpoint counts from the token offset rather than the displayed +// window: carry-forward leaves the window at the session start while the +// tokens before it are already in a checkpoint. +func attachTokens(ctx context.Context, ag agent.Agent, existingState *session.State, transcriptData []byte, transcriptPath string) (checkpointUsage, sessionUsage *agent.TokenUsage, tokenPos int) { + checkpointUsage, tokenPos = strategy.AttachTokenUsage(ctx, ag, existingState, transcriptData, transcriptPath) + if existingState == nil { + return checkpointUsage, checkpointUsage, tokenPos + } + sessionUsage = agent.CalculateTokenUsage(ctx, ag, transcriptData, 0, "") + if sessionUsage != nil && existingState.TokenUsage != nil && existingState.TokenUsage.SubagentTokens != nil { + // The cumulative subagent total, including subagents of a running + // turn that AttachTokenUsage just read. + sessionUsage.SubagentTokens = existingState.TokenUsage.SubagentTokens + sessionUsage.SubagentTokensComplete = existingState.TokenUsage.SubagentTokensComplete + } + return checkpointUsage, sessionUsage, tokenPos +} + // attachCheckpoint is the checkpoint an attach writes into. type attachCheckpoint struct { // repo replaces the caller's handle when fetching the checkpoint refreshed it. @@ -1365,7 +1382,7 @@ func resolveCheckpointID(ctx context.Context, headCommit *object.Commit) (id.Che // If existingState is non-nil, it is updated in place (avoids a redundant disk load). // reviewSkills is the resolved skills list when opts.Review is true; ignored otherwise. // agentHome replaces State.AgentHome; "" clears it. -func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingState *session.State, sessionID string, agentType types.AgentType, transcriptPath, agentHome string, checkpointID id.CheckpointID, meta transcriptMetadata, tokenUsage *agent.TokenUsage, transcriptEnd int, opts attachOptions, reviewSkills []string, seedBase bool) error { +func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingState *session.State, sessionID string, agentType types.AgentType, transcriptPath, agentHome string, checkpointID id.CheckpointID, meta transcriptMetadata, sessionUsage *agent.TokenUsage, tokenPos, transcriptEnd int, opts attachOptions, reviewSkills []string, seedBase bool) error { stateStore, err := session.NewStateStore(ctx) if err != nil { return fmt.Errorf("failed to open session store: %w", err) @@ -1384,6 +1401,11 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS StartedAt: now, } } + // Read before the display offset below moves: TokenStart falls back to it + // for state written before the token offset existed. + tokenStartOnDisk := state.TokenStart() + pendingOnDisk := state.CheckpointTokenUsage + lastCheckpointOnDisk := state.LastCheckpointID // Populate BaseCommit from HEAD if not already set, so the session becomes // active and future commits in the same session receive Entire-Checkpoint trailers. @@ -1420,8 +1442,53 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS if meta.FirstPrompt != "" { state.LastPrompt = meta.FirstPrompt } - if tokenUsage != nil { - state.TokenUsage = tokenUsage + // A hook that ran while attach waited at its prompt recorded newer token + // state than attach read: a condensation moved the offset or stored the + // pending total (a new checkpoint ID), or a Stop added a turn to the + // pending total. Attach's figures are older than the hook's, so they are + // left out: consuming from the stale position would move the offset back + // and have the next checkpoint recount, and clearing the pending total + // would drop a hook-only agent's turn. Turn start clears LastCheckpointID, + // which is not a checkpoint, so only a new non-empty ID counts. + countedFrom := 0 + var pendingRead *agent.TokenUsage + var lastCheckpointRead id.CheckpointID + if existingState != nil { + countedFrom = existingState.TokenStart() + pendingRead = existingState.CheckpointTokenUsage + lastCheckpointRead = existingState.LastCheckpointID + } + switch { + case tokenStartOnDisk != countedFrom, + !reflect.DeepEqual(pendingOnDisk, pendingRead), + lastCheckpointOnDisk != "" && lastCheckpointOnDisk != lastCheckpointRead: + logging.Warn(ctx, "attach: session tokens changed while attach was waiting; leaving token state as the hooks set it", + slog.Int("counted_from", countedFrom), slog.Int("token_start", tokenStartOnDisk)) + default: + if state.Phase.IsActive() { + // The running turn's Stop adds its tokens from turn start, including + // those before this attach, so the whole-transcript total would count + // them twice. Keep the hooks' total; take only the cumulative subagent + // total, which the re-baseline below must cover. + if sessionUsage != nil && sessionUsage.SubagentTokens != nil { + total := agent.TokenUsage{} + if state.TokenUsage != nil { + total = *state.TokenUsage + } + total.SubagentTokens = sessionUsage.SubagentTokens + total.SubagentTokensComplete = sessionUsage.SubagentTokensComplete + state.TokenUsage = &total + } + } else if sessionUsage != nil { + // Without a cumulative subagent total of its own, keep the one hooks + // recorded so the re-baseline doesn't drop it. + if sessionUsage.SubagentTokens == nil && state.TokenUsage != nil { + sessionUsage.SubagentTokens = state.TokenUsage.SubagentTokens + sessionUsage.SubagentTokensComplete = state.TokenUsage.SubagentTokensComplete + } + state.TokenUsage = sessionUsage + } + strategy.ConsumeAttachTokenWindow(state, tokenPos) } if opts.Review { state.Kind = session.KindAgentReview diff --git a/cmd/entire/cli/attach_token_window_test.go b/cmd/entire/cli/attach_token_window_test.go new file mode 100644 index 0000000000..e5b1f28697 --- /dev/null +++ b/cmd/entire/cli/attach_token_window_test.go @@ -0,0 +1,89 @@ +package cli + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" + "github.com/entireio/cli/cmd/entire/cli/session" + "github.com/entireio/cli/cmd/entire/cli/testutil" + "github.com/stretchr/testify/require" +) + +// A hook that checkpoints while attach waits at its prompt moves the token +// offset and records newer totals; attach must not move the offset back or +// overwrite those totals with the older ones it computed before the wait. +func TestSaveAttachSessionState_KeepsTokenOffsetMovedMeanwhile(t *testing.T) { + dir := t.TempDir() + testutil.InitRepo(t, dir) + testutil.WriteFile(t, dir, "f.txt", "init") + testutil.GitAdd(t, dir, "f.txt") + testutil.GitCommit(t, dir, "init") + t.Chdir(dir) + ctx := context.Background() + + store, err := session.NewStateStore(ctx) + require.NoError(t, err) + cpID := id.MustCheckpointID("a1b2c3d4e5f6") + + for _, tc := range []struct { + name string + // hookMovedTo is the token offset on disk when attach saves; -1 means + // the session had no state at all. + hookMovedTo int + // pendingOnDisk is the pending total on disk; attach read 3. + pendingOnDisk int + // checkpointOnDisk is LastCheckpointID on disk; attach read none. + checkpointOnDisk id.CheckpointID + wantStart int + wantPendingOK bool + wantTotal int + }{ + // The session is running: its Stop adds the turn, so attach keeps the hooks' total. + {name: "unchanged", hookMovedTo: 2, pendingOnDisk: 3, wantStart: 4, wantTotal: 12}, + {name: "moved by a hook", hookMovedTo: 6, pendingOnDisk: 3, wantStart: 6, wantPendingOK: true, wantTotal: 12}, + {name: "turn added by a stop", hookMovedTo: 2, pendingOnDisk: 5, wantStart: 2, wantPendingOK: true, wantTotal: 12}, + {name: "condensed in place", hookMovedTo: 2, pendingOnDisk: 3, checkpointOnDisk: id.MustCheckpointID("b1b2c3d4e5f6"), wantStart: 2, wantPendingOK: true, wantTotal: 12}, + {name: "no earlier state", hookMovedTo: -1, wantStart: 4, wantTotal: 9}, + } { + sessionID := "attach-window-" + strings.ReplaceAll(tc.name, " ", "-") + var loaded *session.State + if tc.hookMovedTo >= 0 { + loaded = &session.State{ + SessionID: sessionID, + AgentType: agent.AgentTypeClaudeCode, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 3}, + } + loaded.SetTokenStart(2) + + now := time.Now() + onDisk := &session.State{ + SessionID: sessionID, + AgentType: agent.AgentTypeClaudeCode, + Phase: session.PhaseActive, + StartedAt: now, + LastInteractionTime: &now, + TokenUsage: &agent.TokenUsage{OutputTokens: 12}, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: tc.pendingOnDisk}, + LastCheckpointID: tc.checkpointOnDisk, + } + onDisk.SetTokenStart(tc.hookMovedTo) + require.NoError(t, store.Save(ctx, onDisk)) + } + + // transcriptEnd 7 moves an ended session's display offset, which a + // state without its own token offset would otherwise be read through. + require.NoError(t, saveAttachSessionState(ctx, nil, loaded, sessionID, agent.AgentTypeClaudeCode, + "/t.jsonl", "", cpID, transcriptMetadata{}, &agent.TokenUsage{OutputTokens: 9}, 4, 7, attachOptions{}, nil, false), tc.name) + + got, err := store.Load(ctx, sessionID) + require.NoError(t, err) + require.Equal(t, tc.wantStart, got.TokenStart(), tc.name) + require.Equal(t, tc.wantPendingOK, got.CheckpointTokenUsage != nil, tc.name) + require.NotNil(t, got.TokenUsage, tc.name) + require.Equal(t, tc.wantTotal, got.TokenUsage.OutputTokens, tc.name) + } +} diff --git a/cmd/entire/cli/integration_test/token_scope_test.go b/cmd/entire/cli/integration_test/token_scope_test.go new file mode 100644 index 0000000000..f41305fdd3 --- /dev/null +++ b/cmd/entire/cli/integration_test/token_scope_test.go @@ -0,0 +1,322 @@ +//go:build integration + +package integration + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/checkpoint" + "github.com/entireio/cli/cmd/entire/cli/paths" + "github.com/stretchr/testify/require" +) + +// appendUsageMessage appends an assistant message carrying token usage, the +// shape Claude Code writes, so the transcript's token total is known. +func appendUsageMessage(s *Session, messageID string, tokens int) { + s.TranscriptBuilder.messages = append(s.TranscriptBuilder.messages, map[string]interface{}{ + "uuid": "usage-" + messageID, + "type": roleAssistant, + "message": map[string]interface{}{ + "id": messageID, + "model": "claude-test", + "content": []map[string]interface{}{{"type": blockTypeText, "text": "working"}}, + "usage": map[string]interface{}{"input_tokens": tokens, "output_tokens": tokens}, + }, + }) +} + +// TestTokenScope_PartialCommitDoesNotRecountEarlierTurns: a partial commit +// carries the remaining files forward and restarts the next checkpoint's +// transcript window at line 0 so it shows the conversation behind them. Its +// tokens must still count only the turns no earlier checkpoint has counted; +// the server sums per-checkpoint tokens into the session total. +func TestTokenScope_PartialCommitDoesNotRecountEarlierTurns(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + + // Turn 1 (100 tokens) writes a and b; only a is committed. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("a.txt", "a") + env.WriteFile("b.txt", "b") + appendUsageMessage(s, "msg-1", 100) + s.CreateTranscript("make a and b", []FileChange{{Path: "a.txt", Content: "a"}, {Path: "b.txt", Content: "b"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("a only", "a.txt") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // Turn 2 (7 tokens) writes c; b and c are committed. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("c.txt", "c") + appendUsageMessage(s, "msg-2", 7) + s.CreateTranscript("make c", []FileChange{{Path: "c.txt", Content: "c"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("b and c", "b.txt", "c.txt") + cp2 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp2) + require.NotEqual(t, cp1, cp2) + + // Turn 3 (3 tokens) writes d and e; only d is committed, then e. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("d.txt", "d") + env.WriteFile("e.txt", "e") + appendUsageMessage(s, "msg-3", 3) + s.CreateTranscript("make d and e", []FileChange{{Path: "d.txt", Content: "d"}, {Path: "e.txt", Content: "e"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("d only", "d.txt") + cp3 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp3) + + u1 := readCommittedTokenUsage(t, env, cp1) + u2 := readCommittedTokenUsage(t, env, cp2) + u3 := readCommittedTokenUsage(t, env, cp3) + require.NotNil(t, u1) + require.NotNil(t, u2) + require.NotNil(t, u3) + require.Equal(t, 100, u1.OutputTokens, "checkpoint 1 counts turn 1") + require.Equal(t, 7, u2.OutputTokens, "checkpoint 2 counts turn 2 only, not turn 1 again") + require.Equal(t, 3, u3.OutputTokens, "checkpoint 3 counts turn 3 only") + require.Equal(t, 7, u2.InputTokens) + + // The carried-forward checkpoint still shows the conversation behind b. + transcript, found := env.ReadFileFromBranch(paths.MetadataBranchName, SessionFilePath(cp2, paths.TranscriptFileName)) + require.True(t, found, "checkpoint 2 transcript should exist") + require.Contains(t, transcript, "make a and b") + metaJSON, found := env.ReadFileFromBranch(paths.MetadataBranchName, SessionMetadataPath(cp2)) + require.True(t, found) + var meta checkpoint.Metadata + require.NoError(t, json.Unmarshal([]byte(metaJSON), &meta)) + require.Equal(t, 0, meta.CheckpointTranscriptStart, + "carry-forward keeps the displayed window at the start of the session") +} + +// TestTokenScope_AttachCountsOnlyUncheckpointedTokens: attaching a session +// whose earlier turns are already in a checkpoint stores only the later +// turns' tokens, and the session's next hook checkpoint doesn't count the +// attached turns again. +func TestTokenScope_AttachCountsOnlyUncheckpointedTokens(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + + // Turn 1 (100 tokens) writes a and b; only a is committed. Carry-forward + // clears LastCheckpointID, so attach will write a new checkpoint. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("a.txt", "a") + env.WriteFile("b.txt", "b") + appendUsageMessage(s, "msg-1", 100) + s.CreateTranscript("make a and b", []FileChange{{Path: "a.txt", Content: "a"}, {Path: "b.txt", Content: "b"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("a only", "a.txt") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // Turn 2 (7 tokens) changes no files. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + appendUsageMessage(s, "msg-2", 7) + s.CreateTranscript("explain b", nil) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + + // Attach resolves the transcript from the agent's project dir and records + // that path; keep writing the session there from now on. + transcriptData, err := os.ReadFile(s.TranscriptPath) + require.NoError(t, err) + s.TranscriptPath = filepath.Join(env.ClaudeProjectDir, s.ID+".jsonl") + require.NoError(t, os.WriteFile(s.TranscriptPath, transcriptData, 0o600)) + env.WriteFile("notes.txt", "notes") + env.GitAdd("notes.txt") + env.GitCommit("notes") + + // attach -f amends HEAD, which runs the installed git hooks; they resolve + // "entire" through PATH, so point it at the binary under test. + env.ExtraEnv = append(env.ExtraEnv, + "PATH="+filepath.Dir(getTestBinary())+string(os.PathListSeparator)+os.Getenv("PATH")) + output := env.RunCLI("session", "attach", s.ID, "-a", agentClaudeCode, "-f") + require.Contains(t, output, "Attached session") + attached := env.TryGetLatestCheckpointID() + require.NotEmpty(t, attached) + require.NotEqual(t, cp1, attached) + ua := readCommittedTokenUsage(t, env, attached) + require.NotNil(t, ua) + require.Equal(t, 7, ua.OutputTokens, "attach counts turn 2 only; turn 1 is in checkpoint 1") + + // Turn 3 (3 tokens) writes c; b and c are committed through hooks. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("c.txt", "c") + appendUsageMessage(s, "msg-3", 3) + s.CreateTranscript("make c", []FileChange{{Path: "c.txt", Content: "c"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("b and c", "b.txt", "c.txt") + cp3 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp3) + require.NotEqual(t, attached, cp3) + u3 := readCommittedTokenUsage(t, env, cp3) + require.NotNil(t, u3) + require.Equal(t, 3, u3.OutputTokens, "the next checkpoint counts turn 3 only, not the attached turn 2") +} + +// TestTokenScope_MidTurnCommitTailCountsInNextCheckpoint: when the agent +// commits mid-turn, the rest of the turn is written after condensation. Stop +// moves the displayed window past that tail, but its tokens are in no +// checkpoint yet, so the next checkpoint must count them. +func TestTokenScope_MidTurnCommitTailCountsInNextCheckpoint(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + + require.NoError(t, env.SimulateUserPromptSubmitWithPromptAndTranscriptPath(s.ID, "create and commit", s.TranscriptPath)) + env.WriteFile("feature.go", "package feature\n") + s.TranscriptBuilder.AddUserMessage("create and commit") + appendUsageMessage(s, "msg-1", 100) + toolID := s.TranscriptBuilder.AddToolUse("mcp__acp__Write", "feature.go", "package feature\n") + s.TranscriptBuilder.AddToolResult(toolID) + require.NoError(t, s.TranscriptBuilder.WriteToFile(s.TranscriptPath)) + env.GitCommitWithHooksAsAgent("Add feature", "feature.go") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // The turn continues after the commit (5 tokens), then Stop fires. + appendUsageMessage(s, "msg-tail", 5) + require.NoError(t, s.TranscriptBuilder.WriteToFile(s.TranscriptPath)) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + + // Turn 2 (7 tokens) writes another file, committed by the user. + require.NoError(t, env.SimulateUserPromptSubmitWithPromptAndTranscriptPath(s.ID, "add more", s.TranscriptPath)) + env.WriteFile("more.go", "package feature\n") + appendUsageMessage(s, "msg-2", 7) + s.CreateTranscript("add more", []FileChange{{Path: "more.go", Content: "package feature\n"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("more", "more.go") + cp2 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp2) + require.NotEqual(t, cp1, cp2) + + u1 := readCommittedTokenUsage(t, env, cp1) + u2 := readCommittedTokenUsage(t, env, cp2) + require.NotNil(t, u1) + require.NotNil(t, u2) + require.Equal(t, 100, u1.OutputTokens) + require.Equal(t, 12, u2.OutputTokens, "checkpoint 2 counts turn 1's post-commit tail (5) and turn 2 (7)") +} + +// TestTokenScope_AttachIntoOwnCheckpointKeepsItsTokens: when HEAD's trailer +// already names a checkpoint holding this session, attach leaves that entry +// alone. The entry keeps its tokens, and the turn attach didn't record is +// counted by the session's next checkpoint, so nothing is lost or counted twice. +func TestTokenScope_AttachIntoOwnCheckpointKeepsItsTokens(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("a.txt", "a") + appendUsageMessage(s, "msg-1", 100) + s.CreateTranscript("make a", []FileChange{{Path: "a.txt", Content: "a"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("a", "a.txt") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // Turn 2 (7 tokens) changes no files; turn start clears LastCheckpointID, + // so only the checkpoint's own contents show it already holds the session. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + appendUsageMessage(s, "msg-2", 7) + s.CreateTranscript("explain", nil) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + + transcriptData, err := os.ReadFile(s.TranscriptPath) + require.NoError(t, err) + s.TranscriptPath = filepath.Join(env.ClaudeProjectDir, s.ID+".jsonl") + require.NoError(t, os.WriteFile(s.TranscriptPath, transcriptData, 0o600)) + env.ExtraEnv = append(env.ExtraEnv, + "PATH="+filepath.Dir(getTestBinary())+string(os.PathListSeparator)+os.Getenv("PATH")) + output := env.RunCLI("session", "attach", s.ID, "-a", agentClaudeCode, "-f") + require.Contains(t, output, "is already in checkpoint") + require.Equal(t, cp1, env.TryGetLatestCheckpointID()) + + usage := readCommittedTokenUsage(t, env, cp1) + require.NotNil(t, usage) + require.Equal(t, 100, usage.OutputTokens, "cp1 keeps its 100 tokens") + + // Turn 3 (3 tokens) writes c; its checkpoint also carries turn 2's 7. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("c.txt", "c") + appendUsageMessage(s, "msg-3", 3) + s.CreateTranscript("make c", []FileChange{{Path: "c.txt", Content: "c"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("c", "c.txt") + cp3 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp3) + require.NotEqual(t, cp1, cp3) + u3 := readCommittedTokenUsage(t, env, cp3) + require.NotNil(t, u3) + require.Equal(t, 10, u3.OutputTokens, "the next checkpoint counts turn 2 (7) and turn 3 (3)") +} + +// TestTokenScope_AttachDuringActiveTurnConsumesTokensOnly: attaching a session +// mid-turn stores the turn's tokens so far and marks them counted, so the +// turn's own checkpoint doesn't count them again. The displayed window belongs +// to the running session's hooks and stays where it was. +func TestTokenScope_AttachDuringActiveTurnConsumesTokensOnly(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + env.WriteFile("a.txt", "a") + appendUsageMessage(s, "msg-1", 100) + s.CreateTranscript("make a", []FileChange{{Path: "a.txt", Content: "a"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("a", "a.txt") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // Turn 2 starts and spends 7 tokens; Stop hasn't fired, so it's active. + require.NoError(t, env.SimulateUserPromptSubmit(s.ID)) + appendUsageMessage(s, "msg-2", 7) + s.CreateTranscript("make c", nil) + transcriptData, err := os.ReadFile(s.TranscriptPath) + require.NoError(t, err) + s.TranscriptPath = filepath.Join(env.ClaudeProjectDir, s.ID+".jsonl") + require.NoError(t, os.WriteFile(s.TranscriptPath, transcriptData, 0o600)) + + before, err := env.GetSessionState(s.ID) + require.NoError(t, err) + require.True(t, before.Phase.IsActive()) + + env.WriteFile("notes.txt", "notes") + env.GitAdd("notes.txt") + env.GitCommit("notes") + env.ExtraEnv = append(env.ExtraEnv, + "PATH="+filepath.Dir(getTestBinary())+string(os.PathListSeparator)+os.Getenv("PATH")) + output := env.RunCLI("session", "attach", s.ID, "-a", agentClaudeCode, "-f") + require.Contains(t, output, "Attached session") + attached := env.TryGetLatestCheckpointID() + require.NotEqual(t, cp1, attached) + ua := readCommittedTokenUsage(t, env, attached) + require.NotNil(t, ua) + require.Equal(t, 7, ua.OutputTokens, "attach stores turn 2's tokens so far") + + after, err := env.GetSessionState(s.ID) + require.NoError(t, err) + require.True(t, after.Phase.IsActive(), "attach must not end a running session") + require.Equal(t, before.CheckpointTranscriptStart, after.CheckpointTranscriptStart, + "a running session's displayed window belongs to its hooks") + + // The turn continues (3 tokens), writes c, and is committed through hooks. + env.WriteFile("c.txt", "c") + appendUsageMessage(s, "msg-3", 3) + s.CreateTranscript("make c", []FileChange{{Path: "c.txt", Content: "c"}}) + require.NoError(t, env.SimulateStop(s.ID, s.TranscriptPath)) + env.GitCommitWithHooks("c", "c.txt") + cp3 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp3) + require.NotEqual(t, attached, cp3) + u3 := readCommittedTokenUsage(t, env, cp3) + require.NotNil(t, u3) + require.Equal(t, 3, u3.OutputTokens, "the turn's checkpoint counts only what attach didn't store") +} diff --git a/cmd/entire/cli/integration_test/token_turns_cursor_test.go b/cmd/entire/cli/integration_test/token_turns_cursor_test.go new file mode 100644 index 0000000000..46b877d181 --- /dev/null +++ b/cmd/entire/cli/integration_test/token_turns_cursor_test.go @@ -0,0 +1,215 @@ +//go:build integration + +package integration + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/execx" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// agentTokenStep is what one checkpoint must hold. +type agentTokenStep struct { + id string + step string + want agent.TokenUsage +} + +// tokenFields compares only the four token counters. +func tokenFields(u *agent.TokenUsage) [4]int { + if u == nil { + return [4]int{} + } + return [4]int{u.InputTokens, u.OutputTokens, u.CacheReadTokens, u.CacheCreationTokens} +} + +func addFields(a *agent.TokenUsage, b agent.TokenUsage) { + a.InputTokens += b.InputTokens + a.OutputTokens += b.OutputTokens + a.CacheReadTokens += b.CacheReadTokens + a.CacheCreationTokens += b.CacheCreationTokens +} + +// recordCheckpoint appends the latest checkpoint, requiring it to be new. +func recordCheckpoint(t *testing.T, env *TestEnv, steps []agentTokenStep, step string, want agent.TokenUsage) []agentTokenStep { + t.Helper() + cp := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp, step) + for _, s := range steps { + require.NotEqual(t, s.id, cp, "%s: expected a new checkpoint", step) + } + return append(steps, agentTokenStep{id: cp, step: step, want: want}) +} + +func verifyAgentSteps(t *testing.T, env *TestEnv, steps []agentTokenStep, wantSum agent.TokenUsage) { + t.Helper() + var sum agent.TokenUsage + for _, s := range steps { + got := readCommittedTokenUsage(t, env, s.id) + t.Logf("%s: checkpoint %s tokens=%v want=%v", s.step, s.id, tokenFields(got), tokenFields(&s.want)) + assert.Equal(t, tokenFields(&s.want), tokenFields(got), + "%s: checkpoint must hold exactly the turns it newly covers [in,out,cacheRead,cacheWrite]", s.step) + if got != nil { + addFields(&sum, *got) + } + } + assert.Equal(t, tokenFields(&wantSum), tokenFields(&sum), "checkpoints must sum to the attributable tokens") +} + +func sessionTotal(t *testing.T, env *TestEnv, sessionID string) *agent.TokenUsage { + t.Helper() + st, err := env.GetSessionState(sessionID) + require.NoError(t, err) + require.NotNil(t, st) + return st.TokenUsage +} + +// attachNow commits an unrelated file without hooks and attaches the session. +func attachNow(t *testing.T, env *TestEnv, sessionID, agentName string, n int) { + t.Helper() + name := fmt.Sprintf("notes-%d.txt", n) + env.WriteFile(name, "notes") + env.GitAdd(name) + env.GitCommit("notes " + name) + output := env.RunCLI("session", "attach", sessionID, "-a", agentName, "-f") + require.Contains(t, output, "Attached session", output) +} + +type cursorTokenRun struct { + t *testing.T + env *TestEnv + projectDir string + id string + transcriptPath string + turn int +} + +func newCursorTokenRun(t *testing.T) *cursorTokenRun { + t.Helper() + env := NewFeatureBranchEnv(t) + env.InitEntireWithAgent(agent.AgentNameCursor) + projectDir := t.TempDir() + if resolved, err := filepath.EvalSymlinks(projectDir); err == nil { + projectDir = resolved + } + env.ExtraEnv = append(env.ExtraEnv, + "ENTIRE_TEST_CURSOR_PROJECT_DIR="+projectDir, + "PATH="+filepath.Dir(getTestBinary())+string(os.PathListSeparator)+os.Getenv("PATH")) + id := "cursor-turns-session" + dir := filepath.Join(projectDir, id) + require.NoError(t, os.MkdirAll(dir, 0o755)) + r := &cursorTokenRun{t: t, env: env, projectDir: projectDir, id: id, transcriptPath: filepath.Join(dir, id+".jsonl")} + r.hook("session-start", map[string]any{"model": "cursor-default"}) + return r +} + +func (r *cursorTokenRun) hook(name string, extra map[string]any) { + r.t.Helper() + in := map[string]any{"conversation_id": r.id, "transcript_path": r.transcriptPath} + for k, v := range extra { + in[k] = v + } + data, err := json.Marshal(in) + require.NoError(r.t, err) + cmd := execx.NonInteractive(context.Background(), getTestBinary(), "hooks", "cursor", name) + cmd.Dir = r.env.RepoDir + cmd.Stdin = bytes.NewReader(data) + cmd.Env = r.env.cliEnv() + out, err := cmd.CombinedOutput() + require.NoErrorf(r.t, err, "cursor %s: %s", name, out) +} + +func (r *cursorTokenRun) appendTranscript(s string) { + r.t.Helper() + f, err := os.OpenFile(r.transcriptPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600) + require.NoError(r.t, err) + _, werr := f.WriteString(s) + require.NoError(r.t, f.Close()) + require.NoError(r.t, werr) +} + +func (r *cursorTokenRun) start(files ...string) { + r.t.Helper() + r.turn++ + prompt := fmt.Sprintf("turn %d", r.turn) + r.appendTranscript(fmt.Sprintf(`{"type":"user","text":%q}`+"\n", prompt)) + r.hook("before-submit-prompt", map[string]any{"prompt": prompt}) + for _, f := range files { + r.env.WriteFile(f, prompt+"\n") + } + r.appendTranscript(`{"type":"assistant","text":"ok"}` + "\n") +} + +// stop ends the turn with per-turn tokens; Cursor reports input as the total +// including cache reads and writes. +func (r *cursorTokenRun) stop(u agent.TokenUsage) { + r.t.Helper() + r.hook("stop", map[string]any{ + "model": "cursor-default", + "loop_count": 1, + "input_tokens": u.InputTokens + u.CacheReadTokens + u.CacheCreationTokens, + "output_tokens": u.OutputTokens, + "cache_read_tokens": u.CacheReadTokens, + "cache_write_tokens": u.CacheCreationTokens, + }) +} + +// TestTokenScope_Cursor_EachTurnCountedOnce: Cursor reports tokens only in its +// Stop hook, so attach must neither recount a pending turn nor drop one, and +// must leave the hook-built session total alone, including mid-turn. +func TestTokenScope_Cursor_EachTurnCountedOnce(t *testing.T) { + t.Parallel() + r := newCursorTokenRun(t) + env := r.env + t1 := agent.TokenUsage{InputTokens: 101, OutputTokens: 11, CacheReadTokens: 1001, CacheCreationTokens: 10001} + t2 := agent.TokenUsage{InputTokens: 203, OutputTokens: 23, CacheReadTokens: 2003, CacheCreationTokens: 20003} + t3 := agent.TokenUsage{InputTokens: 307, OutputTokens: 37, CacheReadTokens: 3007, CacheCreationTokens: 30007} + t4 := agent.TokenUsage{InputTokens: 409, OutputTokens: 49, CacheReadTokens: 4009, CacheCreationTokens: 40009} + var all agent.TokenUsage + for _, u := range []agent.TokenUsage{t1, t2, t3, t4} { + addFields(&all, u) + } + var steps []agentTokenStep + + r.start("a.txt") + r.stop(t1) + env.GitCommitWithHooks("turn 1", "a.txt") + steps = recordCheckpoint(t, env, steps, "commit after turn 1", t1) + + r.start("b.txt", "c.txt") + r.stop(t2) + env.GitCommitWithHooks("turn 2 partial", "b.txt") // c.txt carries forward + steps = recordCheckpoint(t, env, steps, "partial commit after turn 2", t2) + + r.start("e.txt") + r.stop(t3) + assert.Equal(t, tokenFields(&agent.TokenUsage{ + InputTokens: t1.InputTokens + t2.InputTokens + t3.InputTokens, OutputTokens: t1.OutputTokens + t2.OutputTokens + t3.OutputTokens, + CacheReadTokens: t1.CacheReadTokens + t2.CacheReadTokens + t3.CacheReadTokens, CacheCreationTokens: t1.CacheCreationTokens + t2.CacheCreationTokens + t3.CacheCreationTokens, + }), tokenFields(sessionTotal(t, env, r.id)), "turn 3's Stop must record its tokens") + before := tokenFields(sessionTotal(t, env, r.id)) + attachNow(t, env, r.id, "cursor", 3) + steps = recordCheckpoint(t, env, steps, "attach after turn 3 (ended)", t3) + assert.Equal(t, before, tokenFields(sessionTotal(t, env, r.id)), "attach must leave the Cursor session total unchanged") + + r.start("d.txt") + before = tokenFields(sessionTotal(t, env, r.id)) + attachNow(t, env, r.id, "cursor", 4) + steps = recordCheckpoint(t, env, steps, "attach while turn 4 runs", agent.TokenUsage{}) + assert.Equal(t, before, tokenFields(sessionTotal(t, env, r.id)), "running attach must leave the Cursor session total unchanged") + r.stop(t4) + env.GitCommitWithHooks("turn 4", "c.txt", "d.txt", "e.txt") + steps = recordCheckpoint(t, env, steps, "commit after turn 4", t4) + + verifyAgentSteps(t, env, steps, all) + assert.Equal(t, tokenFields(&all), tokenFields(sessionTotal(t, env, r.id)), "session total = every turn once") +} diff --git a/cmd/entire/cli/integration_test/token_turns_test.go b/cmd/entire/cli/integration_test/token_turns_test.go new file mode 100644 index 0000000000..9f800fa95e --- /dev/null +++ b/cmd/entire/cli/integration_test/token_turns_test.go @@ -0,0 +1,245 @@ +//go:build integration + +package integration + +import ( + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/paths" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// tokenTurnRun drives one session through turns, commits and attaches, and +// records what each new checkpoint must hold: exactly the tokens of the turns +// (and turn tails) no earlier checkpoint stored. +type tokenTurnRun struct { + t *testing.T + env *TestEnv + s *Session + + turn int + nextSub int + // pending is what the next checkpoint must count. + pendingOut, pendingSub int + total, totalSub int + want []tokenTurnCheckpoint +} + +type tokenTurnCheckpoint struct { + id string + step string + out, sub int +} + +func newTokenTurnRun(t *testing.T) *tokenTurnRun { + t.Helper() + env := NewFeatureBranchEnv(t) + s := env.NewSession() + // Attach finds the transcript in the agent's project dir; keep it there. + s.TranscriptPath = filepath.Join(env.ClaudeProjectDir, s.ID+".jsonl") + env.ExtraEnv = append(env.ExtraEnv, + "PATH="+filepath.Dir(getTestBinary())+string(os.PathListSeparator)+os.Getenv("PATH")) + return &tokenTurnRun{t: t, env: env, s: s} +} + +// start begins a turn that spends tokens and writes files; sub > 0 also runs +// a subagent that spends sub tokens and writes its own file. +func (r *tokenTurnRun) start(tokens, sub int, files ...string) { + r.t.Helper() + r.turn++ + require.NoError(r.t, r.env.SimulateUserPromptSubmitWithTranscriptPath(r.s.ID, r.s.TranscriptPath)) + changes := make([]FileChange, 0, len(files)) + for _, f := range files { + content := fmt.Sprintf("turn %d\n", r.turn) + r.env.WriteFile(f, content) + changes = append(changes, FileChange{Path: f, Content: content}) + } + appendUsageMessage(r.s, fmt.Sprintf("msg-%d", r.turn), tokens) + r.s.CreateTranscript(fmt.Sprintf("turn %d", r.turn), changes) + r.add(tokens, 0) + if sub > 0 { + r.subagent(sub) + } +} + +func (r *tokenTurnRun) subagent(tokens int) { + r.t.Helper() + r.nextSub++ + a := workflowAgentFixture{ + id: fmt.Sprintf("a%016d", r.nextSub), + file: fmt.Sprintf("sub%d.txt", r.nextSub), + usage: map[string]int{"input_tokens": tokens, "output_tokens": tokens}, + } + require.NoError(r.t, r.env.SimulateSubagentStart(SubagentStartInput{ + SessionID: r.s.ID, TranscriptPath: r.s.TranscriptPath, AgentID: a.id, AgentType: "workflow-subagent", + })) + // Each subagent runs in its own Workflow run, which the parent transcript + // names in the Workflow call's result, as Claude Code does. + runID := fmt.Sprintf("wf_tokturn-%d", r.nextSub) + runDir := filepath.Join(paths.SubagentsDir(filepath.Dir(r.s.TranscriptPath), r.s.ID), + paths.SubagentWorkflowsDirName, runID) + transcript := writeWorkflowAgentTranscript(r.t, runDir, r.nextSub, a) + r.env.WriteFile(a.file, "from subagent\n") + r.s.TranscriptBuilder.messages = append(r.s.TranscriptBuilder.messages, map[string]interface{}{ + "uuid": "workflow-result-" + runID, + "type": "user", + "message": map[string]interface{}{"role": "user", "content": []map[string]interface{}{{ + "type": "tool_result", "tool_use_id": "toolu_" + runID, + "content": "Workflow launched in background.\nRun ID: " + runID + "\n", + }}}, + }) + require.NoError(r.t, r.s.TranscriptBuilder.WriteToFile(r.s.TranscriptPath)) + require.NoError(r.t, r.env.SimulateSubagentStop(SubagentStopInput{ + SessionID: r.s.ID, TranscriptPath: r.s.TranscriptPath, AgentID: a.id, + AgentType: "workflow-subagent", AgentTranscriptPath: transcript, + })) + r.add(0, tokens) +} + +// more spends tokens in the running turn without ending it. +func (r *tokenTurnRun) more(tokens int) { + r.t.Helper() + appendUsageMessage(r.s, fmt.Sprintf("msg-%d-more-%d", r.turn, tokens), tokens) + require.NoError(r.t, r.s.TranscriptBuilder.WriteToFile(r.s.TranscriptPath)) + r.add(tokens, 0) +} + +func (r *tokenTurnRun) stop() { + r.t.Helper() + require.NoError(r.t, r.env.SimulateStop(r.s.ID, r.s.TranscriptPath)) +} + +// checkTotal asserts the session total counts every turn once. A turn with no +// file changes skips SaveStep and adds nothing, so call it only after turns +// that changed files. +func (r *tokenTurnRun) checkTotal() { + r.t.Helper() + state, err := r.env.GetSessionState(r.s.ID) + require.NoError(r.t, err) + require.NotNil(r.t, state.TokenUsage) + assert.Equal(r.t, r.total, state.TokenUsage.OutputTokens, "session total after turn %d", r.turn) +} + +func (r *tokenTurnRun) add(out, sub int) { + r.pendingOut += out + r.pendingSub += sub + r.total += out + r.totalSub += sub +} + +// commit commits files through the hooks, condensing the session. +func (r *tokenTurnRun) commit(files ...string) { + r.t.Helper() + r.env.GitCommitWithHooks(fmt.Sprintf("commit after turn %d", r.turn), files...) + r.record(fmt.Sprintf("commit after turn %d", r.turn)) +} + +// attach commits an unrelated file without hooks and attaches the session to it. +func (r *tokenTurnRun) attach() { + r.t.Helper() + name := fmt.Sprintf("notes-%d.txt", len(r.want)) + r.env.WriteFile(name, "notes") + r.env.GitAdd(name) + r.env.GitCommit("notes " + name) + output := r.env.RunCLI("session", "attach", r.s.ID, "-a", agentClaudeCode, "-f") + require.Contains(r.t, output, "Attached session") + r.record(fmt.Sprintf("attach after turn %d", r.turn)) +} + +func (r *tokenTurnRun) record(step string) { + r.t.Helper() + cp := r.env.TryGetLatestCheckpointID() + require.NotEmpty(r.t, cp, step) + for _, w := range r.want { + require.NotEqual(r.t, w.id, cp, "%s: expected a new checkpoint", step) + } + r.want = append(r.want, tokenTurnCheckpoint{id: cp, step: step, out: r.pendingOut, sub: r.pendingSub}) + r.pendingOut, r.pendingSub = 0, 0 +} + +func (r *tokenTurnRun) verify() { + r.t.Helper() + require.Zero(r.t, r.pendingOut+r.pendingSub, "every turn must end up in a checkpoint") + var sumOut, sumSub int + for _, w := range r.want { + u := readCommittedTokenUsage(r.t, r.env, w.id) + require.NotNil(r.t, u, w.step) + gotSub := 0 + if u.SubagentTokens != nil { + gotSub = u.SubagentTokens.OutputTokens + } + assert.Equal(r.t, w.out, u.OutputTokens, "%s: checkpoint must hold exactly the turns it newly covers", w.step) + assert.Equal(r.t, w.sub, gotSub, "%s: subagent tokens", w.step) + sumOut += u.OutputTokens + sumSub += gotSub + } + assert.Equal(r.t, r.total, sumOut, "checkpoints must sum to the session's tokens") + assert.Equal(r.t, r.totalSub, sumSub, "checkpoints must sum to the session's subagent tokens") +} + +// TestTokenScope_EachTurnCountedOnce: every turn's tokens land in exactly one +// checkpoint across hook commits, partial commits (carry-forward), and +// repeated attaches to successive commits, for ended and running sessions. +func TestTokenScope_EachTurnCountedOnce(t *testing.T) { + t.Parallel() + for _, tc := range []struct { + name string + run func(r *tokenTurnRun) + }{ + { + name: "ended session", + run: func(r *tokenTurnRun) { + r.start(10, 0, "a.txt") + r.stop() + r.commit("a.txt") + r.start(20, 0, "b.txt", "c.txt") + r.stop() + r.commit("b.txt") // partial: c.txt carries forward + r.start(30, 0) + r.stop() + r.attach() + r.start(40, 5) + r.stop() + r.attach() + r.start(50, 0, "d.txt") + r.stop() + r.commit("c.txt", "d.txt", "sub1.txt") + }, + }, + { + name: "attach while the turn runs", + run: func(r *tokenTurnRun) { + r.start(10, 0, "a.txt") + r.stop() + r.commit("a.txt") + r.start(20, 0, "b.txt", "c.txt") + r.stop() + r.commit("b.txt") + r.start(30, 0, "e.txt") + r.attach() // running: stores 30 + r.more(3) + r.stop() + r.checkTotal() // the Stop adds 33 to the total, not on top of attach's 30 + r.start(40, 5) + r.attach() // running: stores the 3-token tail, 40 and the subagent's 5 + r.more(4) + r.stop() + r.checkTotal() + r.start(50, 0, "d.txt") + r.stop() + r.commit("c.txt", "d.txt", "e.txt", "sub1.txt") // 4-token tail + 50 + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + r := newTokenTurnRun(t) + tc.run(r) + r.verify() + }) + } +} diff --git a/cmd/entire/cli/review/manifest.go b/cmd/entire/cli/review/manifest.go index f67aa99fc4..fcf61ebcf4 100644 --- a/cmd/entire/cli/review/manifest.go +++ b/cmd/entire/cli/review/manifest.go @@ -439,7 +439,7 @@ func reviewTokenUsageForSession(ctx context.Context, st *session.State, lookup a slog.String("error", err.Error())) return nil } - return agent.CalculateTokenUsage(ctx, ag, transcript, st.CheckpointTranscriptStart, reviewSubagentsDir(st)) + return agent.CalculateTokenUsage(ctx, ag, transcript, st.TokenStart(), reviewSubagentsDir(st)) } func reviewSubagentsDir(st *session.State) string { diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index f0dc6bbc6d..a08bbf1f5c 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -221,6 +221,16 @@ type State struct { // for checkpoint condensation: "everything since last checkpoint". CheckpointTranscriptStart int `json:"checkpoint_transcript_start,omitempty"` + // TokenTranscriptStart is the transcript position up to which this + // session's tokens are already stored in a checkpoint. Carry-forward and + // adopt move CheckpointTranscriptStart back to 0 so the next checkpoint + // shows the whole conversation; they never reset this offset, so that + // checkpoint still counts only new tokens. It moves back only when a + // condensation or attach finds the transcript shorter (rewritten or + // compacted) and records its new end. Nil in state written before the + // field existed; NormalizeAfterLoad fills it. Read it through TokenStart. + TokenTranscriptStart *int `json:"token_transcript_start,omitempty"` + // CheckpointTranscriptSize is the byte size of the transcript at last condensation. // Used for fast "has new content?" checks in PostCommit: compare the git blob size // against this value without reading the full transcript content. @@ -821,6 +831,12 @@ func (s *State) NormalizeAfterLoad(ctx context.Context) { s.CheckpointTranscriptStart = s.TranscriptLinesAtStart } } + // State written before TokenTranscriptStart existed counted tokens from + // CheckpointTranscriptStart, so start there. An older CLI saving a shared + // state drops the field again; the next load lands here the same way. + if s.TokenTranscriptStart == nil { + s.SetTokenStart(s.CheckpointTranscriptStart) + } // Clear deprecated fields so they aren't re-persisted. // Note: this is a one-way migration. If the state is re-saved, older CLI versions // will see 0 for these fields and fall back to scoping from the transcript start. @@ -891,6 +907,45 @@ func (s *State) ClearCondensationAttempt() { s.CondensationAttempt = nil } +// TokenStart returns the transcript position the next checkpoint counts tokens +// from. See TokenTranscriptStart. +func (s *State) TokenStart() int { + if s.TokenTranscriptStart == nil { + return s.CheckpointTranscriptStart + } + return *s.TokenTranscriptStart +} + +// SetTokenStart records that tokens up to pos are stored in a checkpoint. It +// always assigns a fresh pointer so shallow copies of a state never share it. +func (s *State) SetTokenStart(pos int) { + s.TokenTranscriptStart = &pos +} + +// AdvanceCheckpointWindow moves both transcript offsets to pos after a +// condensation stored the transcript and its tokens up to pos. A condensation +// that could not read the transcript reports pos 0 and stored no transcript +// tokens, so the token offset stays put rather than moving back and letting +// the next checkpoint count everything again. +func (s *State) AdvanceCheckpointWindow(pos int) { + tokenStart := pos + if pos == 0 { + tokenStart = s.TokenStart() + } + s.CheckpointTranscriptStart = pos + s.SetTokenStart(tokenStart) +} + +// AdvanceDisplayWindow moves only CheckpointTranscriptStart to pos, for paths +// that change what a checkpoint shows without storing tokens: the turn-end +// advances past already-condensed content, and carry-forward's restart at 0. +// The token offset is pinned first so an unset one cannot follow +// CheckpointTranscriptStart through TokenStart's fallback. +func (s *State) AdvanceDisplayWindow(pos int) { + s.SetTokenStart(s.TokenStart()) + s.CheckpointTranscriptStart = pos +} + // RebaselineSubagentTokens snapshots the current cumulative subagent total // (TokenUsage.SubagentTokens) into SubagentTokensBaseline so the next checkpoint // window's CheckpointTokenUsage.SubagentTokens is rescoped to "since this @@ -898,8 +953,8 @@ func (s *State) ClearCondensationAttempt() { // // The invariant is: every site that starts a fresh checkpoint window by clearing // CheckpointTokenUsage MUST also re-baseline. Callers: the condensation reset -// helper (resetCheckpointWindow) and cross-repo session adoption, which likewise -// opens a fresh target-local window. Sharing this here keeps the two in step. +// helper (resetCheckpointWindow) and attach. Session adoption deliberately does +// not: it continues the source's token window. func (s *State) RebaselineSubagentTokens() { // Legacy agents without a snapshot retain their existing window baseline. if s.TokenUsage == nil && s.AgentType != agent.AgentTypeCodex { diff --git a/cmd/entire/cli/session/state_test.go b/cmd/entire/cli/session/state_test.go index f58bc14281..ae9d946cfe 100644 --- a/cmd/entire/cli/session/state_test.go +++ b/cmd/entire/cli/session/state_test.go @@ -191,28 +191,45 @@ func TestState_NormalizeAfterLoad_JSONRoundTrip(t *testing.T) { json string wantCTS int // CheckpointTranscriptStart wantStep int // StepCount + wantTok int // TokenStart; older state counted tokens from CheckpointTranscriptStart }{ { name: "migrates old condensed_transcript_lines", json: `{"session_id":"s1","condensed_transcript_lines":42,"checkpoint_count":5}`, wantCTS: 42, wantStep: 5, + wantTok: 42, }, { name: "migrates old transcript_lines_at_start", json: `{"session_id":"s1","transcript_lines_at_start":75}`, wantCTS: 75, + wantTok: 75, }, { name: "preserves new field over old", json: `{"session_id":"s1","condensed_transcript_lines":10,"checkpoint_transcript_start":50}`, wantCTS: 50, + wantTok: 50, }, { name: "handles clean new format", json: `{"session_id":"s1","checkpoint_transcript_start":25,"checkpoint_count":3}`, wantCTS: 25, wantStep: 3, + wantTok: 25, + }, + { + name: "keeps token offset ahead of a carry-forward window", + json: `{"session_id":"s1","checkpoint_transcript_start":0,"token_transcript_start":40}`, + wantCTS: 0, + wantTok: 40, + }, + { + name: "keeps an explicit zero token offset", + json: `{"session_id":"s1","checkpoint_transcript_start":30,"token_transcript_start":0}`, + wantCTS: 30, + wantTok: 0, }, } @@ -224,6 +241,8 @@ func TestState_NormalizeAfterLoad_JSONRoundTrip(t *testing.T) { assert.Equal(t, tt.wantCTS, state.CheckpointTranscriptStart) assert.Equal(t, tt.wantStep, state.StepCount) + assert.Equal(t, tt.wantTok, state.TokenStart()) + require.NotNil(t, state.TokenTranscriptStart, "the token offset is persisted after load") assert.Equal(t, 0, state.CondensedTranscriptLines, "deprecated field should be cleared") assert.Equal(t, 0, state.TranscriptLinesAtStart, "deprecated field should be cleared") }) @@ -1280,6 +1299,44 @@ func TestState_RebaselineSubagentTokensPreservesTriState(t *testing.T) { assert.Nil(t, unknown.SubagentTokensBaseline) } +func TestState_AdvanceDisplayWindowKeepsTokenOffset(t *testing.T) { + t.Parallel() + unset := State{CheckpointTranscriptStart: 8} + unset.AdvanceDisplayWindow(12) + assert.Equal(t, 12, unset.CheckpointTranscriptStart) + assert.Equal(t, 8, unset.TokenStart(), "an unset token offset must not follow the display window") + + set := State{CheckpointTranscriptStart: 0} + set.SetTokenStart(8) + set.AdvanceDisplayWindow(12) + assert.Equal(t, 8, set.TokenStart()) +} + +func TestState_AdvanceCheckpointWindowNeverMovesTokenOffsetBack(t *testing.T) { + t.Parallel() + advanced := State{} + advanced.AdvanceCheckpointWindow(12) + assert.Equal(t, 12, advanced.CheckpointTranscriptStart) + assert.Equal(t, 12, advanced.TokenStart()) + + // A condensation that could not read the transcript reports 0 lines. + unreadable := State{CheckpointTranscriptStart: 12} + unreadable.SetTokenStart(12) + unreadable.AdvanceCheckpointWindow(0) + assert.Equal(t, 0, unreadable.CheckpointTranscriptStart) + assert.Equal(t, 12, unreadable.TokenStart(), "tokens already counted must not be counted again") +} + +func TestState_SetTokenStartDoesNotShareAcrossCopies(t *testing.T) { + t.Parallel() + source := State{} + source.SetTokenStart(10) + copied := source + copied.SetTokenStart(20) + assert.Equal(t, 10, source.TokenStart()) + assert.Equal(t, 20, copied.TokenStart()) +} + func TestState_RebaselineSubagentTokensPreservesLegacyNilUsage(t *testing.T) { t.Parallel() for _, agentType := range []types.AgentType{agent.AgentTypeClaudeCode, agent.AgentTypeFactoryAIDroid} { diff --git a/cmd/entire/cli/session_adopt.go b/cmd/entire/cli/session_adopt.go index d90add0f2b..691cf80617 100644 --- a/cmd/entire/cli/session_adopt.go +++ b/cmd/entire/cli/session_adopt.go @@ -716,15 +716,12 @@ func buildAdoptedSessionState(ctx context.Context, source *session.State, home a adopted.TurnCheckpointIDs = nil adopted.LastCheckpointID = id.EmptyCheckpointID adopted.ClearCondensationAttempt() - adopted.CheckpointTokenUsage = nil - // Re-baseline the subagent cumulative for the fresh target-local window. The - // cloned TokenUsage carries the SOURCE session's full cumulative subagent - // total; without re-baselining here, the first post-adopt checkpoint would - // subtract the source's (stale or nil) baseline and over-report — potentially - // the source session's entire subagent usage. Mirrors resetCheckpointWindow's - // baseline capture so the first adopted checkpoint only counts target-side - // subagent growth, consistent with the PromptWindowBase reset below. - adopted.RebaselineSubagentTokens() + // Token accounting continues from the source's last checkpoint, unlike the + // transcript window above: the retired source never condenses again, so + // tokens it used but never checkpointed (TokenTranscriptStart, pending + // CheckpointTokenUsage, the subagent baseline) are counted once here, and + // tokens it did checkpoint are not counted again. + adopted.SetTokenStart(source.TokenStart()) adopted.FullyCondensed = false adopted.UntrackedFilesAtStart = untrackedFiles diff --git a/cmd/entire/cli/session_adopt_test.go b/cmd/entire/cli/session_adopt_test.go index b77896a86e..a6ee5de35c 100644 --- a/cmd/entire/cli/session_adopt_test.go +++ b/cmd/entire/cli/session_adopt_test.go @@ -1034,8 +1034,13 @@ func TestSessionAdopt_ResetsSourceCheckpointWindow(t *testing.T) { if adopted.CondensationAttempt != nil { t.Fatalf("CondensationAttempt = %#v, want nil", adopted.CondensationAttempt) } - if adopted.CheckpointTokenUsage != nil { - t.Fatalf("CheckpointTokenUsage = %#v, want nil for first target checkpoint", adopted.CheckpointTokenUsage) + // Token accounting continues from the source's last checkpoint: the + // retired source never condenses its pending tokens, so the target does. + if adopted.CheckpointTokenUsage == nil || adopted.CheckpointTokenUsage.InputTokens != 100 || adopted.CheckpointTokenUsage.OutputTokens != 25 { + t.Fatalf("CheckpointTokenUsage = %#v, want the source's pending 100/25", adopted.CheckpointTokenUsage) + } + if adopted.TokenStart() != 2 { + t.Fatalf("TokenStart = %d, want the source's 2 so checkpointed turns are not recounted", adopted.TokenStart()) } commitMsgFile := filepath.Join(targetRepo, "COMMIT_EDITMSG") @@ -1089,23 +1094,23 @@ func TestSessionAdopt_ClearsLegacyTranscriptOffsets(t *testing.T) { } } -// TestSessionAdopt_RebaselinesSubagentTokens pins finding 019f5ebf-dc42: cross-repo -// adoption opens a fresh target-local checkpoint window (StepCount=0, -// CheckpointTokenUsage=nil), but the cloned TokenUsage carries the SOURCE -// session's full cumulative subagent total. If SubagentTokensBaseline is not -// re-baselined to that cumulative, the first post-adopt checkpoint subtracts a -// stale/nil baseline and over-reports the source session's subagent usage. -func TestSessionAdopt_RebaselinesSubagentTokens(t *testing.T) { +// TestSessionAdopt_KeepsSourceSubagentBaseline: adoption retires the source +// session, which never condenses again, so subagent tokens it used after its +// last checkpoint must be counted by the first adopted checkpoint. The adopted +// state keeps the source's baseline (its last checkpoint) rather than +// re-baselining to the current cumulative, which would drop them. +func TestSessionAdopt_KeepsSourceSubagentBaseline(t *testing.T) { for _, tc := range []struct { name string sourceBaseline *agent.TokenUsage + wantIn int + wantOut int }{ - // Source never condensed: baseline is nil, so the first adopted - // checkpoint would report the entire cumulative subagent total. - {name: "never-condensed-source", sourceBaseline: nil}, - // Source condensed at an earlier window: its baseline is stale relative - // to the current cumulative and must not carry into the target window. - {name: "previously-condensed-source", sourceBaseline: &agent.TokenUsage{InputTokens: 200, OutputTokens: 100, APICallCount: 2}}, + // Source never condensed: no subagent tokens are checkpointed yet, so + // the first adopted checkpoint counts all of them. + {name: "never-condensed-source", sourceBaseline: nil, wantIn: 500, wantOut: 250}, + // Source condensed earlier: only growth since that checkpoint counts. + {name: "previously-condensed-source", sourceBaseline: &agent.TokenUsage{InputTokens: 200, OutputTokens: 100, APICallCount: 2}, wantIn: 300, wantOut: 150}, } { t.Run(tc.name, func(t *testing.T) { targetRepo := setupAdoptRepo(t) @@ -1131,19 +1136,9 @@ func TestSessionAdopt_RebaselinesSubagentTokens(t *testing.T) { t.Fatalf("buildAdoptedSessionState failed: %v", err) } - if adopted.SubagentTokensBaseline == nil { - t.Fatal("adopted SubagentTokensBaseline = nil, want re-baselined to the cumulative subagent total") - } - if adopted.SubagentTokensBaseline.InputTokens != 500 || adopted.SubagentTokensBaseline.OutputTokens != 250 { - t.Fatalf("adopted SubagentTokensBaseline = %#v, want cumulative subagent total 500/250", - adopted.SubagentTokensBaseline) - } - - // The first post-adopt checkpoint delta (cumulative - baseline) must be - // zero: adoption should count only target-side subagent growth. delta := types.SubtractTokenUsage(adopted.TokenUsage.SubagentTokens, adopted.SubagentTokensBaseline) - if delta.InputTokens != 0 || delta.OutputTokens != 0 || delta.APICallCount != 0 { - t.Fatalf("first post-adopt subagent delta = %#v, want zero", delta) + if delta.InputTokens != tc.wantIn || delta.OutputTokens != tc.wantOut { + t.Fatalf("first post-adopt subagent delta = %#v, want %d/%d", delta, tc.wantIn, tc.wantOut) } }) } diff --git a/cmd/entire/cli/strategy/attach_token_usage_test.go b/cmd/entire/cli/strategy/attach_token_usage_test.go new file mode 100644 index 0000000000..b7e0f5e3a3 --- /dev/null +++ b/cmd/entire/cli/strategy/attach_token_usage_test.go @@ -0,0 +1,51 @@ +package strategy + +import ( + "context" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/claudecode" + "github.com/stretchr/testify/require" +) + +const attachTestTranscript = `{"type":"assistant","uuid":"u1","message":{"id":"m1","usage":{"input_tokens":10,"output_tokens":10}}}` + "\n" + +// attachTwoTurnTranscript's first line is what an earlier checkpoint counted; +// the second is new since TokenStart 1. +const attachTwoTurnTranscript = `{"type":"assistant","uuid":"u1","message":{"id":"m1","usage":{"input_tokens":100,"output_tokens":100}}}` + "\n" + + `{"type":"assistant","uuid":"u2","message":{"id":"m2","usage":{"input_tokens":10,"output_tokens":10}}}` + "\n" + +func TestAttachTokenUsage_KeepsPendingSubagentTokens(t *testing.T) { + t.Parallel() + state := &SessionState{ + SessionID: "s", + AgentType: agent.AgentTypeClaudeCode, + TokenUsage: &agent.TokenUsage{OutputTokens: 10, SubagentTokens: &agent.TokenUsage{OutputTokens: 50}}, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 10, SubagentTokens: &agent.TokenUsage{OutputTokens: 50}}, + } + usage, pos := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, state, []byte(attachTestTranscript), "") + require.NotNil(t, usage.SubagentTokens, "the attach checkpoint must carry the pending subagent tokens") + require.Equal(t, 50, usage.SubagentTokens.OutputTokens) + + ConsumeAttachTokenWindow(state, pos) + require.Nil(t, state.CheckpointTokenUsage) + require.Equal(t, 1, state.TokenStart()) +} + +func TestAttachTokenUsage_CountsFromTokenStart(t *testing.T) { + t.Parallel() + withState := &SessionState{SessionID: "s", AgentType: agent.AgentTypeClaudeCode} + withState.SetTokenStart(1) + usage, pos := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, withState, []byte(attachTwoTurnTranscript), "") + require.Equal(t, 10, usage.OutputTokens, "tokens before TokenStart are in an earlier checkpoint") + require.Equal(t, 2, pos) + + usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, nil, []byte(attachTwoTurnTranscript), "") + require.Equal(t, 110, usage.OutputTokens, "without state the whole transcript is counted") + + // A state recreated from scratch (cleanup, resume) has counted nothing. + fresh := &SessionState{SessionID: "s", AgentType: agent.AgentTypeClaudeCode} + usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, fresh, []byte(attachTwoTurnTranscript), "") + require.Equal(t, 110, usage.OutputTokens) +} diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 3387357cd1..0f47611b38 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1084,7 +1084,7 @@ func buildSessionMetrics(state *SessionState) *cpkg.SessionMetrics { // 1. Session-state backfill from the freshly-extracted transcript: Copilot // CLI writes session.shutdown after the hooks return, so by condensation // time the authoritative full-session total is recoverable while -// checkpoint metadata stays scoped to CheckpointTranscriptStart. +// checkpoint metadata stays scoped to the token offset (TokenStart). // 2. Accumulated per-checkpoint usage (state.CheckpointTokenUsage, reset at // every condensation). This is what carries out-of-band token counts // (e.g. Antigravity, whose transcript has no token data — SaveStep @@ -1101,7 +1101,7 @@ func resolveCondensedTokenUsage(ctx context.Context, ag agent.Agent, state *Sess // Backfill session state token usage from the freshly-extracted transcript. // Copilot CLI writes session.shutdown after the hooks return, so by condensation // time we can recover the authoritative full-session total from the transcript - // while keeping checkpoint metadata scoped to CheckpointTranscriptStart. The + // while keeping checkpoint metadata scoped to the token offset (TokenStart). The // recompute drops SubagentTokens (subagentsDir=""); the helper preserves the // cumulative subagent total across the backfill so resetCheckpointWindow's // baseline does not regress to nil (finding 019f5ebf-a57e). @@ -1117,6 +1117,45 @@ func resolveCondensedTokenUsage(ctx context.Context, ag agent.Agent, state *Sess } } +// AttachTokenUsage returns the tokens `entire session attach` stores for a +// session's transcript and the transcript end position in the agent's offset +// units. With existing state it counts only what no checkpoint has counted yet +// (from TokenStart), falling back to the pending hook-reported usage the way +// condensation does; without state it counts the whole transcript. The caller +// records pos with ConsumeAttachTokenWindow once the checkpoint is written. +// While a turn runs, its subagents' tokens aren't in the pending usage yet +// (Stop records them), so they are read live as a mid-turn condensation does; +// that also moves state.TokenUsage's cumulative subagent total, which the +// caller must keep so the re-baseline covers what this checkpoint stored. +func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte, transcriptPath string) (*agent.TokenUsage, int) { + var usage *agent.TokenUsage + if state == nil { + usage = agent.CalculateTokenUsage(ctx, ag, transcript, 0, "") + } else { + usage = calculateLiveTranscriptTokenUsage(ctx, ag, transcript, state, transcriptPath) + } + if state != nil { + if !hasTokenUsageData(usage) && hasTokenUsageData(state.CheckpointTokenUsage) { + usage = accumulateTokenUsage(nil, state.CheckpointTokenUsage) + } else { + // The transcript recompute reads no subagent transcripts; take the + // pending window's subagent total, as condensation does. + usage = fillMissingSubagentTokensFrom(usage, state.CheckpointTokenUsage) + } + } + return usage, countTranscriptItems(ag.Type(), string(transcript)) +} + +// ConsumeAttachTokenWindow marks the tokens an attach checkpoint stored as +// counted: the token offset moves to pos and the pending window is reset, so +// the session's next checkpoint does not count them again. The displayed +// transcript window (CheckpointTranscriptStart) is left alone. +func ConsumeAttachTokenWindow(state *SessionState, pos int) { + state.CheckpointTokenUsage = nil + state.RebaselineSubagentTokens() + state.SetTokenStart(pos) +} + func hasTokenUsageData(usage *agent.TokenUsage) bool { if usage == nil { return false @@ -1348,7 +1387,8 @@ func (s *ManualCommitStrategy) extractSessionDataFrom(ctx context.Context, ag ag // Resolve files touched: prefers hook-populated state, falls back to transcript extraction data.FilesTouched = s.resolveFilesTouched(ctx, state) - // Calculate token usage from the checkpoint-scoped transcript portion. + // Calculate token usage from the token offset: everything no earlier + // checkpoint counted, even when carry-forward widened the displayed window. // Skill events annotate the stored raw transcript, which is full-session, so // extract them from offset 0; consumers can filter by checkpoint_transcript_start // if they only render the checkpoint-scoped slice. @@ -1436,7 +1476,7 @@ func resolvePendingTranscriptOffset(ctx context.Context, ag agent.Agent, state * slog.Int("old_offset", state.CheckpointTranscriptStart), slog.Int("new_offset", pos), ) - state.CheckpointTranscriptStart = pos + state.AdvanceDisplayWindow(pos) state.TranscriptOffsetPending = false } @@ -1510,7 +1550,7 @@ func calculateLiveTranscriptTokenUsage( transcriptPath string, ) *agent.TokenUsage { subagentsDir := liveSubagentsDir(ag, state, transcriptPath) - usage := agent.CalculateTokenUsage(ctx, ag, transcript, state.CheckpointTranscriptStart, subagentsDir) + usage := agent.CalculateTokenUsage(ctx, ag, transcript, state.TokenStart(), subagentsDir) if usage == nil || usage.SubagentTokens == nil { return usage } @@ -2059,7 +2099,7 @@ func (s *ManualCommitStrategy) CondenseSessionByID(ctx context.Context, sessionI ) resetCheckpointWindow(state) - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.CheckpointTranscriptSize = result.TranscriptSizeBaseline state.Phase = session.PhaseIdle state.LastCheckpointID = result.CheckpointID @@ -2225,7 +2265,7 @@ func (s *ManualCommitStrategy) CondenseAndMarkFullyCondensed(ctx context.Context } resetCheckpointWindow(state) - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.LastCheckpointID = result.CheckpointID state.FullyCondensed = true // Phase stays ENDED — do NOT set to IDLE diff --git a/cmd/entire/cli/strategy/manual_commit_hooks.go b/cmd/entire/cli/strategy/manual_commit_hooks.go index 9f4a563bd5..159c78e825 100644 --- a/cmd/entire/cli/strategy/manual_commit_hooks.go +++ b/cmd/entire/cli/strategy/manual_commit_hooks.go @@ -1799,7 +1799,7 @@ func (s *ManualCommitStrategy) condenseAndUpdateState( // work was not. Keep its existing SaveStep count, which keeps the // session pending (State.HasPendingWork) until a home commit consumes it. state.StepCount = pendingStepCount - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.CheckpointTranscriptSize = result.TranscriptSizeBaseline logging.Info(logCtx, "session guest-condensed from a sibling worktree; home worktree state untouched", slog.String("strategy", "manual-commit"), @@ -1814,7 +1814,7 @@ func (s *ManualCommitStrategy) condenseAndUpdateState( newHead := head.Hash().String() state.BaseCommit = newHead resetCheckpointWindow(state) - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.CheckpointTranscriptSize = result.TranscriptSizeBaseline state.FilesTouched = nil state.TouchedFileHashes = nil @@ -3154,7 +3154,7 @@ func advanceCheckpointTranscriptStartToTurnEnd(ctx context.Context, state *Sessi slog.Int("old_offset", state.CheckpointTranscriptStart), slog.Int("new_offset", pos), ) - state.CheckpointTranscriptStart = pos + state.AdvanceDisplayWindow(pos) advanced = true } } @@ -3521,7 +3521,9 @@ func filesChangedInCommitFallback(ctx context.Context, headTree, parentTree *obj // sessionHasNewContent through the stored-transcript check. func carryForwardRemainingFiles(logCtx context.Context, state *SessionState) { state.StepCount = 1 - state.CheckpointTranscriptStart = 0 + // Only the displayed window restarts; the token offset stays where the + // last condensation left it so this checkpoint counts only new tokens. + state.AdvanceDisplayWindow(0) state.CheckpointTranscriptSize = 0 // Carry-forward deliberately restarts the offset at 0; a pending turn-end // advance from before the carry-forward must not re-apply on top of it. diff --git a/cmd/entire/cli/strategy/manual_commit_types.go b/cmd/entire/cli/strategy/manual_commit_types.go index b535b8406a..0045b60fab 100644 --- a/cmd/entire/cli/strategy/manual_commit_types.go +++ b/cmd/entire/cli/strategy/manual_commit_types.go @@ -81,7 +81,7 @@ type ExtractedSessionData struct { FullTranscriptLines int // Total line count in full transcript Prompts []string // User prompts from the current checkpoint portion FilesTouched []string - TokenUsage *agent.TokenUsage // Token usage calculated from transcript (since CheckpointTranscriptStart) + TokenUsage *agent.TokenUsage // Token usage calculated from transcript (since the token offset, TokenStart) // SkillEvents are this condensation's extracted events. Transient — the // durable ledger, and the per-hook cost of carrying it, is // session.SessionState.SkillEvents; see its size note. diff --git a/docs/architecture/copilot-token-validation.md b/docs/architecture/copilot-token-validation.md index f6f479dec3..2c80fcd121 100644 --- a/docs/architecture/copilot-token-validation.md +++ b/docs/architecture/copilot-token-validation.md @@ -17,7 +17,7 @@ Copilot CLI exposes two different kinds of token data in `events.jsonl`: That means Entire must treat the same transcript in two different ways: -- Checkpoint metadata must stay scoped to `CheckpointTranscriptStart` +- Checkpoint metadata must stay scoped to the token offset (`SessionState.TokenStart()`), which starts where the previous checkpoint's tokens ended - Session state for `entire status` should use the full-session aggregate once `session.shutdown` exists If this logic regresses, earlier checkpoints can suddenly show the same token count as the whole session. diff --git a/docs/development/checkpoint-implementation.md b/docs/development/checkpoint-implementation.md index 3088bcab24..2e44a24947 100644 --- a/docs/development/checkpoint-implementation.md +++ b/docs/development/checkpoint-implementation.md @@ -55,7 +55,7 @@ The manual-commit strategy (`manual_commit*.go`) does not modify the active bran - Uses the `post-rewrite` Git hook to keep local session linkage aligned after amend/rebase rewrites - Builds git trees in-memory using go-git plumbing APIs - **Location-independent transcript resolution** - transcript paths are always computed dynamically from the current repo location (via `agent.GetSessionDir` + `agent.ResolveSessionFile`), never stored in checkpoint metadata. This ensures log restore (`RestoreLogsOnly`) works after repo relocation or across machines. -- **Token usage scoping** - `SessionState.TokenUsage` is the session-wide total used by `entire status`; `SessionState.CheckpointTokenUsage` is the pending checkpoint delta since the last condensation. Checkpoint metadata must stay scoped to `CheckpointTranscriptStart` or the pending checkpoint delta. Cursor tokens come only from stop-hook payloads, while Copilot CLI can also backfill full-session totals from `session.shutdown`. Condensation's full-session backfill recompute runs with `subagentsDir=""` and so drops `SubagentTokens`; `fillMissingSubagentTokensFrom` refills it from the already-rescoped `state.CheckpointTokenUsage`. The checkpoint-scoped recompute (`calculateLiveTranscriptTokenUsage`) reads the still-available subagent transcripts only when no checkpoint-scoped subagent total already exists, the agent supports that extraction, and a real subagent directory exists. It subtracts `SubagentTokensBaseline` for checkpoint metadata and keeps the cumulative snapshot on `state.TokenUsage` so the reset advances the next baseline; an empty delta stays nil. The scan is substantially more expensive for subagent-heavy sessions, so every gate is load-bearing. The store sums those scoped values across a checkpoint's sessions via `types.AddTokenUsage` (the single token-summing primitive — do not hand-roll another; a field-by-field copy is how the nested total came to be dropped in the first place). +- **Token usage scoping** - `SessionState.TokenUsage` is the session-wide total used by `entire status`; `SessionState.CheckpointTokenUsage` is the pending checkpoint delta since the last condensation. Checkpoint token metadata must stay scoped to `SessionState.TokenStart()` (`TokenTranscriptStart`) or the pending checkpoint delta. `TokenTranscriptStart` is set only where tokens are stored or carried over, and moves back only when a condensation or attach records the end of a transcript that got shorter (rewritten or compacted): condensation (`AdvanceCheckpointWindow`, which leaves it alone when the transcript could not be read), attach (`ConsumeAttachTokenWindow`) and adopt (`SetTokenStart`); unlike `CheckpointTranscriptStart`, carry-forward and adopt never move it back and the turn-end advance after a mid-turn commit leaves it behind, so a checkpoint whose displayed window restarts at 0 still counts each token once and a mid-turn commit's post-commit tail lands in the next checkpoint. `session attach` counts from it too and records the transcript end afterwards (`strategy.AttachTokenUsage` / `ConsumeAttachTokenWindow`); adopt carries the source's offset, pending delta and subagent baseline because the retired source never condenses again. Cursor tokens come only from stop-hook payloads, while Copilot CLI can also backfill full-session totals from `session.shutdown`. Condensation's full-session backfill recompute runs with `subagentsDir=""` and so drops `SubagentTokens`; `fillMissingSubagentTokensFrom` refills it from the already-rescoped `state.CheckpointTokenUsage`. The checkpoint-scoped recompute (`calculateLiveTranscriptTokenUsage`, counting from `TokenStart()`) reads the still-available subagent transcripts only when no checkpoint-scoped subagent total already exists, the agent supports that extraction, and a real subagent directory exists. It subtracts `SubagentTokensBaseline` for checkpoint metadata and keeps the cumulative snapshot on `state.TokenUsage` so the reset advances the next baseline; an empty delta stays nil. The scan is substantially more expensive for subagent-heavy sessions, so every gate is load-bearing. The store sums those scoped values across a checkpoint's sessions via `types.AddTokenUsage` (the single token-summing primitive — do not hand-roll another; a field-by-field copy is how the nested total came to be dropped in the first place). - Tracks session state in `.git/entire-sessions/` (shared across worktrees) - **Commit-to-session linking is identity-first** (`strategy/session_identity.go`): identity comes from `SessionState.Owner`, the `proclive.Identity` that `captureSessionOwner` already records on every turn start (first non-transient ancestor — proclive skips shells, `entire` itself, and the Go toolchain, so a human commit typed in the same terminal never matches). Commit hooks snapshot their own ancestry once (`proclive.CurrentAncestry`) and match every candidate against it in memory (`Ancestry.Depth`) — one hostname/boot-id/proc walk per commit, not one per session state — linking the commit to the session whose agent process is an ancestor — in any worktree (nearest ancestor wins, so a nested agent beats the outer agent that spawned it, and only a tie at equal depth falls to the latest interaction; host/boot/start-time guards defeat PID reuse and cross-machine matches; Windows cannot introspect and falls back to worktree matching). The identity match is UNIONED with the worktree-matched set, never a replacement: a commit condenses every session with pending content in its worktree. Any session matched outside its home worktree is guest-linked — whether identity-matched or selected by the pre-existing single-worktree fallback — and is condensed and linked without mutating worktree-coupled state (`BaseCommit`) from the foreign worktree (`isSessionHomeWorktree`). Worktree matching is always computed (it is the sole mechanism for commits with no agent ancestry): imported sessions never link, and multi-worktree ambiguity is filtered to recently-interacting sessions (15 min) before declining. This deliberately turns some former ambiguity declines into a best-candidate link; `recentSessionWindow` is a correctness tradeoff because a session in a long-running build or tool call can age out and leave the other recent worktree to win. The stderr hint naming `entire session adopt` fires only from the commit-linking path, and only when identity matching could not rescue the commit either. **Squashes inherit their trailers**: when SQUASH_MSG exists in the per-worktree git dir, prepare-commit-msg copies its `Entire-Checkpoint` trailers into the message only when a staged path is one a commit Git recorded in that message changed (`inheritSquashedCheckpointTrailers`); a squash committed with `-m` reports `message`, not `squash`, and an abandoned squash's stale message is ignored. Inherited trailers are links: matching still runs and a session with pending work is stamped after them, and PostCommit condenses only into the trailer without a checkpoint (`pickCondensationTarget`, the last such one, since the stamp is appended) among those prepare did not record as inherited (`recordInheritedTrailers` / `takeInheritedTrailers`, a marker in the per-worktree git dir tied to the commit's parent), rechecks that target's existence before writing, and refuses to write into a checkpoint the session did not stamp or is not amending (`stampedByAnotherCommit`). **Redone commits inherit too** (`inheritReplacedCommitsTrailers`, `replacedCommits`): the dropped work is read from HEAD's reflog (per-worktree `logs/HEAD`, tail-bounded), skipping resets that did not move HEAD; dropped commits are those only the old tip reaches (`git rev-list ` excluding HEAD, remotes and every branch not pointing at the tip); a dropped commit is inherited when every staged file it changed matches the old tip byte for byte; commits since the reset keep the redo open only if each redid dropped work; an amend fold (`reset --soft` + `commit --amend`) inherits too, recorded against HEAD's parent. Under `go test`, `session.NewStateStore` and `NewStateStoreForWorktree` refuse to open outside the temp root so non-isolated tests fail loudly instead of leaking fixture sessions into a real repo. - **Reclaiming sessions whose agent vanished** - not every agent fires a session-end hook, and any agent can be killed before its hook runs, so a session can be left un-finalized forever. `SessionState.Owner` — the same fingerprint commit linking matches above — is captured at every turn start by `captureSessionOwner`, and `State.OwnerExited()` reports it gone via `proclive.Check`. `finalizeExitedSessions` sweeps those inside `entire doctor` and `__sweep_sessions`, ending them exactly as a clean stop would. **Not `entire status`**: status is a read-only observer — it reads through `StateStore.ListReadOnly` and reports a dead-owner session as `exited` without finalizing it, because asking what is happening must not change what is happening. Status therefore shows such a session until doctor or the sweeper reaches it. **`OwnerExited` deliberately covers IDLE as well as ACTIVE** — an agent that finishes its last turn and then quits leaves IDLE, so gating on ACTIVE alone missed the common case; only already-finalized sessions are excluded, per the shared `State.IsEnded()` predicate. Liveness is Unknown on Windows and for cross-host state, where behaviour degrades to the `StuckActiveThreshold` timeout. Because the sweep runs inside interactive commands, its eager condensing is capped by `sweepCondenseBudget` across the whole sweep: every candidate is always marked ENDED (a single atomic rename — that is what stops `entire status` reporting it as live), while condensing runs only while the budget lasts, so a multi-day backlog drains over successive invocations instead of stalling one. Skipping a condense is the existing fail-open path — PostCommit retries, and `doctor` reports the session as "ended with uncondensed checkpoint data".