From 4ea92d80801db8e7530634dc0ca5bae6190c2306 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Mon, 5 Oct 2026 20:33:03 -0700 Subject: [PATCH 01/12] fix(strategy): count each checkpoint's tokens once CheckpointTranscriptStart chose both which part of the transcript a checkpoint shows and where its token count starts. Carry-forward after a partial commit and session adopt reset it to 0 so the next checkpoint shows the whole conversation, which also made that checkpoint count every token from the start of the session again. Session attach always counted from 0 and never recorded how far it got, and the turn-end advance after a mid-turn commit skipped that turn's remaining tokens entirely. Add a separate offset, TokenTranscriptStart, that advances only when a condensation stores tokens. Token counting in condensation, the live transcript path, review sessions and attach reads it through TokenStart(). Attach records the transcript end afterwards and resets the pending window. Adopt keeps the source's token offset, pending total and subagent baseline, because the retired source never condenses again. State written by an older CLI starts from CheckpointTranscriptStart, as before. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M47MB89PBDYQTEFR7ACGRJCS --- cmd/entire/cli/attach.go | 23 +- .../cli/integration_test/token_scope_test.go | 204 ++++++++++++++++++ cmd/entire/cli/review/manifest.go | 2 +- cmd/entire/cli/session/state.go | 40 +++- cmd/entire/cli/session/state_test.go | 29 +++ cmd/entire/cli/session_adopt.go | 15 +- cmd/entire/cli/session_adopt_test.go | 49 ++--- .../cli/strategy/condensation_prompts_test.go | 2 +- .../strategy/manual_commit_condensation.go | 47 +++- .../cli/strategy/manual_commit_hooks.go | 4 +- docs/development/checkpoint-implementation.md | 2 +- 11 files changed, 361 insertions(+), 56 deletions(-) create mode 100644 cmd/entire/cli/integration_test/token_scope_test.go diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index a4b2a5849b..6ac91dd2fb 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -329,7 +329,13 @@ 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, 0, "") + // The checkpoint stores only tokens no earlier checkpoint of this session + // counted; session state keeps the whole-transcript total. + tokenUsage, tokenPos := strategy.AttachTokenUsage(logCtx, ag, existingState, transcriptData) + sessionUsage := tokenUsage + if existingState != nil { + sessionUsage = agent.CalculateTokenUsage(logCtx, ag, transcriptData, 0, "") + } // attach writes checkpoints and historically never configured // redaction; a scanner-config failure must fail the attach. @@ -373,7 +379,7 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa } // Create or update session state. - if err := saveAttachSessionState(logCtx, repo, existingState, sessionID, ag.Type(), transcriptPath, checkpointID, meta, tokenUsage, opts, reviewSkills); err != nil { + if err := saveAttachSessionState(logCtx, repo, existingState, sessionID, ag.Type(), transcriptPath, checkpointID, meta, sessionUsage, tokenPos, opts, reviewSkills); err != nil { logging.Warn(logCtx, "failed to save session state", "error", err) } @@ -698,7 +704,7 @@ func resolveCheckpointID(ctx context.Context, headCommit *object.Commit) (id.Che // saveAttachSessionState creates or updates the session state file for the attached session. // 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. -func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingState *session.State, sessionID string, agentType types.AgentType, transcriptPath string, checkpointID id.CheckpointID, meta transcriptMetadata, tokenUsage *agent.TokenUsage, opts attachOptions, reviewSkills []string) error { +func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingState *session.State, sessionID string, agentType types.AgentType, transcriptPath string, checkpointID id.CheckpointID, meta transcriptMetadata, sessionUsage *agent.TokenUsage, tokenPos int, opts attachOptions, reviewSkills []string) error { stateStore, err := session.NewStateStore(ctx) if err != nil { return fmt.Errorf("failed to open session store: %w", err) @@ -743,9 +749,16 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS if meta.FirstPrompt != "" { state.LastPrompt = meta.FirstPrompt } - if tokenUsage != nil { - state.TokenUsage = tokenUsage + if sessionUsage != nil { + // Attach reads no subagent transcripts; keep the cumulative subagent + // total hooks recorded so the re-baseline below doesn't drop it. + if 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 state.ReviewSkills = reviewSkills 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..5dd20c2dd9 --- /dev/null +++ b/cmd/entire/cli/integration_test/token_scope_test.go @@ -0,0 +1,204 @@ +//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.GitCommitWithShadowHooks("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.GitCommitWithShadowHooks("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.GitCommitWithShadowHooks("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.GitCommitWithShadowHooks("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.GitCommitWithShadowHooks("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.GitCommitWithShadowHooksAsAgent("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.GitCommitWithShadowHooks("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)") +} 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 1ccb02c977..a2009286e3 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -227,6 +227,14 @@ 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; this offset never moves back, so that + // checkpoint still counts only new tokens. 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. @@ -822,6 +830,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. @@ -906,6 +920,28 @@ 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. +func (s *State) AdvanceCheckpointWindow(pos int) { + s.CheckpointTranscriptStart = pos + s.SetTokenStart(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 @@ -913,8 +949,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 14cde2a9d1..4efd88db07 100644 --- a/cmd/entire/cli/session/state_test.go +++ b/cmd/entire/cli/session/state_test.go @@ -227,28 +227,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, }, } @@ -260,6 +277,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") }) @@ -1316,6 +1335,16 @@ func TestState_RebaselineSubagentTokensPreservesTriState(t *testing.T) { assert.Nil(t, unknown.SubagentTokensBaseline) } +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 fd2a89fcd9..bb134b0a8a 100644 --- a/cmd/entire/cli/session_adopt.go +++ b/cmd/entire/cli/session_adopt.go @@ -592,15 +592,12 @@ func buildAdoptedSessionState(ctx context.Context, source *session.State, source adopted.LastCheckpointID = id.EmptyCheckpointID adopted.ClearCondensationAttempt() adopted.LastCheckpointCommitHash = "" - 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 6361c84501..52ad2a110c 100644 --- a/cmd/entire/cli/session_adopt_test.go +++ b/cmd/entire/cli/session_adopt_test.go @@ -1053,8 +1053,13 @@ func TestSessionAdopt_ResetsSourceCheckpointWindow(t *testing.T) { if adopted.LastCheckpointCommitHash != "" { t.Fatalf("LastCheckpointCommitHash = %q, want empty", adopted.LastCheckpointCommitHash) } - 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") @@ -1108,23 +1113,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) @@ -1150,19 +1155,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/condensation_prompts_test.go b/cmd/entire/cli/strategy/condensation_prompts_test.go index 3ec5cbe3dd..d8e4c1a97c 100644 --- a/cmd/entire/cli/strategy/condensation_prompts_test.go +++ b/cmd/entire/cli/strategy/condensation_prompts_test.go @@ -83,7 +83,7 @@ func TestExtractSessionData_ResolvesPromptsWithEmptyTranscript(t *testing.T) { s := &ManualCommitStrategy{} data, err := s.extractSessionData(context.Background(), repo, hash, sessionID, nil, - agent.AgentTypeAntigravity, "", 0, false) + agent.AgentTypeAntigravity, "", 0, 0, false) require.NoError(t, err) require.Empty(t, data.Transcript, "sanity: this is the empty-transcript case") require.Equal(t, []string{"the prompt that must survive"}, data.Prompts, diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index a7cd5c1e65..7a65184e04 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1004,7 +1004,7 @@ func (s *ManualCommitStrategy) extractOrCreateSessionData(ctx context.Context, r case hasShadowBranch: // Shadow branch exists (from SaveStep commits) — extract transcript and // metadata from the branch tree, preferring the live transcript if fresher. - data, err := s.extractSessionData(ctx, repo, shadowHash, state.SessionID, state.FilesTouched, state.AgentType, state.TranscriptPath, state.CheckpointTranscriptStart, state.Phase.IsActive()) + data, err := s.extractSessionData(ctx, repo, shadowHash, state.SessionID, state.FilesTouched, state.AgentType, state.TranscriptPath, state.CheckpointTranscriptStart, state.TokenStart(), state.Phase.IsActive()) if err != nil { return nil, fmt.Errorf("failed to extract session data: %w", err) } @@ -1243,6 +1243,34 @@ 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. +func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte) (*agent.TokenUsage, int) { + start := 0 + if state != nil { + start = state.TokenStart() + } + usage := agent.CalculateTokenUsage(ctx, ag, transcript, start, "") + if state != nil && !hasTokenUsageData(usage) && hasTokenUsageData(state.CheckpointTokenUsage) { + usage = accumulateTokenUsage(nil, 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 @@ -1525,7 +1553,8 @@ func committedFilesExcludingMetadata(committedFiles map[string]struct{}) []strin // This handles the case where SaveStep was skipped (no code changes) but the transcript // continued growing — the shadow branch copy would be stale. // checkpointTranscriptStart is the line offset (JSONL agents) or message index (OpenCode) where the current checkpoint began. -func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git.Repository, shadowRef plumbing.Hash, sessionID string, filesTouched []string, agentType types.AgentType, liveTranscriptPath string, checkpointTranscriptStart int, isActive bool) (*ExtractedSessionData, error) { +// tokenStart is the position tokens are counted from (SessionState.TokenStart), in the same units. +func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git.Repository, shadowRef plumbing.Hash, sessionID string, filesTouched []string, agentType types.AgentType, liveTranscriptPath string, checkpointTranscriptStart, tokenStart int, isActive bool) (*ExtractedSessionData, error) { ag, _ := agent.GetByAgentType(agentType) //nolint:errcheck // ag may be nil for unknown agent types; callers use type assertions so nil is safe commit, err := repo.CommitObject(shadowRef) if err != nil { @@ -1609,7 +1638,8 @@ func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git // Use tracked files from session state (not all files in tree) data.FilesTouched = filesTouched - // 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. @@ -1621,7 +1651,7 @@ func (s *ManualCommitStrategy) extractSessionData(ctx context.Context, repo *git // cumulative snapshot needing the same rescoping SaveStep already did. // CondenseSession fills the already-rescoped window total in instead; // see fillMissingSubagentTokensFrom. - data.TokenUsage = agent.CalculateTokenUsage(ctx, ag, data.Transcript, checkpointTranscriptStart, "") + data.TokenUsage = agent.CalculateTokenUsage(ctx, ag, data.Transcript, tokenStart, "") data.SkillEvents = agent.ExtractSkillEvents(ctx, ag, data.Transcript, 0) } @@ -1684,7 +1714,8 @@ func (s *ManualCommitStrategy) extractSessionDataFromLiveTranscript(ctx context. // 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. @@ -1846,7 +1877,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 } @@ -2166,7 +2197,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 @@ -2341,7 +2372,7 @@ func (s *ManualCommitStrategy) CondenseAndMarkFullyCondensed(ctx context.Context } resetCheckpointWindow(state) - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.LastCheckpointID = result.CheckpointID state.LastCheckpointCommitHash = state.BaseCommit state.RealignAttributionBase(state.BaseCommit) diff --git a/cmd/entire/cli/strategy/manual_commit_hooks.go b/cmd/entire/cli/strategy/manual_commit_hooks.go index 89d5aed005..e8fa9d0a99 100644 --- a/cmd/entire/cli/strategy/manual_commit_hooks.go +++ b/cmd/entire/cli/strategy/manual_commit_hooks.go @@ -1882,7 +1882,7 @@ func (s *ManualCommitStrategy) condenseAndUpdateState( // content was not. Keep its existing SaveStep count as the ownership // pin used by cleanup and doctor 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; shadow state untouched", slog.String("strategy", "manual-commit"), @@ -1901,7 +1901,7 @@ func (s *ManualCommitStrategy) condenseAndUpdateState( state.BaseCommit = newHead state.RealignAttributionBase(newHead) resetCheckpointWindow(state) - state.CheckpointTranscriptStart = result.TotalTranscriptLines + state.AdvanceCheckpointWindow(result.TotalTranscriptLines) state.CheckpointTranscriptSize = result.TranscriptSizeBaseline // Clear attribution tracking — condensation already used these values diff --git a/docs/development/checkpoint-implementation.md b/docs/development/checkpoint-implementation.md index 6ae6e4f869..306f5d0857 100644 --- a/docs/development/checkpoint-implementation.md +++ b/docs/development/checkpoint-implementation.md @@ -53,7 +53,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`. Shadow-branch condensation's transcript recompute runs with `subagentsDir=""` and so drops `SubagentTokens`; `fillMissingSubagentTokensFrom` refills it from the already-rescoped `state.CheckpointTokenUsage`. A live mid-turn condensation when no shadow branch resolves instead 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` advances only when a condensation stores tokens (`AdvanceCheckpointWindow`); 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`. Shadow-branch condensation's transcript recompute runs with `subagentsDir=""` and so drops `SubagentTokens`; `fillMissingSubagentTokensFrom` refills it from the already-rescoped `state.CheckpointTokenUsage`. A live mid-turn condensation when no shadow branch resolves instead 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`, shadow-branch realignment) 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". From 183d68731a7451f8952dbbddecfed38a2f98eaf4 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Mon, 5 Oct 2026 20:43:12 -0700 Subject: [PATCH 02/12] fix(session): pin the token offset when only the display window advances The turn-end advances after a mid-turn commit move CheckpointTranscriptStart without storing tokens. An unset TokenTranscriptStart falls back to CheckpointTranscriptStart, so it would follow the display window forward and skip that turn's tail tokens. AdvanceDisplayWindow pins the token offset before moving the window; both turn-end sites use it. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M47MXV4A361ASRFC7NM860TX --- cmd/entire/cli/session/state.go | 9 +++++++++ cmd/entire/cli/session/state_test.go | 13 +++++++++++++ .../cli/strategy/manual_commit_condensation.go | 2 +- cmd/entire/cli/strategy/manual_commit_hooks.go | 2 +- 4 files changed, 24 insertions(+), 2 deletions(-) diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index a2009286e3..3a389f7ec2 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -942,6 +942,15 @@ func (s *State) AdvanceCheckpointWindow(pos int) { s.SetTokenStart(pos) } +// AdvanceDisplayWindow moves only CheckpointTranscriptStart to pos, for the +// turn-end advances that skip already-condensed content without storing its +// tokens. The token offset is pinned first so an unset one cannot follow +// CheckpointTranscriptStart forward 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 diff --git a/cmd/entire/cli/session/state_test.go b/cmd/entire/cli/session/state_test.go index 4efd88db07..422ec46136 100644 --- a/cmd/entire/cli/session/state_test.go +++ b/cmd/entire/cli/session/state_test.go @@ -1335,6 +1335,19 @@ 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_SetTokenStartDoesNotShareAcrossCopies(t *testing.T) { t.Parallel() source := State{} diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 7a65184e04..d6072ceb77 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1803,7 +1803,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 } diff --git a/cmd/entire/cli/strategy/manual_commit_hooks.go b/cmd/entire/cli/strategy/manual_commit_hooks.go index e8fa9d0a99..52526a0ec7 100644 --- a/cmd/entire/cli/strategy/manual_commit_hooks.go +++ b/cmd/entire/cli/strategy/manual_commit_hooks.go @@ -3486,7 +3486,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 } } From ee6b820454c86d067d57dc587240d486971a5e1b Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Wed, 7 Oct 2026 11:22:01 -0700 Subject: [PATCH 03/12] fix(attach): keep tokens already stored in the replaced entry and pending subagent tokens Attach writing into HEAD's checkpoint replaces this session's entry. It now reads that entry first and keeps its stored tokens, adding only the tokens since TokenStart; without state the whole-transcript count already covers them. Attach also fills the pending window's subagent tokens before consuming the window, as condensation does, so they land in the attach checkpoint instead of nowhere. The existing-entry read moves into attachReplacedEntryTokens; review attaches still refuse an existing entry and fail on a read error. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4BSKPVWH6285NNJ70J7DD7M --- cmd/entire/cli/attach.go | 82 ++++++++++++------- .../cli/integration_test/token_scope_test.go | 39 +++++++++ .../cli/strategy/attach_token_usage_test.go | 41 ++++++++++ .../strategy/manual_commit_condensation.go | 19 ++++- 4 files changed, 148 insertions(+), 33 deletions(-) create mode 100644 cmd/entire/cli/strategy/attach_token_usage_test.go diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index 6ac91dd2fb..14b03f20eb 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -302,25 +302,11 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa return err } - // Defense-in-depth guard: the earlier existingState.LastCheckpointID - // check only fires when the session's state file records its - // checkpoint. A session already stored in the HEAD checkpoint but - // whose state is missing/stale (state file deleted, never written, - // condensed without LastCheckpointID update, or pulled from a remote - // that wasn't reflected locally) would bypass that guard. - // findSessionIndex matches by SessionID — without this check, a - // review-attach on such a session silently overwrites the existing - // session's metadata in the checkpoint. - if opts.Review && isExistingCheckpoint { - exists, readErr := checkpointHasSessionMetadata(ctx, repo, refs, checkpointID, sessionID) - if readErr != nil { - return fmt.Errorf("failed to check checkpoint %s for session %s: %w", checkpointID.String(), sessionID, readErr) - } - if exists { - return fmt.Errorf( - "session %s is already recorded in checkpoint %s; rewriting an existing checkpoint as a review is not supported yet", - sessionID, checkpointID.String(), - ) + var replacedUsage *agent.TokenUsage + if isExistingCheckpoint { + replacedUsage, err = attachReplacedEntryTokens(ctx, logCtx, repo, refs, checkpointID, sessionID, opts.Review) + if err != nil { + return err } } @@ -331,7 +317,7 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa // The checkpoint stores only tokens no earlier checkpoint of this session // counted; session state keeps the whole-transcript total. - tokenUsage, tokenPos := strategy.AttachTokenUsage(logCtx, ag, existingState, transcriptData) + tokenUsage, tokenPos := strategy.AttachTokenUsage(logCtx, ag, existingState, transcriptData, replacedUsage) sessionUsage := tokenUsage if existingState != nil { sessionUsage = agent.CalculateTokenUsage(logCtx, ag, transcriptData, 0, "") @@ -460,31 +446,67 @@ func attachSummaryLine(meta transcriptMetadata, tokenUsage *agent.TokenUsage) st return strings.Join(parts, " · ") } -// checkpointHasSessionMetadata reports whether sessionID has existing metadata -// at Primary. Reads target Primary directly, not refs.Read, because this guard -// must reflect what the next write would target. -func checkpointHasSessionMetadata(ctx context.Context, repo *git.Repository, refs cpkg.PersistentRefs, checkpointID id.CheckpointID, sessionID string) (bool, error) { +// attachReplacedEntryTokens checks the HEAD checkpoint attach is about to write +// into for an entry this session already has, and returns that entry's stored +// tokens (nil when there is none). +// +// Defense-in-depth guard: the existingState.LastCheckpointID check only fires +// when the session's state file records its checkpoint. A session already +// stored in the HEAD checkpoint but whose state is missing/stale (state file +// deleted, never written, condensed without LastCheckpointID update, or pulled +// from a remote that wasn't reflected locally) would bypass that guard. +// findSessionIndex matches by SessionID — without this check, a review-attach +// on such a session silently overwrites the existing session's metadata in the +// checkpoint. +// +// A plain attach overwrites that entry too, so the tokens it already stores +// must carry over or no checkpoint would count them. +func attachReplacedEntryTokens(ctx, logCtx context.Context, repo *git.Repository, refs cpkg.PersistentRefs, checkpointID id.CheckpointID, sessionID string, review bool) (*agent.TokenUsage, error) { + existing, found, err := checkpointSessionMetadata(ctx, repo, refs, checkpointID, sessionID) + switch { + case err != nil && review: + return nil, fmt.Errorf("failed to check checkpoint %s for session %s: %w", checkpointID.String(), sessionID, err) + case err != nil: + logging.Warn(logCtx, "attach: could not read the checkpoint entry it replaces; its stored tokens are not carried over", + slog.String("checkpoint_id", checkpointID.String()), + slog.String("error", err.Error())) + return nil, nil //nolint:nilnil // no carried-over tokens is the documented fallback + case found && review: + return nil, fmt.Errorf( + "session %s is already recorded in checkpoint %s; rewriting an existing checkpoint as a review is not supported yet", + sessionID, checkpointID.String(), + ) + case found: + return existing.TokenUsage, nil + } + return nil, nil //nolint:nilnil // nil means the checkpoint has no entry for this session +} + +// checkpointSessionMetadata returns sessionID's existing metadata at Primary. +// Reads target Primary directly, not refs.Read, because callers must see what +// the next write would replace. +func checkpointSessionMetadata(ctx context.Context, repo *git.Repository, refs cpkg.PersistentRefs, checkpointID id.CheckpointID, sessionID string) (*cpkg.Metadata, bool, error) { store, err := openAttachStore(ctx, repo, refs.PrimaryAsLocalRead()) if err != nil { - return false, err + return nil, false, err } summary, err := store.Read(ctx, checkpointID) if err != nil { - return false, fmt.Errorf("read checkpoint summary: %w", err) + return nil, false, fmt.Errorf("read checkpoint summary: %w", err) } if summary == nil { - return false, nil + return nil, false, nil } for i := range summary.Sessions { metadata, err := store.ReadSessionMetadata(ctx, checkpointID, i) if err != nil { - return false, fmt.Errorf("read session %d metadata: %w", i, err) + return nil, false, fmt.Errorf("read session %d metadata: %w", i, err) } if metadata != nil && metadata.SessionID == sessionID { - return true, nil + return metadata, true, nil } } - return false, nil + return nil, false, nil } // getHeadCommit returns the HEAD commit object. diff --git a/cmd/entire/cli/integration_test/token_scope_test.go b/cmd/entire/cli/integration_test/token_scope_test.go index 5dd20c2dd9..0d82cea506 100644 --- a/cmd/entire/cli/integration_test/token_scope_test.go +++ b/cmd/entire/cli/integration_test/token_scope_test.go @@ -202,3 +202,42 @@ func TestTokenScope_MidTurnCommitTailCountsInNextCheckpoint(t *testing.T) { 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 rewrites that +// session's entry. The entry's stored tokens must survive, plus the new turn. +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.GitCommitWithShadowHooks("a", "a.txt") + cp1 := env.TryGetLatestCheckpointID() + require.NotEmpty(t, cp1) + + // Turn 2 (7 tokens) changes no files; turn start clears LastCheckpointID, + // so attach writes into cp1 instead of returning early. + 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, "Attached session") + require.Equal(t, cp1, env.TryGetLatestCheckpointID(), "attach should write into HEAD's checkpoint") + + usage := readCommittedTokenUsage(t, env, cp1) + require.NotNil(t, usage) + require.Equal(t, 107, usage.OutputTokens, "cp1 keeps its 100 tokens and adds turn 2's 7") +} 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..a31711fa7b --- /dev/null +++ b/cmd/entire/cli/strategy/attach_token_usage_test.go @@ -0,0 +1,41 @@ +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" + +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), nil) + 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_ReplacedEntry(t *testing.T) { + t.Parallel() + replaced := &agent.TokenUsage{InputTokens: 100, OutputTokens: 100} + + withState := &SessionState{SessionID: "s", AgentType: agent.AgentTypeClaudeCode} + usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, withState, []byte(attachTestTranscript), replaced) + require.Equal(t, 110, usage.OutputTokens, "with state, the replaced entry's tokens precede TokenStart and are kept") + + usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, nil, []byte(attachTestTranscript), replaced) + require.Equal(t, 10, usage.OutputTokens, "without state the whole transcript is counted, which already covers the entry") +} diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index d6072ceb77..8a2f32ffb5 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1249,14 +1249,27 @@ func resolveCondensedTokenUsage(ctx context.Context, ag agent.Agent, state *Sess // (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. -func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte) (*agent.TokenUsage, int) { +// replaced is the usage already stored in the checkpoint entry this attach +// overwrites (same checkpoint, same session), or nil. With state, that usage +// lies before TokenStart, so it is kept and the new tokens are added to it; +// without state the whole transcript is counted, which already includes it. +func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte, replaced *agent.TokenUsage) (*agent.TokenUsage, int) { start := 0 if state != nil { start = state.TokenStart() } usage := agent.CalculateTokenUsage(ctx, ag, transcript, start, "") - if state != nil && !hasTokenUsageData(usage) && hasTokenUsageData(state.CheckpointTokenUsage) { - usage = accumulateTokenUsage(nil, state.CheckpointTokenUsage) + 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) + } + } + if state != nil && hasTokenUsageData(replaced) { + usage = accumulateTokenUsage(accumulateTokenUsage(nil, replaced), usage) } return usage, countTranscriptItems(ag.Type(), string(transcript)) } From 070cba9ad0ab2aacb4afb2b2b19a1ba26c03119b Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Wed, 7 Oct 2026 11:30:25 -0700 Subject: [PATCH 04/12] fix(strategy): pin the token offset when carry-forward restarts the display window Carry-forward set CheckpointTranscriptStart to 0 directly. An unset TokenTranscriptStart would follow it through TokenStart's fallback and bring the recount back; AdvanceDisplayWindow(0) pins it first, like the turn-end advances. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4BT334KYX2QZYXY99PCF0XJ --- cmd/entire/cli/session/state.go | 9 +++++---- cmd/entire/cli/strategy/manual_commit_hooks.go | 4 +++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index 3a389f7ec2..aa5e69a8be 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -942,10 +942,11 @@ func (s *State) AdvanceCheckpointWindow(pos int) { s.SetTokenStart(pos) } -// AdvanceDisplayWindow moves only CheckpointTranscriptStart to pos, for the -// turn-end advances that skip already-condensed content without storing its -// tokens. The token offset is pinned first so an unset one cannot follow -// CheckpointTranscriptStart forward through TokenStart's fallback. +// 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 diff --git a/cmd/entire/cli/strategy/manual_commit_hooks.go b/cmd/entire/cli/strategy/manual_commit_hooks.go index 52526a0ec7..42abbe8ed8 100644 --- a/cmd/entire/cli/strategy/manual_commit_hooks.go +++ b/cmd/entire/cli/strategy/manual_commit_hooks.go @@ -3899,7 +3899,9 @@ func (s *ManualCommitStrategy) carryForwardToNewShadowBranch( // An alternative would be incremental checkpoints (only new content since last condensation), // but this would complicate checkpoint retrieval and require careful tracking of dependencies. 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. From 2bd00767146d6545e49a87171a3f31b189e55853 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Thu, 8 Oct 2026 13:39:46 -0700 Subject: [PATCH 05/12] Add stored subagent tokens when attach replaces an entry AttachTokenUsage merged the replaced checkpoint entry with accumulateTokenUsage, which replaces SubagentTokens because it was written for cumulative snapshots. Both sides here are window deltas, so use types.AddTokenUsage to add the nested subagent totals. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4EKWN5WQQMEA5BW95K9GPC4 --- .../cli/strategy/attach_token_usage_test.go | 16 ++++++++++++++++ .../cli/strategy/manual_commit_condensation.go | 3 ++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/cmd/entire/cli/strategy/attach_token_usage_test.go b/cmd/entire/cli/strategy/attach_token_usage_test.go index a31711fa7b..5bec5b7b77 100644 --- a/cmd/entire/cli/strategy/attach_token_usage_test.go +++ b/cmd/entire/cli/strategy/attach_token_usage_test.go @@ -39,3 +39,19 @@ func TestAttachTokenUsage_ReplacedEntry(t *testing.T) { usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, nil, []byte(attachTestTranscript), replaced) require.Equal(t, 10, usage.OutputTokens, "without state the whole transcript is counted, which already covers the entry") } + +func TestAttachTokenUsage_ReplacedEntryKeepsSubagentTokens(t *testing.T) { + t.Parallel() + // Both subagent totals are window deltas: the replaced entry's is what that + // checkpoint stored, the pending one is what the session added since. + replaced := &agent.TokenUsage{OutputTokens: 100, SubagentTokens: &agent.TokenUsage{OutputTokens: 50}} + state := &SessionState{ + SessionID: "s", + AgentType: agent.AgentTypeClaudeCode, + CheckpointTokenUsage: &agent.TokenUsage{SubagentTokens: &agent.TokenUsage{OutputTokens: 5}}, + } + usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, state, []byte(attachTestTranscript), replaced) + require.Equal(t, 110, usage.OutputTokens) + require.NotNil(t, usage.SubagentTokens) + require.Equal(t, 55, usage.SubagentTokens.OutputTokens, "the replaced entry's subagent tokens must be added, not overwritten") +} diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 8a2f32ffb5..bb8e471d89 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1269,7 +1269,8 @@ func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, } } if state != nil && hasTokenUsageData(replaced) { - usage = accumulateTokenUsage(accumulateTokenUsage(nil, replaced), usage) + // Both sides are window deltas, so subagent totals add too. + usage = types.AddTokenUsage(replaced, usage) } return usage, countTranscriptItems(ag.Type(), string(transcript)) } From 47b540fde1db10a07443d02509916c508632468a Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Thu, 8 Oct 2026 13:49:26 -0700 Subject: [PATCH 06/12] Remove stray patch backup file Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4EMEBSKX0SARDNN8YHG8PS8 --- cmd/entire/cli/session_adopt_test.go.orig | 1763 --------------------- 1 file changed, 1763 deletions(-) delete mode 100644 cmd/entire/cli/session_adopt_test.go.orig diff --git a/cmd/entire/cli/session_adopt_test.go.orig b/cmd/entire/cli/session_adopt_test.go.orig deleted file mode 100644 index 78d190d554..0000000000 --- a/cmd/entire/cli/session_adopt_test.go.orig +++ /dev/null @@ -1,1763 +0,0 @@ -package cli - -import ( - "bytes" - "context" - "encoding/json" - "os" - "path/filepath" - "runtime" - "strings" - "testing" - "time" - - "github.com/entireio/cli/cmd/entire/cli/agent" - "github.com/entireio/cli/cmd/entire/cli/agent/types" - "github.com/entireio/cli/cmd/entire/cli/checkpoint/id" - "github.com/entireio/cli/cmd/entire/cli/internal/flock" - "github.com/entireio/cli/cmd/entire/cli/paths" - "github.com/entireio/cli/cmd/entire/cli/proclive" - "github.com/entireio/cli/cmd/entire/cli/session" - "github.com/entireio/cli/cmd/entire/cli/strategy" - "github.com/entireio/cli/cmd/entire/cli/testutil" -) - -func TestSessionAdopt_HelpDistinguishesForceAndYes(t *testing.T) { - cmd := newAdoptCmd() - var stdout bytes.Buffer - cmd.SetOut(&stdout) - cmd.SetArgs([]string{"--help"}) - - if err := cmd.ExecuteContext(context.Background()); err != nil { - t.Fatalf("expected help to render without error, got: %v", err) - } - - out := stdout.String() - for _, want := range []string{ - "--force", - "replace an existing local state file for the same session", - "--yes", - "confirm same-store adoption and replacement without prompting", - } { - if !strings.Contains(out, want) { - t.Fatalf("help missing %q:\n%s", want, out) - } - } - if strings.Count(out, "replace an existing local state file for the same session") != 1 { - t.Fatalf("--force and --yes should not share replacement help text:\n%s", out) - } -} - -func TestSessionAdopt_MovesExternalSessionIntoCurrentWorktree(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-session-001" - transcriptPath := claudeAdoptTranscriptPath(t, sourceRepo, sessionID) - if err := os.MkdirAll(filepath.Dir(transcriptPath), 0o750); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(transcriptPath, []byte(`{"type":"user","message":{"role":"user","content":"update target file"},"uuid":"u1"}`+"\n"), 0o600); err != nil { - t.Fatal(err) - } - - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - lastInteraction := time.Now().Add(-1 * time.Minute) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - TranscriptPath: transcriptPath, - LastPrompt: "update target file", - FilesTouched: []string{"source-only.txt"}, - TurnCheckpointIDs: []string{"abc123def456"}, - AttachedManually: true, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected adopted session state in target repo") - } - if adopted.WorktreePath != targetRepo { - t.Fatalf("WorktreePath = %q, want %q", adopted.WorktreePath, targetRepo) - } - if adopted.BaseCommit != testutil.GetHeadHash(t, targetRepo) { - t.Fatalf("BaseCommit = %q, want target HEAD", adopted.BaseCommit) - } - if adopted.TranscriptPath != transcriptPath { - t.Fatalf("TranscriptPath = %q, want %q", adopted.TranscriptPath, transcriptPath) - } - if adopted.AttachedManually { - t.Fatal("adopted active sessions should not be marked manually attached") - } - if len(adopted.FilesTouched) != 1 || adopted.FilesTouched[0] != "feature.txt" { - t.Fatalf("FilesTouched = %v, want [feature.txt]", adopted.FilesTouched) - } - if len(adopted.TurnCheckpointIDs) != 0 { - t.Fatalf("TurnCheckpointIDs = %v, want empty target-local checkpoint bookkeeping", adopted.TurnCheckpointIDs) - } - if !bytes.Contains(out.Bytes(), []byte("Adopted session")) { - t.Fatalf("output = %q, want adoption confirmation", out.String()) - } - if !bytes.Contains(out.Bytes(), []byte("Review tracked files before committing")) { - t.Fatalf("output = %q, want tracked-file attribution warning", out.String()) - } -} - -func TestSessionAdopt_ExternalStoreRetiresSourceSession(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-external-retire-source" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - LastPrompt: "continue work in target repo", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected adopted target session state") - } - if adopted.Phase != session.PhaseActive || adopted.EndedAt != nil { - t.Fatalf("target state Phase/EndedAt = %q/%v, want active/nil", adopted.Phase, adopted.EndedAt) - } - - sourceAfter, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if sourceAfter == nil { - t.Fatal("expected source session state to remain as a retired record") - } - if sourceAfter.Phase != session.PhaseEnded { - t.Fatalf("source Phase = %q, want ended", sourceAfter.Phase) - } - if sourceAfter.EndedAt == nil { - t.Fatal("source EndedAt = nil, want retirement timestamp") - } - if isAdoptableSourceSession(sourceAfter) { - t.Fatalf("source state remains adoptable after external adoption: %#v", sourceAfter) - } - - t.Chdir(sourceRepo) - sourceAgent := &mockLifecycleAgent{name: agent.AgentNameClaudeCode, agentType: agent.AgentTypeClaudeCode} - if err := handleLifecycleSessionStart(context.Background(), sourceAgent, &agent.Event{ - Type: agent.SessionStart, - SessionID: sessionID, - }); err != nil { - t.Fatalf("SessionStart in the adopted-away source repo should no-op without disrupting the hook, got: %v", err) - } - sourceAfterSessionStart, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if sourceAfterSessionStart == nil { - entries, readErr := os.ReadDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if readErr != nil { - t.Fatalf("source state disappeared after SessionStart; read state dir: %v", readErr) - } - names := make([]string, 0, len(entries)) - for _, entry := range entries { - names = append(names, entry.Name()) - } - t.Fatalf("source state disappeared after SessionStart; state dir contains %v", names) - } - if sourceAfterSessionStart.Phase != session.PhaseEnded { - t.Fatalf("source Phase after SessionStart = %q, want ended", sourceAfterSessionStart.Phase) - } - if sourceAfterSessionStart.EndedAt == nil { - t.Fatal("source EndedAt after SessionStart = nil, want retirement timestamp") - } - - err = strategy.NewManualCommitStrategy().InitializeSession( - context.Background(), - sessionID, - agent.AgentTypeClaudeCode, - "", - "source prompt after adoption", - "", - ) - if err != nil { - t.Fatalf("InitializeSession in the adopted-away source repo should no-op without disrupting the hook, got: %v", err) - } - - sourceAfterTurnStart, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if sourceAfterTurnStart.Phase != session.PhaseEnded { - t.Fatalf("source Phase after rejected TurnStart = %q, want ended", sourceAfterTurnStart.Phase) - } - if sourceAfterTurnStart.EndedAt == nil { - t.Fatal("source EndedAt after rejected TurnStart = nil, want retirement timestamp") - } -} - -func TestSessionAdopt_ExternalStoreRollsBackTargetWhenSourceRetireFails(t *testing.T) { - if runtime.GOOS == windowsGOOS { - t.Skip("uses POSIX directory permissions to force source save failure") - } - - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-retire-rollback" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStateDir := filepath.Join(sourceRepo, ".git", session.SessionStateDirName) - sourceStore := session.NewStateStoreWithDir(sourceStateDir) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - LastPrompt: "move this session", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - if err := targetStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-10 * time.Minute), - Phase: session.PhaseIdle, - BaseCommit: testutil.GetHeadHash(t, targetRepo), - WorktreePath: targetRepo, - LastPrompt: "preexisting target state", - }); err != nil { - t.Fatal(err) - } - - _, _, sourceCommonDir, err := stateStoreForWorktree(context.Background(), sourceRepo) - if err != nil { - t.Fatal(err) - } - _, _, targetCommonDir, err := stateStoreForWorktree(context.Background(), targetRepo) - if err != nil { - t.Fatal(err) - } - - info, err := os.Stat(sourceStateDir) - if err != nil { - t.Fatal(err) - } - restoreSourceStateDir := func() error { - return os.Chmod(sourceStateDir, info.Mode().Perm()) - } - if err := os.Chmod(sourceStateDir, 0o500); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { - if err := restoreSourceStateDir(); err != nil { - t.Logf("restore source state dir permissions: %v", err) - } - }) - - _, _, err = adoptFromExternalSessionStore( - context.Background(), - sourceStore, - sourceRepo, - sourceCommonDir, - targetStore, - targetCommonDir, - sessionID, - adoptOptions{Force: true}, - ) - if err := restoreSourceStateDir(); err != nil { - t.Fatalf("restore source state dir permissions: %v", err) - } - if err == nil { - t.Fatal("adoptFromExternalSessionStore succeeded, want source-retire failure") - } - if !strings.Contains(err.Error(), "retire source session state") { - t.Fatalf("adoptFromExternalSessionStore error = %v, want source-retire failure", err) - } - - loadedTarget, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loadedTarget == nil { - t.Fatal("target rollback removed preexisting state, want restore") - } - if loadedTarget.LastPrompt != "preexisting target state" { - t.Fatalf("target LastPrompt after rollback = %q, want preexisting target state", loadedTarget.LastPrompt) - } - if loadedTarget.Phase != session.PhaseIdle { - t.Fatalf("target Phase after rollback = %q, want idle", loadedTarget.Phase) - } - - sourceAfter, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if sourceAfter == nil || sourceAfter.Phase != session.PhaseActive { - t.Fatalf("source state after failed adoption = %#v, want original active state", sourceAfter) - } -} - -func TestSessionAdopt_ExternalStoreClearsNewTargetWhenSourceRetireFails(t *testing.T) { - if runtime.GOOS == windowsGOOS { - t.Skip("uses POSIX directory permissions to force source save failure") - } - - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-retire-clear-target" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStateDir := filepath.Join(sourceRepo, ".git", session.SessionStateDirName) - sourceStore := session.NewStateStoreWithDir(sourceStateDir) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - LastPrompt: "move this session", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - _, _, sourceCommonDir, err := stateStoreForWorktree(context.Background(), sourceRepo) - if err != nil { - t.Fatal(err) - } - _, _, targetCommonDir, err := stateStoreForWorktree(context.Background(), targetRepo) - if err != nil { - t.Fatal(err) - } - - info, err := os.Stat(sourceStateDir) - if err != nil { - t.Fatal(err) - } - restoreSourceStateDir := func() error { - return os.Chmod(sourceStateDir, info.Mode().Perm()) - } - if err := os.Chmod(sourceStateDir, 0o500); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { - if err := restoreSourceStateDir(); err != nil { - t.Logf("restore source state dir permissions: %v", err) - } - }) - - _, _, err = adoptFromExternalSessionStore( - context.Background(), - sourceStore, - sourceRepo, - sourceCommonDir, - targetStore, - targetCommonDir, - sessionID, - adoptOptions{Force: true}, - ) - if err := restoreSourceStateDir(); err != nil { - t.Fatalf("restore source state dir permissions: %v", err) - } - if err == nil { - t.Fatal("adoptFromExternalSessionStore succeeded, want source-retire failure") - } - if !strings.Contains(err.Error(), "retire source session state") { - t.Fatalf("adoptFromExternalSessionStore error = %v, want source-retire failure", err) - } - - loadedTarget, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loadedTarget != nil { - t.Fatalf("target state after rollback = %#v, want nil", loadedTarget) - } -} - -func TestSessionAdopt_ClearsSourceOwner(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-clear-owner" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - Owner: &proclive.Identity{PID: os.Getpid(), Start: "source-owner"}, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected adopted session state in target repo") - } - if adopted.Owner != nil { - t.Fatalf("Owner = %#v, want nil so source process liveness cannot finalize adopted session", adopted.Owner) - } -} - -func TestSessionAdopt_RejectsUnexpectedSourceTranscriptPath(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-reject-transcript" - transcriptPath := filepath.Join(t.TempDir(), sessionID+".jsonl") - if err := os.WriteFile(transcriptPath, []byte(`{"type":"user"}`+"\n"), 0o600); err != nil { - t.Fatal(err) - } - - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - lastInteraction := time.Now().Add(-1 * time.Minute) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - TranscriptPath: transcriptPath, - LastPrompt: "update target file", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err == nil { - t.Fatal("runAdopt succeeded, want transcript-path refusal") - } - if !strings.Contains(err.Error(), "unexpected transcript path") { - t.Fatalf("runAdopt error = %v, want unexpected transcript path", err) - } - - targetStore, storeErr := session.NewStateStore(context.Background()) - if storeErr != nil { - t.Fatal(storeErr) - } - adopted, loadErr := targetStore.Load(context.Background(), sessionID) - if loadErr != nil { - t.Fatal(loadErr) - } - if adopted != nil { - t.Fatalf("target state was written despite transcript-path refusal: %#v", adopted) - } -} - -func TestSessionAdopt_ExternalStoreRejectsSourceEndedAfterInitialSelection(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-external-source-stale" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - }); err != nil { - t.Fatal(err) - } - if _, err := selectAdoptSourceSession(context.Background(), sourceStore, sourceRepo, sessionID); err != nil { - t.Fatalf("initial source selection failed: %v", err) - } - - endedAt := time.Now() - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - EndedAt: &endedAt, - Phase: session.PhaseIdle, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - _, _, sourceCommonDir, err := stateStoreForWorktree(context.Background(), sourceRepo) - if err != nil { - t.Fatal(err) - } - _, _, targetCommonDir, err := stateStoreForWorktree(context.Background(), targetRepo) - if err != nil { - t.Fatal(err) - } - - _, _, err = adoptFromExternalSessionStore( - context.Background(), - sourceStore, - sourceRepo, - sourceCommonDir, - targetStore, - targetCommonDir, - sessionID, - adoptOptions{Force: true}, - ) - if err == nil { - t.Fatal("adoptFromExternalSessionStore succeeded from stale ended source, want refusal") - } - if !strings.Contains(err.Error(), "ended or fully condensed") { - t.Fatalf("adoptFromExternalSessionStore error = %v, want ended-session refusal", err) - } -} - -func TestSessionAdopt_ExternalStoreChecksTargetStateAfterLockWait(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-external-target-race" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - _, _, sourceCommonDir, err := stateStoreForWorktree(context.Background(), sourceRepo) - if err != nil { - t.Fatal(err) - } - _, _, targetCommonDir, err := stateStoreForWorktree(context.Background(), targetRepo) - if err != nil { - t.Fatal(err) - } - - lockPath := filepath.Join(targetCommonDir, "entire-session-locks", sessionID+".lock") - if err := os.MkdirAll(filepath.Dir(lockPath), 0o750); err != nil { - t.Fatal(err) - } - release, err := flock.Acquire(lockPath) - if err != nil { - t.Fatal(err) - } - - done := make(chan error, 1) - go func() { - _, _, adoptErr := adoptFromExternalSessionStore( - context.Background(), - sourceStore, - sourceRepo, - sourceCommonDir, - targetStore, - targetCommonDir, - sessionID, - adoptOptions{}, - ) - done <- adoptErr - }() - - select { - case err := <-done: - release() - t.Fatalf("adoptFromExternalSessionStore finished before target lock released: %v", err) - case <-time.After(100 * time.Millisecond): - } - - if err := targetStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now(), - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, targetRepo), - WorktreePath: targetRepo, - LastPrompt: "concurrent target state", - }); err != nil { - release() - t.Fatal(err) - } - release() - - err = <-done - if err == nil { - t.Fatal("adoptFromExternalSessionStore succeeded, want existing target refusal") - } - if !strings.Contains(err.Error(), "already tracked in this repo") { - t.Fatalf("adoptFromExternalSessionStore error = %v, want existing-state refusal", err) - } - - loaded, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loaded.LastPrompt != "concurrent target state" { - t.Fatalf("target state LastPrompt = %q, want concurrent target state", loaded.LastPrompt) - } -} - -func TestSessionAdopt_EnablesPrepareCommitMsgTrailer(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-trailer-001" - targetRelPath := "src/feature.go" - targetAbsPath := filepath.Join(targetRepo, targetRelPath) - - transcriptPath := claudeAdoptTranscriptPath(t, sourceRepo, sessionID) - if err := os.MkdirAll(filepath.Dir(transcriptPath), 0o750); err != nil { - t.Fatal(err) - } - transcript := `{"type":"human","message":{"content":"write feature.go"}} -{"type":"assistant","message":{"content":[{"type":"tool_use","name":"Write","input":{"file_path":"` + targetAbsPath + `","content":"package src\n"}}]}} -` - if err := os.WriteFile(transcriptPath, []byte(transcript), 0o600); err != nil { - t.Fatal(err) - } - stale := time.Now().Add(-3 * time.Minute) - if err := os.Chtimes(transcriptPath, stale, stale); err != nil { - t.Fatal(err) - } - - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - TranscriptPath: transcriptPath, - LastPrompt: "write feature.go", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, targetRelPath, "package src\n") - testutil.GitAdd(t, targetRepo, targetRelPath) - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - commitMsgFile := filepath.Join(targetRepo, "COMMIT_EDITMSG") - if err := os.WriteFile(commitMsgFile, []byte("add feature\n"), 0o600); err != nil { - t.Fatal(err) - } - - if err := strategy.NewManualCommitStrategy().PrepareCommitMsg(context.Background(), commitMsgFile, ""); err != nil { - t.Fatalf("PrepareCommitMsg failed: %v", err) - } - - content, err := os.ReadFile(commitMsgFile) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(string(content), "Entire-Checkpoint:") { - t.Fatalf("commit message = %q, want Entire-Checkpoint trailer", string(content)) - } -} - -func TestSessionAdopt_IdleSourceSurvivesPrepareCommitMsgTrailer(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-idle-source" - targetRelPath := "src/idle.go" - targetAbsPath := filepath.Join(targetRepo, targetRelPath) - transcriptPath := claudeAdoptTranscriptPath(t, sourceRepo, sessionID) - if err := os.MkdirAll(filepath.Dir(transcriptPath), 0o750); err != nil { - t.Fatal(err) - } - transcript := `{"type":"human","message":{"content":"write idle.go"}} -{"type":"assistant","message":{"content":[{"type":"tool_use","name":"Write","input":{"file_path":"` + targetAbsPath + `","content":"package src\n"}}]}} -` - if err := os.WriteFile(transcriptPath, []byte(transcript), 0o600); err != nil { - t.Fatal(err) - } - - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseIdle, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - TranscriptPath: transcriptPath, - LastPrompt: "write idle.go", - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, targetRelPath, "package src\n") - testutil.GitAdd(t, targetRepo, targetRelPath) - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected adopted session state") - } - if adopted.Phase != session.PhaseActive { - t.Fatalf("Phase = %q, want active so commit hooks do not sweep adopted state", adopted.Phase) - } - if adopted.EndedAt != nil { - t.Fatalf("EndedAt = %v, want nil", adopted.EndedAt) - } - - commitMsgFile := filepath.Join(targetRepo, "COMMIT_EDITMSG") - if err := os.WriteFile(commitMsgFile, []byte("add idle feature\n"), 0o600); err != nil { - t.Fatal(err) - } - if err := strategy.NewManualCommitStrategy().PrepareCommitMsg(context.Background(), commitMsgFile, ""); err != nil { - t.Fatalf("PrepareCommitMsg failed: %v", err) - } - content, err := os.ReadFile(commitMsgFile) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(string(content), "Entire-Checkpoint:") { - t.Fatalf("commit message = %q, want Entire-Checkpoint trailer", string(content)) - } -} - -func TestSessionAdopt_RejectsEndedAtSourceSession(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-ended-at" - endedAt := time.Now().Add(-30 * time.Second) - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - EndedAt: &endedAt, - Phase: session.PhaseIdle, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err == nil { - t.Fatal("runAdopt succeeded, want ended-session refusal") - } - if !strings.Contains(err.Error(), "ended or fully condensed") { - t.Fatalf("runAdopt error = %v, want ended-session refusal", err) - } - - _, err = selectAdoptSourceSession(context.Background(), sourceStore, sourceRepo, "") - if err == nil { - t.Fatal("selectAdoptSourceSession succeeded, want no recent active sessions") - } - if !strings.Contains(err.Error(), "no recent active sessions") { - t.Fatalf("selectAdoptSourceSession error = %v, want no recent active sessions", err) - } -} - -func TestSessionAdopt_ResetsSourceCheckpointWindow(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sessionID := "test-adopt-reset-window" - targetRelPath := "src/feature.go" - targetAbsPath := filepath.Join(targetRepo, targetRelPath) - - transcriptPath := claudeAdoptTranscriptPath(t, sourceRepo, sessionID) - if err := os.MkdirAll(filepath.Dir(transcriptPath), 0o750); err != nil { - t.Fatal(err) - } - transcript := `{"type":"human","message":{"content":"first source prompt"},"uuid":"source-user"} -{"type":"assistant","message":{"content":"source response"},"uuid":"source-assistant"} -{"type":"human","message":{"content":"write target feature"},"uuid":"target-user"} -{"type":"assistant","message":{"content":[{"type":"tool_use","name":"Write","input":{"file_path":"` + targetAbsPath + `","content":"package src\n"}}]},"uuid":"target-assistant"} -` - if err := os.WriteFile(transcriptPath, []byte(transcript), 0o600); err != nil { - t.Fatal(err) - } - - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - TranscriptPath: transcriptPath, - LastPrompt: "write target feature", - StepCount: 4, - SessionDurationMs: 120_000, - SessionTurnCount: 7, - ContextTokens: 42_000, - ContextWindowSize: 200_000, - CheckpointTranscriptStart: 2, - CheckpointTranscriptSize: 1234, - CondensedTranscriptLines: 2, //nolint:staticcheck // legacy field, asserted so migration keeps working - TranscriptLinesAtStart: 2, //nolint:staticcheck // legacy field, asserted so migration keeps working - TranscriptIdentifierAtStart: "source-assistant", - TurnID: "source-turn", - TurnCheckpointIDs: []string{"abc123def456"}, - LastCheckpointID: id.MustCheckpointID("abc123def456"), - CondensationAttempt: &session.CondensationAttempt{ - CheckpointID: id.MustCheckpointID("fedcba987654"), - RecoveryPending: true, - }, - CheckpointTokenUsage: &agent.TokenUsage{InputTokens: 100, OutputTokens: 25, APICallCount: 1}, - UntrackedFilesAtStart: []string{"source-only.txt"}, - PromptWindowBase: 3, - PromptWindowResetPending: true, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, targetRelPath, "package src\n") - testutil.GitAdd(t, targetRepo, targetRelPath) - testutil.WriteFile(t, targetRepo, "target-notes.txt", "user notes\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected adopted session state in target repo") - } - if adopted.StepCount != 0 { - t.Fatalf("StepCount = %d, want 0 for first target checkpoint", adopted.StepCount) - } - if adopted.CheckpointTranscriptStart != 0 { - t.Fatalf("CheckpointTranscriptStart = %d, want 0", adopted.CheckpointTranscriptStart) - } - if adopted.CheckpointTranscriptSize != 0 { - t.Fatalf("CheckpointTranscriptSize = %d, want 0", adopted.CheckpointTranscriptSize) - } - if adopted.TranscriptIdentifierAtStart != "" { - t.Fatalf("TranscriptIdentifierAtStart = %q, want empty", adopted.TranscriptIdentifierAtStart) - } - if adopted.SessionDurationMs != 120_000 { - t.Fatalf("SessionDurationMs = %d, want preserved source duration", adopted.SessionDurationMs) - } - if adopted.SessionTurnCount != 7 { - t.Fatalf("SessionTurnCount = %d, want preserved source turn count", adopted.SessionTurnCount) - } - if adopted.ContextTokens != 42_000 { - t.Fatalf("ContextTokens = %d, want preserved source context tokens", adopted.ContextTokens) - } - if adopted.ContextWindowSize != 200_000 { - t.Fatalf("ContextWindowSize = %d, want preserved source context window size", adopted.ContextWindowSize) - } - if adopted.PromptWindowBase != adopted.SessionTurnCount { - t.Fatalf("PromptWindowBase = %d, want current SessionTurnCount %d", adopted.PromptWindowBase, adopted.SessionTurnCount) - } - if adopted.PromptWindowResetPending { - t.Fatal("PromptWindowResetPending = true, want false for adopted target window") - } - if len(adopted.TurnCheckpointIDs) != 0 { - t.Fatalf("TurnCheckpointIDs = %v, want empty", adopted.TurnCheckpointIDs) - } - if adopted.TurnID != "" { - t.Fatalf("TurnID = %q, want empty target-local turn ID", adopted.TurnID) - } - if len(adopted.UntrackedFilesAtStart) != 1 || adopted.UntrackedFilesAtStart[0] != "target-notes.txt" { - t.Fatalf("UntrackedFilesAtStart = %v, want target worktree snapshot [target-notes.txt]", adopted.UntrackedFilesAtStart) - } - if !adopted.LastCheckpointID.IsEmpty() { - t.Fatalf("LastCheckpointID = %s, want empty", adopted.LastCheckpointID.String()) - } - 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) - } - - commitMsgFile := filepath.Join(targetRepo, "COMMIT_EDITMSG") - if err := os.WriteFile(commitMsgFile, []byte("add target feature\n"), 0o600); err != nil { - t.Fatal(err) - } - if err := strategy.NewManualCommitStrategy().PrepareCommitMsg(context.Background(), commitMsgFile, ""); err != nil { - t.Fatalf("PrepareCommitMsg failed: %v", err) - } - content, err := os.ReadFile(commitMsgFile) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(string(content), "Entire-Checkpoint:") { - t.Fatalf("commit message = %q, want Entire-Checkpoint trailer", string(content)) - } -} - -func TestSessionAdopt_ClearsLegacyTranscriptOffsets(t *testing.T) { - targetRepo := setupAdoptRepo(t) - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - adopted, _, err := buildAdoptedSessionState(context.Background(), &session.State{ - SessionID: "test-adopt-legacy-offsets", - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - Phase: session.PhaseActive, - BaseCommit: "source-head", - WorktreePath: "/source/repo", - CheckpointTranscriptStart: 9, - CondensedTranscriptLines: 9, //nolint:staticcheck // legacy field, asserted so migration keeps working - TranscriptLinesAtStart: 9, //nolint:staticcheck // legacy field, asserted so migration keeps working - }, "") - if err != nil { - t.Fatalf("buildAdoptedSessionState failed: %v", err) - } - if adopted.CheckpointTranscriptStart != 0 { - t.Fatalf("CheckpointTranscriptStart = %d, want 0", adopted.CheckpointTranscriptStart) - } - - encoded, err := json.Marshal(adopted) - if err != nil { - t.Fatal(err) - } - if bytes.Contains(encoded, []byte("condensed_transcript_lines")) { - t.Fatalf("adopted state JSON contains condensed_transcript_lines: %s", encoded) - } - if bytes.Contains(encoded, []byte("transcript_lines_at_start")) { - t.Fatalf("adopted state JSON contains transcript_lines_at_start: %s", encoded) - } -} - -// 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) { - for _, tc := range []struct { - name string - sourceBaseline *agent.TokenUsage - }{ - // 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}}, - } { - t.Run(tc.name, func(t *testing.T) { - targetRepo := setupAdoptRepo(t) - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - adopted, _, err := buildAdoptedSessionState(context.Background(), &session.State{ - SessionID: "test-adopt-subagent-baseline-" + tc.name, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - Phase: session.PhaseActive, - BaseCommit: "source-head", - WorktreePath: "/source/repo", - TokenUsage: &agent.TokenUsage{ - InputTokens: 1000, - OutputTokens: 500, - APICallCount: 10, - SubagentTokens: &agent.TokenUsage{InputTokens: 500, OutputTokens: 250, APICallCount: 5}, - }, - SubagentTokensBaseline: tc.sourceBaseline, - }, "") - if err != nil { - 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) - } - }) - } -} - -func TestSessionAdopt_PreservesReviewAndInvestigateMetadata(t *testing.T) { - for _, tc := range []struct { - name string - kind session.Kind - }{ - {name: "review", kind: session.KindAgentReview}, - {name: "investigate", kind: session.KindAgentInvestigate}, - } { - t.Run(tc.name, func(t *testing.T) { - targetRepo := setupAdoptRepo(t) - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - adopted, _, err := buildAdoptedSessionState(context.Background(), &session.State{ - SessionID: "test-adopt-kind-" + tc.name, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - Phase: session.PhaseActive, - Kind: tc.kind, - ReviewSkills: []string{"/review"}, - ReviewPrompt: "review this branch", - InvestigateRunID: "abcdef012345", - InvestigateTopic: "Why is adoption misclassified?", - BaseCommit: "source-head", - WorktreePath: "/source/repo", - LastCheckpointID: id.MustCheckpointID("abc123def456"), - TurnCheckpointIDs: []string{"abc123def456"}, - PromptWindowBase: 3, - SessionTurnCount: 7, - AttachedManually: true, - }, "") - if err != nil { - t.Fatalf("buildAdoptedSessionState failed: %v", err) - } - - if adopted.Kind != tc.kind { - t.Fatalf("Kind = %q, want %q", adopted.Kind, tc.kind) - } - if len(adopted.ReviewSkills) != 1 || adopted.ReviewSkills[0] != "/review" { - t.Fatalf("ReviewSkills = %v, want [/review]", adopted.ReviewSkills) - } - if adopted.ReviewPrompt != "review this branch" { - t.Fatalf("ReviewPrompt = %q, want review prompt", adopted.ReviewPrompt) - } - if adopted.InvestigateRunID != "abcdef012345" { - t.Fatalf("InvestigateRunID = %q, want source run ID", adopted.InvestigateRunID) - } - if adopted.InvestigateTopic != "Why is adoption misclassified?" { - t.Fatalf("InvestigateTopic = %q, want source topic", adopted.InvestigateTopic) - } - }) - } -} - -func TestSessionAdopt_CloneSourceStateDoesNotShareMutableFields(t *testing.T) { - lastInteraction := time.Now().Add(-1 * time.Minute) - endedAt := time.Now() - source := &session.State{ - SessionID: "test-adopt-deep-copy", - StartedAt: time.Now().Add(-5 * time.Minute), - EndedAt: &endedAt, - LastInteractionTime: &lastInteraction, - ReviewSkills: []string{"/review"}, - TurnCheckpointIDs: []string{"source-checkpoint"}, - UntrackedFilesAtStart: []string{"untracked.txt"}, - FilesTouched: []string{"source.txt"}, - TaskRecords: []session.TaskRecord{{ - ToolUseID: "toolu_1", - DeclaredTranscriptPath: "source.jsonl", - Files: []string{"task.txt"}, - TokenUsage: &agent.TokenUsage{InputTokens: 5}, - }}, - SubagentInventory: []session.SubagentInventoryEntry{{AgentID: "child", ObservedTurnIDs: []string{"turn-1"}}}, - TokenUsage: &agent.TokenUsage{ - InputTokens: 1, - SubagentTokens: &agent.TokenUsage{ - OutputTokens: 2, - }, - }, - SkillEvents: []agent.SkillEvent{ - { - ID: "skill-event", - TranscriptAnchor: &agent.SkillEventTranscriptAnchor{ - EntryIDs: []string{"entry-1"}, - }, - Native: map[string]string{"tool": "skill"}, - }, - }, - } - - adopted := cloneAdoptSourceState(source) - *adopted.EndedAt = endedAt.Add(1 * time.Hour) - *adopted.LastInteractionTime = lastInteraction.Add(1 * time.Hour) - adopted.ReviewSkills[0] = "/changed" - adopted.TurnCheckpointIDs[0] = "changed-checkpoint" - adopted.UntrackedFilesAtStart[0] = "changed-untracked.txt" - adopted.TaskRecords[0].DeclaredTranscriptPath = "" - adopted.TaskRecords[0].Files[0] = "changed-task.txt" - adopted.TaskRecords[0].TokenUsage.InputTokens = 99 - adopted.SubagentInventory[0].ObservedTurnIDs[0] = "changed-turn" - adopted.FilesTouched[0] = "changed-source.txt" - adopted.TokenUsage.SubagentTokens.OutputTokens = 99 - adopted.SkillEvents[0].TranscriptAnchor.EntryIDs[0] = "changed-entry" - adopted.SkillEvents[0].Native["tool"] = "changed-skill" - - if !source.EndedAt.Equal(endedAt) { - t.Fatalf("source EndedAt was mutated: %v", source.EndedAt) - } - if !source.LastInteractionTime.Equal(lastInteraction) { - t.Fatalf("source LastInteractionTime was mutated: %v", source.LastInteractionTime) - } - if source.ReviewSkills[0] != "/review" { - t.Fatalf("source ReviewSkills = %v, want unchanged", source.ReviewSkills) - } - if source.TurnCheckpointIDs[0] != "source-checkpoint" { - t.Fatalf("source TurnCheckpointIDs = %v, want unchanged", source.TurnCheckpointIDs) - } - if source.UntrackedFilesAtStart[0] != "untracked.txt" { - t.Fatalf("source UntrackedFilesAtStart = %v, want unchanged", source.UntrackedFilesAtStart) - } - if source.FilesTouched[0] != "source.txt" { - t.Fatalf("source FilesTouched = %v, want unchanged", source.FilesTouched) - } - if record := source.TaskRecords[0]; record.DeclaredTranscriptPath != "source.jsonl" || - record.Files[0] != "task.txt" || record.TokenUsage.InputTokens != 5 { - t.Fatalf("source TaskRecords = %+v, want unchanged", source.TaskRecords) - } - if source.SubagentInventory[0].ObservedTurnIDs[0] != "turn-1" { - t.Fatalf("source SubagentInventory = %+v, want unchanged", source.SubagentInventory) - } - if source.TokenUsage.SubagentTokens.OutputTokens != 2 { - t.Fatalf("source TokenUsage.SubagentTokens.OutputTokens = %d, want unchanged", source.TokenUsage.SubagentTokens.OutputTokens) - } - if source.SkillEvents[0].TranscriptAnchor.EntryIDs[0] != "entry-1" { - t.Fatalf("source SkillEvents entry IDs = %v, want unchanged", source.SkillEvents[0].TranscriptAnchor.EntryIDs) - } - if source.SkillEvents[0].Native["tool"] != "skill" { - t.Fatalf("source SkillEvents native = %v, want unchanged", source.SkillEvents[0].Native) - } -} - -func TestSessionAdopt_FromSubdirectoryReadsSourceStore(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetRepo := setupAdoptRepo(t) - - sourceSubdir := filepath.Join(sourceRepo, "nested", "dir") - if err := os.MkdirAll(sourceSubdir, 0o750); err != nil { - t.Fatal(err) - } - - sessionID := "test-adopt-from-subdir" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err := runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceSubdir, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed from source subdir: %v", err) - } -} - -func TestSessionAdopt_FiltersSharedSourceStoreByFromWorktree(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - siblingWorktree := filepath.Join(t.TempDir(), "sibling-worktree") - runAdoptGit(t, sourceRepo, "worktree", "add", siblingWorktree, "-b", "sibling-worktree") - resolvedSiblingWorktree, err := filepath.EvalSymlinks(siblingWorktree) - if err != nil { - t.Fatal(err) - } - siblingWorktree = resolvedSiblingWorktree - t.Cleanup(func() { - runAdoptGit(t, sourceRepo, "worktree", "remove", siblingWorktree, "--force") - }) - targetRepo := setupAdoptRepo(t) - - sourceWorktreeID, err := paths.GetWorktreeID(sourceRepo) - if err != nil { - t.Fatal(err) - } - siblingWorktreeID, err := paths.GetWorktreeID(siblingWorktree) - if err != nil { - t.Fatal(err) - } - - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: "source-worktree-session", - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - WorktreeID: sourceWorktreeID, - }); err != nil { - t.Fatal(err) - } - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: "sibling-worktree-session", - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, siblingWorktree), - WorktreePath: siblingWorktree, - WorktreeID: siblingWorktreeID, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n") - t.Chdir(targetRepo) - - var out bytes.Buffer - err = runAdopt(context.Background(), &out, "", adoptOptions{ - FromWorktree: sourceRepo, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - targetStore, err := session.NewStateStore(context.Background()) - if err != nil { - t.Fatal(err) - } - adopted, err := targetStore.Load(context.Background(), "source-worktree-session") - if err != nil { - t.Fatal(err) - } - if adopted == nil { - t.Fatal("expected source worktree session to be adopted") - } - if wrong, err := targetStore.Load(context.Background(), "sibling-worktree-session"); err != nil { - t.Fatal(err) - } else if wrong != nil { - t.Fatalf("adopted sibling worktree session unexpectedly: %#v", wrong) - } -} - -func TestSessionAdopt_RejectsSourceSessionWithoutWorktreeMetadata(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - - sessionID := "missing-worktree-metadata" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - }); err != nil { - t.Fatal(err) - } - - _, err := selectAdoptSourceSession(context.Background(), sourceStore, sourceRepo, sessionID) - if err == nil { - t.Fatal("selectAdoptSourceSession succeeded for explicit session without worktree metadata, want refusal") - } - if !strings.Contains(err.Error(), "belongs to") || !strings.Contains(err.Error(), "unknown") { - t.Fatalf("selectAdoptSourceSession error = %v, want missing-worktree ownership refusal", err) - } - - _, err = selectAdoptSourceSession(context.Background(), sourceStore, sourceRepo, "") - if err == nil { - t.Fatal("selectAdoptSourceSession auto-selected session without worktree metadata, want no candidate") - } - if !strings.Contains(err.Error(), "no recent active sessions") { - t.Fatalf("selectAdoptSourceSession error = %v, want no recent active sessions", err) - } -} - -func TestStateStoreForWorktreeIgnoresGitStderrOnSuccess(t *testing.T) { - if runtime.GOOS == windowsGOOS { - t.Skip("uses a POSIX shell script fake git") - } - - fakeBin := t.TempDir() - fakeGit := filepath.Join(fakeBin, "git") - script := `#!/bin/sh -printf 'advice: noisy git warning\n' >&2 -printf '%s\n%s\n' "$FAKE_WORKTREE_ROOT" "$FAKE_GIT_COMMON_DIR" -` - if err := os.WriteFile(fakeGit, []byte(script), 0o755); err != nil { - t.Fatal(err) - } - - sourceRoot := filepath.Join(t.TempDir(), "source") - commonDir := filepath.Join(t.TempDir(), "common.git") - t.Setenv("PATH", fakeBin+string(os.PathListSeparator)+os.Getenv("PATH")) - t.Setenv("FAKE_WORKTREE_ROOT", sourceRoot) - t.Setenv("FAKE_GIT_COMMON_DIR", commonDir) - - _, gotSourceRoot, gotCommonDir, err := stateStoreForWorktree(context.Background(), ".") - if err != nil { - t.Fatalf("stateStoreForWorktree failed: %v", err) - } - if gotSourceRoot != sourceRoot { - t.Fatalf("sourceRoot = %q, want %q", gotSourceRoot, sourceRoot) - } - if gotCommonDir != filepath.Clean(commonDir) { - t.Fatalf("commonDir = %q, want %q", gotCommonDir, filepath.Clean(commonDir)) - } -} - -func TestStateStoreForWorktreePreservesGitCommonDirSymlink(t *testing.T) { - if runtime.GOOS == windowsGOOS { - t.Skip("uses a POSIX shell script fake git") - } - - fakeBin := t.TempDir() - fakeGit := filepath.Join(fakeBin, "git") - script := `#!/bin/sh -printf '%s\n%s\n' "$FAKE_WORKTREE_ROOT" "$FAKE_GIT_COMMON_DIR" -` - if err := os.WriteFile(fakeGit, []byte(script), 0o755); err != nil { - t.Fatal(err) - } - - sourceRoot := filepath.Join(t.TempDir(), "source") - realCommonDir := filepath.Join(t.TempDir(), "real-common.git") - if err := os.MkdirAll(realCommonDir, 0o750); err != nil { - t.Fatal(err) - } - commonDirLink := filepath.Join(t.TempDir(), "common-link.git") - if err := os.Symlink(realCommonDir, commonDirLink); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - t.Setenv("PATH", fakeBin+string(os.PathListSeparator)+os.Getenv("PATH")) - t.Setenv("FAKE_WORKTREE_ROOT", sourceRoot) - t.Setenv("FAKE_GIT_COMMON_DIR", commonDirLink) - - _, _, gotCommonDir, err := stateStoreForWorktree(context.Background(), ".") - if err != nil { - t.Fatalf("stateStoreForWorktree failed: %v", err) - } - if gotCommonDir != filepath.Clean(commonDirLink) { - t.Fatalf("commonDir = %q, want git-reported symlink path %q", gotCommonDir, filepath.Clean(commonDirLink)) - } -} - -func TestSameAdoptStoreCanonicalizesGitCommonDirSymlinks(t *testing.T) { - if runtime.GOOS == windowsGOOS { - t.Skip("symlink path canonicalization is POSIX-only in this test") - } - - realCommonDir := filepath.Join(t.TempDir(), "real-common.git") - if err := os.MkdirAll(realCommonDir, 0o750); err != nil { - t.Fatal(err) - } - commonDirLink := filepath.Join(t.TempDir(), "common-link.git") - if err := os.Symlink(realCommonDir, commonDirLink); err != nil { - t.Skipf("symlinks unavailable: %v", err) - } - - if !sameAdoptStore(commonDirLink, realCommonDir) { - t.Fatalf("sameAdoptStore(%q, %q) = false, want true", commonDirLink, realCommonDir) - } -} - -func TestSessionAdopt_SameStoreReloadsSourceStateUnderLock(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetWorktree := filepath.Join(t.TempDir(), "target-worktree") - runAdoptGit(t, sourceRepo, "worktree", "add", targetWorktree, "-b", "target-worktree") - resolvedTargetWorktree, err := filepath.EvalSymlinks(targetWorktree) - if err != nil { - t.Fatal(err) - } - targetWorktree = resolvedTargetWorktree - t.Cleanup(func() { - runAdoptGit(t, sourceRepo, "worktree", "remove", targetWorktree, "--force") - }) - - sourceWorktreeID, err := paths.GetWorktreeID(sourceRepo) - if err != nil { - t.Fatal(err) - } - targetWorktreeID, err := paths.GetWorktreeID(targetWorktree) - if err != nil { - t.Fatal(err) - } - - sessionID := "test-adopt-same-store-reload" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - WorktreeID: sourceWorktreeID, - LastPrompt: "stale prompt", - SessionTurnCount: 1, - }); err != nil { - t.Fatal(err) - } - staleSelected, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - WorktreeID: sourceWorktreeID, - LastPrompt: "fresh hook prompt", - SessionTurnCount: 9, - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetWorktree, "feature.txt", "agent change\n") - testutil.GitAdd(t, targetWorktree, "feature.txt") - t.Chdir(targetWorktree) - - adopted, _, err := adoptFromSameSessionStore(context.Background(), sourceRepo, staleSelected, adoptOptions{ - Force: true, - }) - if err != nil { - t.Fatalf("adoptFromSameSessionStore failed: %v", err) - } - if adopted.LastPrompt != "fresh hook prompt" { - t.Fatalf("adopted LastPrompt = %q, want fresh hook prompt", adopted.LastPrompt) - } - if adopted.SessionTurnCount != 9 { - t.Fatalf("adopted SessionTurnCount = %d, want fresh source value", adopted.SessionTurnCount) - } - - loaded, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loaded.WorktreePath != targetWorktree { - t.Fatalf("WorktreePath = %q, want %q", loaded.WorktreePath, targetWorktree) - } - if loaded.WorktreeID != targetWorktreeID { - t.Fatalf("WorktreeID = %q, want %q", loaded.WorktreeID, targetWorktreeID) - } - if loaded.LastPrompt != "fresh hook prompt" { - t.Fatalf("loaded LastPrompt = %q, want fresh hook prompt", loaded.LastPrompt) - } - if loaded.SessionTurnCount != 9 { - t.Fatalf("loaded SessionTurnCount = %d, want fresh source value", loaded.SessionTurnCount) - } -} - -func TestSessionAdopt_MovesSameStoreSessionIntoCurrentWorktree(t *testing.T) { - sourceRepo := setupAdoptRepo(t) - targetWorktree := filepath.Join(t.TempDir(), "target-worktree") - runAdoptGit(t, sourceRepo, "worktree", "add", targetWorktree, "-b", "target-worktree") - resolvedTargetWorktree, err := filepath.EvalSymlinks(targetWorktree) - if err != nil { - t.Fatal(err) - } - targetWorktree = resolvedTargetWorktree - t.Cleanup(func() { - runAdoptGit(t, sourceRepo, "worktree", "remove", targetWorktree, "--force") - }) - - sourceWorktreeID, err := paths.GetWorktreeID(sourceRepo) - if err != nil { - t.Fatal(err) - } - targetWorktreeID, err := paths.GetWorktreeID(targetWorktree) - if err != nil { - t.Fatal(err) - } - - sessionID := "test-adopt-same-store" - lastInteraction := time.Now().Add(-1 * time.Minute) - sourceStore := session.NewStateStoreWithDir(filepath.Join(sourceRepo, ".git", session.SessionStateDirName)) - if err := sourceStore.Save(context.Background(), &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - StartedAt: time.Now().Add(-5 * time.Minute), - LastInteractionTime: &lastInteraction, - Phase: session.PhaseActive, - BaseCommit: testutil.GetHeadHash(t, sourceRepo), - WorktreePath: sourceRepo, - WorktreeID: sourceWorktreeID, - StepCount: 4, - CheckpointTranscriptStart: 2, - LastCheckpointID: id.MustCheckpointID("abc123def456"), - }); err != nil { - t.Fatal(err) - } - - testutil.WriteFile(t, targetWorktree, "feature.txt", "agent change\n") - testutil.GitAdd(t, targetWorktree, "feature.txt") - t.Chdir(targetWorktree) - - var out bytes.Buffer - err = runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - }) - if err == nil { - t.Fatal("runAdopt succeeded without --force, want existing same-store state refusal") - } - if !strings.Contains(err.Error(), "already tracked in this repo") { - t.Fatalf("runAdopt error = %v, want existing-state refusal", err) - } - - loaded, err := sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loaded.WorktreePath != sourceRepo { - t.Fatalf("WorktreePath changed without --force: %q", loaded.WorktreePath) - } - - err = runAdopt(context.Background(), &out, sessionID, adoptOptions{ - FromWorktree: sourceRepo, - Force: true, - }) - if err != nil { - t.Fatalf("runAdopt failed: %v", err) - } - - loaded, err = sourceStore.Load(context.Background(), sessionID) - if err != nil { - t.Fatal(err) - } - if loaded.WorktreePath != targetWorktree { - t.Fatalf("WorktreePath = %q, want %q", loaded.WorktreePath, targetWorktree) - } - if loaded.WorktreeID != targetWorktreeID { - t.Fatalf("WorktreeID = %q, want %q", loaded.WorktreeID, targetWorktreeID) - } - if loaded.BaseCommit != testutil.GetHeadHash(t, targetWorktree) { - t.Fatalf("BaseCommit = %q, want target HEAD", loaded.BaseCommit) - } - if loaded.StepCount != 0 { - t.Fatalf("StepCount = %d, want reset target-local checkpoint state", loaded.StepCount) - } - if loaded.CheckpointTranscriptStart != 0 { - t.Fatalf("CheckpointTranscriptStart = %d, want reset target-local transcript window", loaded.CheckpointTranscriptStart) - } - if !loaded.LastCheckpointID.IsEmpty() { - t.Fatalf("LastCheckpointID = %s, want empty target-local checkpoint ID", loaded.LastCheckpointID.String()) - } - - commitMsgFile := filepath.Join(targetWorktree, "COMMIT_EDITMSG") - if err := os.WriteFile(commitMsgFile, []byte("add same-store feature\n"), 0o600); err != nil { - t.Fatal(err) - } - if err := strategy.NewManualCommitStrategy().PrepareCommitMsg(context.Background(), commitMsgFile, ""); err != nil { - t.Fatalf("PrepareCommitMsg failed: %v", err) - } - content, err := os.ReadFile(commitMsgFile) - if err != nil { - t.Fatal(err) - } - if !strings.Contains(string(content), "Entire-Checkpoint:") { - t.Fatalf("commit message = %q, want Entire-Checkpoint trailer", string(content)) - } -} - -func setupAdoptRepo(t *testing.T) string { - t.Helper() - - repoDir := t.TempDir() - testutil.InitRepo(t, repoDir) - testutil.WriteFile(t, repoDir, "init.txt", "init\n") - testutil.GitAdd(t, repoDir, "init.txt") - testutil.GitCommit(t, repoDir, "init") - enableEntire(t, repoDir) - realRepoDir, err := filepath.EvalSymlinks(repoDir) - if err != nil { - t.Fatal(err) - } - return realRepoDir -} - -func claudeAdoptTranscriptPath(t *testing.T, sourceRepo, sessionID string) string { - t.Helper() - - transcriptDir := filepath.Join(sourceRepo, ".claude", "projects", "adopt-test") - t.Setenv("ENTIRE_TEST_CLAUDE_PROJECT_DIR", transcriptDir) - return filepath.Join(transcriptDir, sessionID+".jsonl") -} - -func runAdoptGit(t *testing.T, dir string, args ...string) { - t.Helper() - - testutil.RunGit(t, dir, args...) -} From 0e8c80bc1b5816e388be7f021858ec4daabaef10 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Thu, 8 Oct 2026 16:12:09 -0700 Subject: [PATCH 07/12] Don't double-count attach or move the token offset back Attach added the replaced checkpoint entry's tokens on top of a count from TokenStart even when the state was recreated from scratch (TokenStart 0), so the entry was counted twice. Only keep it when tokens were already checkpointed; otherwise count the whole transcript, as without state. A condensation that could not read the transcript reports 0 lines, and AdvanceCheckpointWindow(0) moved the token offset back to 0, so the next checkpoint recounted everything. Leave the token offset alone then. Also point docs and comments that still named CheckpointTranscriptStart as the token scope at the token offset. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4EWKP5QA813YPJNBMTS6W7C --- cmd/entire/cli/session/state.go | 11 ++++++++-- cmd/entire/cli/session/state_test.go | 15 +++++++++++++ .../cli/strategy/attach_token_usage_test.go | 21 +++++++++++++++---- .../strategy/manual_commit_condensation.go | 14 +++++++------ .../cli/strategy/manual_commit_types.go | 2 +- docs/architecture/copilot-token-validation.md | 2 +- docs/development/checkpoint-implementation.md | 2 +- 7 files changed, 52 insertions(+), 15 deletions(-) diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index b47c1fa5f1..9c3e9cb8a9 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -921,10 +921,17 @@ func (s *State) SetTokenStart(pos int) { } // AdvanceCheckpointWindow moves both transcript offsets to pos after a -// condensation stored the transcript and its tokens up to pos. +// 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(pos) + s.SetTokenStart(tokenStart) } // AdvanceDisplayWindow moves only CheckpointTranscriptStart to pos, for paths diff --git a/cmd/entire/cli/session/state_test.go b/cmd/entire/cli/session/state_test.go index 12d100a3a7..ae9d946cfe 100644 --- a/cmd/entire/cli/session/state_test.go +++ b/cmd/entire/cli/session/state_test.go @@ -1312,6 +1312,21 @@ func TestState_AdvanceDisplayWindowKeepsTokenOffset(t *testing.T) { 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{} diff --git a/cmd/entire/cli/strategy/attach_token_usage_test.go b/cmd/entire/cli/strategy/attach_token_usage_test.go index 5bec5b7b77..6ffba775fb 100644 --- a/cmd/entire/cli/strategy/attach_token_usage_test.go +++ b/cmd/entire/cli/strategy/attach_token_usage_test.go @@ -11,6 +11,11 @@ import ( 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 replaced entry's tokens); 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{ @@ -33,11 +38,18 @@ func TestAttachTokenUsage_ReplacedEntry(t *testing.T) { replaced := &agent.TokenUsage{InputTokens: 100, OutputTokens: 100} withState := &SessionState{SessionID: "s", AgentType: agent.AgentTypeClaudeCode} - usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, withState, []byte(attachTestTranscript), replaced) + withState.SetTokenStart(1) + usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, withState, []byte(attachTwoTurnTranscript), replaced) require.Equal(t, 110, usage.OutputTokens, "with state, the replaced entry's tokens precede TokenStart and are kept") - usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, nil, []byte(attachTestTranscript), replaced) - require.Equal(t, 10, usage.OutputTokens, "without state the whole transcript is counted, which already covers the entry") + usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, nil, []byte(attachTwoTurnTranscript), replaced) + require.Equal(t, 110, usage.OutputTokens, "without state the whole transcript is counted, which already covers the entry") + + // A state recreated from scratch (cleanup, resume) has counted nothing, so + // the whole transcript is counted, as without state, not added to the entry. + fresh := &SessionState{SessionID: "s", AgentType: agent.AgentTypeClaudeCode} + usage, _ = AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, fresh, []byte(attachTwoTurnTranscript), replaced) + require.Equal(t, 110, usage.OutputTokens, "a fresh state must not count the replaced entry twice") } func TestAttachTokenUsage_ReplacedEntryKeepsSubagentTokens(t *testing.T) { @@ -50,7 +62,8 @@ func TestAttachTokenUsage_ReplacedEntryKeepsSubagentTokens(t *testing.T) { AgentType: agent.AgentTypeClaudeCode, CheckpointTokenUsage: &agent.TokenUsage{SubagentTokens: &agent.TokenUsage{OutputTokens: 5}}, } - usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, state, []byte(attachTestTranscript), replaced) + state.SetTokenStart(1) + usage, _ := AttachTokenUsage(context.Background(), &claudecode.ClaudeCodeAgent{}, state, []byte(attachTwoTurnTranscript), replaced) require.Equal(t, 110, usage.OutputTokens) require.NotNil(t, usage.SubagentTokens) require.Equal(t, 55, usage.SubagentTokens.OutputTokens, "the replaced entry's subagent tokens must be added, not overwritten") diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 32542459c6..b16e109b71 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1116,7 +1116,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 @@ -1133,7 +1133,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). @@ -1156,9 +1156,11 @@ func resolveCondensedTokenUsage(ctx context.Context, ag agent.Agent, state *Sess // condensation does; without state it counts the whole transcript. The caller // records pos with ConsumeAttachTokenWindow once the checkpoint is written. // replaced is the usage already stored in the checkpoint entry this attach -// overwrites (same checkpoint, same session), or nil. With state, that usage -// lies before TokenStart, so it is kept and the new tokens are added to it; -// without state the whole transcript is counted, which already includes it. +// overwrites (same checkpoint, same session), or nil. When the state shows +// tokens were already checkpointed (TokenStart > 0), that usage lies before +// TokenStart, so it is kept and the new tokens are added to it. Otherwise +// (no state, or a state recreated from scratch after cleanup or resume) the +// whole transcript is counted, which already includes it. func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte, replaced *agent.TokenUsage) (*agent.TokenUsage, int) { start := 0 if state != nil { @@ -1174,7 +1176,7 @@ func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, usage = fillMissingSubagentTokensFrom(usage, state.CheckpointTokenUsage) } } - if state != nil && hasTokenUsageData(replaced) { + if state != nil && start > 0 && hasTokenUsageData(replaced) { // Both sides are window deltas, so subagent totals add too. usage = types.AddTokenUsage(replaced, usage) } 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 cd4178202d..582bca087e 100644 --- a/docs/development/checkpoint-implementation.md +++ b/docs/development/checkpoint-implementation.md @@ -56,7 +56,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 token metadata must stay scoped to `SessionState.TokenStart()` (`TokenTranscriptStart`) or the pending checkpoint delta. `TokenTranscriptStart` advances only when a condensation stores tokens (`AdvanceCheckpointWindow`); 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). +- **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` only moves forward and is set only where tokens are stored or carried over: 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". From 36c29a7c1da8735fc8647509e6878139bc78f14d Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 10:39:51 -0700 Subject: [PATCH 08/12] attach: don't move the token offset back past a hook's checkpoint saveAttachSessionState reloads the session from disk, so a hook that checkpointed while attach waited at its prompt has already moved the token offset. Consuming attach's window from the stale position moved it back and had the next checkpoint recount those tokens. Consume only when the offset is still where attach counted from; otherwise leave it. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4GVZXYFGM6SXT3S93J4TEVE --- cmd/entire/cli/attach.go | 14 ++++- cmd/entire/cli/attach_token_window_test.go | 63 ++++++++++++++++++++++ 2 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 cmd/entire/cli/attach_token_window_test.go diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index 770d4525f1..6bd7dfb523 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -1434,7 +1434,19 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS } state.TokenUsage = sessionUsage } - strategy.ConsumeAttachTokenWindow(state, tokenPos) + // A hook that checkpointed while attach waited at its prompt has already + // moved the token offset past what attach counted from; consuming from the + // stale position would move it back and have the next checkpoint recount. + countedFrom := 0 + if existingState != nil { + countedFrom = existingState.TokenStart() + } + if state.TokenStart() == countedFrom { + strategy.ConsumeAttachTokenWindow(state, tokenPos) + } else { + logging.Warn(ctx, "attach: session tokens were checkpointed while attach was waiting; leaving the token offset as the hooks set it", + slog.Int("counted_from", countedFrom), slog.Int("token_start", state.TokenStart())) + } if opts.Review { state.Kind = session.KindAgentReview state.ReviewSkills = reviewSkills 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..43ed135f41 --- /dev/null +++ b/cmd/entire/cli/attach_token_window_test.go @@ -0,0 +1,63 @@ +package cli + +import ( + "context" + "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; attach must not move it back to the position it counted to. +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) + const sessionID = "attach-window-session" + cpID := id.MustCheckpointID("a1b2c3d4e5f6") + + for _, tc := range []struct { + name string + hookMovedTo int + wantStart int + wantPendingOK bool + }{ + {name: "unchanged", hookMovedTo: 2, wantStart: 4}, + {name: "moved by a hook", hookMovedTo: 6, wantStart: 6, wantPendingOK: true}, + } { + loaded := &session.State{SessionID: sessionID, AgentType: agent.AgentTypeClaudeCode} + loaded.SetTokenStart(2) + + now := time.Now() + onDisk := &session.State{ + SessionID: sessionID, + AgentType: agent.AgentTypeClaudeCode, + Phase: session.PhaseActive, + StartedAt: now, + LastInteractionTime: &now, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 3}, + } + onDisk.SetTokenStart(tc.hookMovedTo) + require.NoError(t, store.Save(ctx, onDisk)) + + require.NoError(t, saveAttachSessionState(ctx, nil, loaded, sessionID, agent.AgentTypeClaudeCode, + "/t.jsonl", "", cpID, transcriptMetadata{}, &agent.TokenUsage{OutputTokens: 9}, 4, 4, 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) + } +} From 7498826cc3751c2527c7ab1da17ad23205e254b9 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 11:04:09 -0700 Subject: [PATCH 09/12] attach: count a running turn's subagent tokens, and test each turn once Attaching mid-turn left the turn's subagent tokens out of the attach checkpoint: hooks record them in the pending usage only at Stop, so the next checkpoint picked them up instead. Attach now reads them live the way a mid-turn condensation does, and keeps the cumulative subagent total so the re-baseline covers what the attach checkpoint stored. TestTokenScope_EachTurnCountedOnce runs a session through hook commits, a partial commit (carry-forward), two attaches to successive commits and a subagent turn, for an ended and a running session, and checks that each checkpoint holds exactly the turns it newly covers and that the checkpoints sum to the session's tokens. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4GXCE3389JHD45NHRKZNDY9 --- cmd/entire/cli/attach.go | 37 ++- .../cli/integration_test/token_turns_test.go | 232 ++++++++++++++++++ .../cli/strategy/attach_token_usage_test.go | 8 +- .../strategy/manual_commit_condensation.go | 15 +- 4 files changed, 270 insertions(+), 22 deletions(-) create mode 100644 cmd/entire/cli/integration_test/token_turns_test.go diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index 6bd7dfb523..7705835789 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -366,16 +366,7 @@ func runAttach(ctx context.Context, w, errW io.Writer, sessionID string, agentNa return fmt.Errorf("failed to get git author: %w", err) } - // The checkpoint stores only tokens no earlier checkpoint of this session - // counted, 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. Session state keeps the whole-transcript - // total. - tokenUsage, tokenPos := strategy.AttachTokenUsage(logCtx, ag, existingState, transcriptData) - sessionUsage := tokenUsage - if existingState != nil { - sessionUsage = 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. @@ -420,6 +411,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. @@ -1426,9 +1437,9 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS state.LastPrompt = meta.FirstPrompt } if sessionUsage != nil { - // Attach reads no subagent transcripts; keep the cumulative subagent - // total hooks recorded so the re-baseline below doesn't drop it. - if state.TokenUsage != nil { + // Without a cumulative subagent total of its own, keep the one hooks + // recorded so the re-baseline below doesn't drop it. + if sessionUsage.SubagentTokens == nil && state.TokenUsage != nil { sessionUsage.SubagentTokens = state.TokenUsage.SubagentTokens sessionUsage.SubagentTokensComplete = state.TokenUsage.SubagentTokensComplete } 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..c806706dc8 --- /dev/null +++ b/cmd/entire/cli/integration_test/token_turns_test.go @@ -0,0 +1,232 @@ +//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.SimulateUserPromptSubmit(r.s.ID)) + 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)) +} + +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) + r.attach() // running: stores 30 + r.more(3) + r.stop() + r.start(40, 5) + r.attach() // running: stores the 3-token tail, 40 and the subagent's 5 + r.more(4) + r.stop() + r.start(50, 0, "d.txt") + r.stop() + r.commit("c.txt", "d.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/strategy/attach_token_usage_test.go b/cmd/entire/cli/strategy/attach_token_usage_test.go index 430b2f4bfc..b7e0f5e3a3 100644 --- a/cmd/entire/cli/strategy/attach_token_usage_test.go +++ b/cmd/entire/cli/strategy/attach_token_usage_test.go @@ -24,7 +24,7 @@ func TestAttachTokenUsage_KeepsPendingSubagentTokens(t *testing.T) { 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)) + 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) @@ -37,15 +37,15 @@ 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)) + 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)) + 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)) + 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 d3d698fa4c..0f47611b38 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1123,12 +1123,17 @@ func resolveCondensedTokenUsage(ctx context.Context, ag agent.Agent, state *Sess // (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. -func AttachTokenUsage(ctx context.Context, ag agent.Agent, state *SessionState, transcript []byte) (*agent.TokenUsage, int) { - start := 0 - if state != nil { - start = state.TokenStart() +// 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) } - usage := agent.CalculateTokenUsage(ctx, ag, transcript, start, "") if state != nil { if !hasTokenUsageData(usage) && hasTokenUsageData(state.CheckpointTokenUsage) { usage = accumulateTokenUsage(nil, state.CheckpointTokenUsage) From 732b87c157c605310a9ca0caa05b8de79e9983e3 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 11:33:37 -0700 Subject: [PATCH 10/12] attach: leave all token state alone when a hook checkpointed meanwhile The race guard skipped only the offset move, so attach still overwrote the session's token total with the older figure it computed before its prompt. Skip both. Also read the on-disk token offset before the display offset moves: TokenStart falls back to CheckpointTranscriptStart for a state without its own token offset, so a brand-new session looked raced and its token window was never consumed. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4GZ2C3JW2V9XPMR034Y6RFB --- cmd/entire/cli/attach.go | 36 ++++++++------- cmd/entire/cli/attach_token_window_test.go | 51 ++++++++++++++-------- 2 files changed, 53 insertions(+), 34 deletions(-) diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index 7705835789..db10016bec 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -1400,6 +1400,9 @@ 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() // Populate BaseCommit from HEAD if not already set, so the session becomes // active and future commits in the same session receive Entire-Checkpoint trailers. @@ -1436,27 +1439,30 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS if meta.FirstPrompt != "" { state.LastPrompt = meta.FirstPrompt } - if sessionUsage != nil { - // Without a cumulative subagent total of its own, keep the one hooks - // recorded so the re-baseline below doesn't drop it. - if sessionUsage.SubagentTokens == nil && state.TokenUsage != nil { - sessionUsage.SubagentTokens = state.TokenUsage.SubagentTokens - sessionUsage.SubagentTokensComplete = state.TokenUsage.SubagentTokensComplete - } - state.TokenUsage = sessionUsage - } // A hook that checkpointed while attach waited at its prompt has already - // moved the token offset past what attach counted from; consuming from the - // stale position would move it back and have the next checkpoint recount. + // moved the token offset past what attach counted from, and recorded newer + // totals. 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. countedFrom := 0 if existingState != nil { countedFrom = existingState.TokenStart() } - if state.TokenStart() == countedFrom { + switch { + case tokenStartOnDisk != countedFrom: + logging.Warn(ctx, "attach: session tokens were checkpointed while attach was waiting; leaving token state as the hooks set it", + slog.Int("counted_from", countedFrom), slog.Int("token_start", tokenStartOnDisk)) + default: + 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) - } else { - logging.Warn(ctx, "attach: session tokens were checkpointed while attach was waiting; leaving the token offset as the hooks set it", - slog.Int("counted_from", countedFrom), slog.Int("token_start", state.TokenStart())) } 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 index 43ed135f41..f45b4a3e52 100644 --- a/cmd/entire/cli/attach_token_window_test.go +++ b/cmd/entire/cli/attach_token_window_test.go @@ -13,7 +13,8 @@ import ( ) // A hook that checkpoints while attach waits at its prompt moves the token -// offset; attach must not move it back to the position it counted to. +// 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) @@ -25,39 +26,51 @@ func TestSaveAttachSessionState_KeepsTokenOffsetMovedMeanwhile(t *testing.T) { store, err := session.NewStateStore(ctx) require.NoError(t, err) - const sessionID = "attach-window-session" cpID := id.MustCheckpointID("a1b2c3d4e5f6") for _, tc := range []struct { - name string + name string + // hookMovedTo is the token offset on disk when attach saves; -1 means + // the session had no state at all. hookMovedTo int wantStart int wantPendingOK bool + wantTotal int }{ - {name: "unchanged", hookMovedTo: 2, wantStart: 4}, - {name: "moved by a hook", hookMovedTo: 6, wantStart: 6, wantPendingOK: true}, + {name: "unchanged", hookMovedTo: 2, wantStart: 4, wantTotal: 9}, + {name: "moved by a hook", hookMovedTo: 6, wantStart: 6, wantPendingOK: true, wantTotal: 12}, + {name: "no earlier state", hookMovedTo: -1, wantStart: 4, wantTotal: 9}, } { - loaded := &session.State{SessionID: sessionID, AgentType: agent.AgentTypeClaudeCode} - loaded.SetTokenStart(2) - - now := time.Now() - onDisk := &session.State{ - SessionID: sessionID, - AgentType: agent.AgentTypeClaudeCode, - Phase: session.PhaseActive, - StartedAt: now, - LastInteractionTime: &now, - CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 3}, + sessionID := "attach-window-" + map[int]string{2: "unchanged", 6: "moved", -1: "new"}[tc.hookMovedTo] + var loaded *session.State + if tc.hookMovedTo >= 0 { + loaded = &session.State{SessionID: sessionID, AgentType: agent.AgentTypeClaudeCode} + 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: 3}, + } + onDisk.SetTokenStart(tc.hookMovedTo) + require.NoError(t, store.Save(ctx, onDisk)) } - 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, 4, attachOptions{}, nil, false), tc.name) + "/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) } } From 47e01ffa4a4db65e301e3beefe72ad0f66084a4f Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 13:32:59 -0700 Subject: [PATCH 11/12] attach: also treat a changed pending total or new checkpoint as hook activity Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4H5WYHBXPY1YJJJJHDEK133 --- cmd/entire/cli/attach.go | 26 +++++++++++----- cmd/entire/cli/attach_token_window_test.go | 30 +++++++++++++------ cmd/entire/cli/session/state.go | 6 ++-- docs/development/checkpoint-implementation.md | 2 +- 4 files changed, 45 insertions(+), 19 deletions(-) diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index db10016bec..a4f7ddfae0 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" @@ -1403,6 +1404,8 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS // 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. @@ -1439,18 +1442,27 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS if meta.FirstPrompt != "" { state.LastPrompt = meta.FirstPrompt } - // A hook that checkpointed while attach waited at its prompt has already - // moved the token offset past what attach counted from, and recorded newer - // totals. 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. + // 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: - logging.Warn(ctx, "attach: session tokens were checkpointed while attach was waiting; leaving token state as the hooks set it", + 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 sessionUsage != nil { diff --git a/cmd/entire/cli/attach_token_window_test.go b/cmd/entire/cli/attach_token_window_test.go index f45b4a3e52..6c3db1ea4a 100644 --- a/cmd/entire/cli/attach_token_window_test.go +++ b/cmd/entire/cli/attach_token_window_test.go @@ -2,6 +2,7 @@ package cli import ( "context" + "strings" "testing" "time" @@ -32,19 +33,29 @@ func TestSaveAttachSessionState_KeepsTokenOffsetMovedMeanwhile(t *testing.T) { name string // hookMovedTo is the token offset on disk when attach saves; -1 means // the session had no state at all. - hookMovedTo int - wantStart int - wantPendingOK bool - wantTotal int + 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 }{ - {name: "unchanged", hookMovedTo: 2, wantStart: 4, wantTotal: 9}, - {name: "moved by a hook", hookMovedTo: 6, wantStart: 6, wantPendingOK: true, wantTotal: 12}, + {name: "unchanged", hookMovedTo: 2, pendingOnDisk: 3, wantStart: 4, wantTotal: 9}, + {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-" + map[int]string{2: "unchanged", 6: "moved", -1: "new"}[tc.hookMovedTo] + sessionID := "attach-window-" + strings.ReplaceAll(tc.name, " ", "-") var loaded *session.State if tc.hookMovedTo >= 0 { - loaded = &session.State{SessionID: sessionID, AgentType: agent.AgentTypeClaudeCode} + loaded = &session.State{ + SessionID: sessionID, + AgentType: agent.AgentTypeClaudeCode, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 3}, + } loaded.SetTokenStart(2) now := time.Now() @@ -55,7 +66,8 @@ func TestSaveAttachSessionState_KeepsTokenOffsetMovedMeanwhile(t *testing.T) { StartedAt: now, LastInteractionTime: &now, TokenUsage: &agent.TokenUsage{OutputTokens: 12}, - CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: 3}, + CheckpointTokenUsage: &agent.TokenUsage{OutputTokens: tc.pendingOnDisk}, + LastCheckpointID: tc.checkpointOnDisk, } onDisk.SetTokenStart(tc.hookMovedTo) require.NoError(t, store.Save(ctx, onDisk)) diff --git a/cmd/entire/cli/session/state.go b/cmd/entire/cli/session/state.go index 9c3e9cb8a9..a08bbf1f5c 100644 --- a/cmd/entire/cli/session/state.go +++ b/cmd/entire/cli/session/state.go @@ -224,8 +224,10 @@ type State struct { // 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; this offset never moves back, so that - // checkpoint still counts only new tokens. Nil in state written before the + // 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"` diff --git a/docs/development/checkpoint-implementation.md b/docs/development/checkpoint-implementation.md index 51d558fb1b..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 token metadata must stay scoped to `SessionState.TokenStart()` (`TokenTranscriptStart`) or the pending checkpoint delta. `TokenTranscriptStart` only moves forward and is set only where tokens are stored or carried over: 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). +- **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". From 4997a48eb1befc7846bf87fed065846661b4b7b1 Mon Sep 17 00:00:00 2001 From: Peyton Montei Date: Fri, 9 Oct 2026 18:52:21 -0700 Subject: [PATCH 12/12] attach: leave a running session's token total to its hooks A running attach set the session total to the whole-transcript total, which already includes the running turn's tokens so far. The turn's Stop then adds its delta from turn start, so the part before the attach was counted twice. Copilot CLI heals at the next condensation; Claude Code and the other transcript agents kept the overcount. While the session is active, attach now keeps the hooks' total and moves only the cumulative subagent total, which the re-baseline needs. Ended sessions are unchanged. Co-Authored-By: Claude Opus 5.5 Entire-Checkpoint: 01M4HR5QV3AK8FVSQESQVVBFMH --- cmd/entire/cli/attach.go | 16 +- cmd/entire/cli/attach_token_window_test.go | 3 +- .../token_turns_cursor_test.go | 215 ++++++++++++++++++ .../cli/integration_test/token_turns_test.go | 19 +- 4 files changed, 248 insertions(+), 5 deletions(-) create mode 100644 cmd/entire/cli/integration_test/token_turns_cursor_test.go diff --git a/cmd/entire/cli/attach.go b/cmd/entire/cli/attach.go index a4f7ddfae0..a6a3f1e90b 100644 --- a/cmd/entire/cli/attach.go +++ b/cmd/entire/cli/attach.go @@ -1465,7 +1465,21 @@ func saveAttachSessionState(ctx context.Context, repo *git.Repository, existingS 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 sessionUsage != nil { + 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 { diff --git a/cmd/entire/cli/attach_token_window_test.go b/cmd/entire/cli/attach_token_window_test.go index 6c3db1ea4a..e5b1f28697 100644 --- a/cmd/entire/cli/attach_token_window_test.go +++ b/cmd/entire/cli/attach_token_window_test.go @@ -42,7 +42,8 @@ func TestSaveAttachSessionState_KeepsTokenOffsetMovedMeanwhile(t *testing.T) { wantPendingOK bool wantTotal int }{ - {name: "unchanged", hookMovedTo: 2, pendingOnDisk: 3, wantStart: 4, wantTotal: 9}, + // 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}, 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 index c806706dc8..9f800fa95e 100644 --- a/cmd/entire/cli/integration_test/token_turns_test.go +++ b/cmd/entire/cli/integration_test/token_turns_test.go @@ -51,7 +51,7 @@ func newTokenTurnRun(t *testing.T) *tokenTurnRun { func (r *tokenTurnRun) start(tokens, sub int, files ...string) { r.t.Helper() r.turn++ - require.NoError(r.t, r.env.SimulateUserPromptSubmit(r.s.ID)) + 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) @@ -113,6 +113,17 @@ func (r *tokenTurnRun) stop() { 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 @@ -208,17 +219,19 @@ func TestTokenScope_EachTurnCountedOnce(t *testing.T) { r.start(20, 0, "b.txt", "c.txt") r.stop() r.commit("b.txt") - r.start(30, 0) + 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", "sub1.txt") // 4-token tail + 50 + r.commit("c.txt", "d.txt", "e.txt", "sub1.txt") // 4-token tail + 50 }, }, } {