From 170267dd153ff17adbf840c8b4a4ac91c47133d4 Mon Sep 17 00:00:00 2001 From: Muskan Paliwal Date: Fri, 11 Sep 2026 17:09:47 +0530 Subject: [PATCH 1/3] fix(strategy): preserve session token totals on condensation Entire-Checkpoint: 01M2848G1CNPKDA8YJDA1MDFJA --- .../strategy/manual_commit_condensation.go | 4 - .../manual_commit_condensation_test.go | 92 +++++++++++++++++++ 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 012f6a9c99..d9f2ba042b 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1190,10 +1190,6 @@ func sessionStateBackfillTokenUsage(ctx context.Context, ag agent.Agent, agentTy return checkpointUsage } - if checkpointUsage != nil && checkpointUsage.InputTokens > 0 { - return checkpointUsage - } - return nil } diff --git a/cmd/entire/cli/strategy/manual_commit_condensation_test.go b/cmd/entire/cli/strategy/manual_commit_condensation_test.go index 7915b10af7..f0bff06696 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation_test.go @@ -258,6 +258,83 @@ func TestCountTranscriptItems_CursorEmpty(t *testing.T) { } } +func TestNonCopilotCondensationPreservesSessionTokenUsage(t *testing.T) { + t.Parallel() + + sessionUsage := &agent.TokenUsage{ + InputTokens: 10_000, + OutputTokens: 999, + CacheReadTokens: 2_000, + CacheCreationTokens: 500, + APICallCount: 42, + } + state := &SessionState{ + SessionID: "s1", + AgentType: agent.AgentTypeClaudeCode, + TokenUsage: sessionUsage, + } + checkpointUsage := &agent.TokenUsage{ + InputTokens: 100, + OutputTokens: 10, + CacheReadTokens: 20, + CacheCreationTokens: 5, + APICallCount: 1, + } + + applyBackfilledSessionTokenUsage(t.Context(), nil, state, nil, checkpointUsage) + + require.Equal(t, sessionUsage, state.TokenUsage) +} + +func TestCondenseSessionByID_NonCopilotPreservesSessionTokenUsage(t *testing.T) { //nolint:paralleltest // uses t.Chdir + dir := setupGitRepo(t) + t.Chdir(dir) + + s := &ManualCommitStrategy{} + sessionID := "non-copilot-token-usage" + metadataDir := ".entire/metadata/" + sessionID + metadataDirAbs := filepath.Join(dir, metadataDir) + require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + + transcript := strings.Join([]string{ + `{"type":"human","uuid":"u1","message":{"content":"hello"}}`, + `{"type":"assistant","uuid":"u2","message":{"id":"msg_001","usage":{"input_tokens":100,"output_tokens":10}}}`, + }, "\n") + "\n" + require.NoError(t, os.WriteFile( + filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644, + )) + require.NoError(t, os.WriteFile(filepath.Join(dir, "test.txt"), []byte("agent content"), 0o644)) + + sessionUsage := &agent.TokenUsage{ + InputTokens: 10_000, + OutputTokens: 999, + CacheReadTokens: 2_000, + CacheCreationTokens: 500, + APICallCount: 42, + } + require.NoError(t, s.SaveStep(t.Context(), StepContext{ + SessionID: sessionID, + ModifiedFiles: []string{"test.txt"}, + MetadataDir: metadataDir, + CommitMessage: "Checkpoint 1", + AuthorName: "Test", + AuthorEmail: "test@test.com", + AgentType: agent.AgentTypeClaudeCode, + TokenUsage: sessionUsage, + })) + + state, err := s.loadSessionState(t.Context(), sessionID) + require.NoError(t, err) + require.Equal(t, sessionUsage, state.TokenUsage) + + require.NoError(t, s.CondenseSessionByID(t.Context(), sessionID)) + + state, err = s.loadSessionState(t.Context(), sessionID) + require.NoError(t, err) + require.Equal(t, sessionUsage, state.TokenUsage) + require.Nil(t, state.CheckpointTokenUsage) +} + func TestSessionStateBackfillTokenUsage_CopilotUsesZeroInputSessionAggregate(t *testing.T) { t.Parallel() @@ -284,6 +361,21 @@ func TestSessionStateBackfillTokenUsage_CopilotUsesZeroInputSessionAggregate(t * require.Equal(t, 3, backfillUsage.APICallCount) } +func TestSessionStateBackfillTokenUsage_CopilotFallsBackToCheckpointUsage(t *testing.T) { + t.Parallel() + + checkpointUsage := &agent.TokenUsage{ + OutputTokens: 25, + APICallCount: 1, + } + + backfillUsage := sessionStateBackfillTokenUsage( + t.Context(), nil, agent.AgentTypeCopilotCLI, nil, checkpointUsage, + ) + + require.Same(t, checkpointUsage, backfillUsage) +} + func TestSessionStateBackfillModel_PiReadsModelFromTranscript(t *testing.T) { t.Parallel() From 5c44abe7ce8f2de66eff36941841e4a3035c6037 Mon Sep 17 00:00:00 2001 From: Muskan Paliwal Date: Fri, 11 Sep 2026 21:13:26 +0530 Subject: [PATCH 2/3] test(strategy): cover session token backfill boundaries Entire-Checkpoint: 01M28J6MX4M92Y88PWYE3WPRFK --- .../manual_commit_condensation_test.go | 49 +++++++++++++++++++ .../cli/strategy/subagent_tokens_test.go | 21 ++------ 2 files changed, 53 insertions(+), 17 deletions(-) diff --git a/cmd/entire/cli/strategy/manual_commit_condensation_test.go b/cmd/entire/cli/strategy/manual_commit_condensation_test.go index f0bff06696..5a60611171 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation_test.go @@ -286,6 +286,24 @@ func TestNonCopilotCondensationPreservesSessionTokenUsage(t *testing.T) { require.Equal(t, sessionUsage, state.TokenUsage) } +func TestNonCopilotCondensationDoesNotPromoteCheckpointUsage(t *testing.T) { + t.Parallel() + + state := &SessionState{ + SessionID: "s1", + AgentType: agent.AgentTypeClaudeCode, + } + checkpointUsage := &agent.TokenUsage{ + InputTokens: 100, + OutputTokens: 10, + APICallCount: 1, + } + + applyBackfilledSessionTokenUsage(t.Context(), nil, state, nil, checkpointUsage) + + require.Nil(t, state.TokenUsage) +} + func TestCondenseSessionByID_NonCopilotPreservesSessionTokenUsage(t *testing.T) { //nolint:paralleltest // uses t.Chdir dir := setupGitRepo(t) t.Chdir(dir) @@ -376,6 +394,37 @@ func TestSessionStateBackfillTokenUsage_CopilotFallsBackToCheckpointUsage(t *tes require.Same(t, checkpointUsage, backfillUsage) } +func TestApplyBackfilledSessionTokenUsage_CopilotPreservesSubagentTotal(t *testing.T) { + t.Parallel() + + checkpointUsage := &agent.TokenUsage{ + OutputTokens: 25, + APICallCount: 1, + } + state := &SessionState{ + AgentType: agent.AgentTypeCopilotCLI, + TokenUsage: &agent.TokenUsage{ + InputTokens: 1_000, + SubagentTokens: &agent.TokenUsage{ + InputTokens: 200, + OutputTokens: 50, + APICallCount: 2, + }, + }, + } + + applyBackfilledSessionTokenUsage(t.Context(), nil, state, nil, checkpointUsage) + + require.Equal(t, 25, state.TokenUsage.OutputTokens) + require.Equal(t, 1, state.TokenUsage.APICallCount) + require.Equal(t, &agent.TokenUsage{ + InputTokens: 200, + OutputTokens: 50, + APICallCount: 2, + }, state.TokenUsage.SubagentTokens) + require.Nil(t, checkpointUsage.SubagentTokens) +} + func TestSessionStateBackfillModel_PiReadsModelFromTranscript(t *testing.T) { t.Parallel() diff --git a/cmd/entire/cli/strategy/subagent_tokens_test.go b/cmd/entire/cli/strategy/subagent_tokens_test.go index 86cd25e96d..98559f1024 100644 --- a/cmd/entire/cli/strategy/subagent_tokens_test.go +++ b/cmd/entire/cli/strategy/subagent_tokens_test.go @@ -518,9 +518,6 @@ func TestCalculateLiveTranscriptTokenUsage_RescopesSubagentCumulativeTotal(t *te require.Equal(t, 200, state.TokenUsage.SubagentTokens.InputTokens, "session state must retain the cumulative snapshot for the next baseline") - applyBackfilledSessionTokenUsage(t.Context(), ag, state, mainTranscript, usage) - require.Equal(t, 200, state.TokenUsage.SubagentTokens.InputTokens, - "main-token backfill must not replace the cumulative snapshot with the checkpoint delta") state.RebaselineSubagentTokens() require.NoError(t, os.WriteFile(subagentPath, []byte(`{"type":"assistant","uuid":"a-sub","message":{"id":"msg_sub","type":"message","role":"assistant","content":[{"type":"text","text":"done"}],"usage":{"input_tokens":260,"output_tokens":35}}} `), 0o644)) @@ -627,15 +624,8 @@ func TestCondenseSessionByID_CapturesSubagentBaselineViaRealResetPath(t *testing metadataDir := ".entire/metadata/" + sessionID metadataDirAbs := filepath.Join(dir, metadataDir) require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) - // The assistant line carries real usage data (message.id + usage). Real - // Claude Code transcripts always do, which makes sessionStateBackfillTokenUsage - // fire during condensation (its InputTokens > 0 branch) and overwrite - // state.TokenUsage with the transcript-recomputed value — which is computed - // with subagentsDir="" and therefore drops SubagentTokens. This is what makes - // this test guard the REAL condensation path: without preserving the - // cumulative subagent total across the backfill, resetCheckpointWindow would - // snapshot a nil baseline and the next checkpoint would re-report the full - // cumulative subagent total (finding 019f5ebf-a57e). + // Checkpoint-scoped transcript usage must not replace the cumulative session + // total here, so the reset can retain the subagent baseline. transcript := `{"type":"human","message":{"content":"do the thing"}} {"type":"assistant","uuid":"a1","message":{"id":"m1","usage":{"input_tokens":300,"output_tokens":150}}} ` @@ -715,11 +705,8 @@ func TestCondenseSessionByID_CapturesSubagentBaselineViaRealResetPath(t *testing require.Equal(t, 60, summary.TokenUsage.SubagentTokens.OutputTokens) } -// TestWithSubagentTokensFrom_DoesNotMutateInput guards the copy semantics directly. -// The condensation tests cannot: applyBackfilledSessionTokenUsage already hands back -// a copy on that path, so a mutate-in-place implementation passes them. Mutating -// would overwrite the session-wide cumulative with a window delta and make -// resetCheckpointWindow snapshot a too-small baseline for the next window. +// Session and checkpoint token snapshots can share pointers, so replacement must +// not mutate either input. func TestSubagentCoverageSurvivesBackfill(t *testing.T) { t.Parallel() incomplete := false From b7f56e460af08db15c74e2bba77d9392d57f6fc6 Mon Sep 17 00:00:00 2001 From: Muskan Paliwal Date: Sat, 12 Sep 2026 12:28:19 +0530 Subject: [PATCH 3/3] refactor(strategy): simplify session token backfill Entire-Checkpoint: 01M2A6HV96SV0206QW63R1V23N --- cmd/entire/cli/strategy/manual_commit_condensation.go | 8 ++++++-- .../cli/strategy/manual_commit_condensation_test.go | 10 +++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index d9f2ba042b..92e52a00f1 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -1178,7 +1178,11 @@ func applyBackfilledSessionTokenUsage(ctx context.Context, ag agent.Agent, state // sessionStateBackfillTokenUsage returns the best session-level token usage to // persist in session state after condensation. func sessionStateBackfillTokenUsage(ctx context.Context, ag agent.Agent, agentType types.AgentType, transcript []byte, checkpointUsage *agent.TokenUsage) *agent.TokenUsage { - if agentType == agent.AgentTypeCopilotCLI && len(transcript) > 0 { + if agentType != agent.AgentTypeCopilotCLI { + return nil + } + + if len(transcript) > 0 { fullSessionUsage := agent.CalculateTokenUsage(ctx, ag, transcript, 0, "") if hasTokenUsageData(fullSessionUsage) { return fullSessionUsage @@ -1186,7 +1190,7 @@ func sessionStateBackfillTokenUsage(ctx context.Context, ag agent.Agent, agentTy logging.Debug(ctx, "copilot-cli: full-session token read produced no data, falling back to checkpoint usage") } - if agentType == agent.AgentTypeCopilotCLI && hasTokenUsageData(checkpointUsage) { + if hasTokenUsageData(checkpointUsage) { return checkpointUsage } diff --git a/cmd/entire/cli/strategy/manual_commit_condensation_test.go b/cmd/entire/cli/strategy/manual_commit_condensation_test.go index 5a60611171..f3308fed72 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation_test.go @@ -310,18 +310,14 @@ func TestCondenseSessionByID_NonCopilotPreservesSessionTokenUsage(t *testing.T) s := &ManualCommitStrategy{} sessionID := "non-copilot-token-usage" - metadataDir := ".entire/metadata/" + sessionID - metadataDirAbs := filepath.Join(dir, metadataDir) - require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + metadataDir := paths.SessionMetadataDirFromSessionID(sessionID) transcript := strings.Join([]string{ `{"type":"human","uuid":"u1","message":{"content":"hello"}}`, `{"type":"assistant","uuid":"u2","message":{"id":"msg_001","usage":{"input_tokens":100,"output_tokens":10}}}`, }, "\n") + "\n" - require.NoError(t, os.WriteFile( - filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644, - )) - require.NoError(t, os.WriteFile(filepath.Join(dir, "test.txt"), []byte("agent content"), 0o644)) + testutil.WriteFile(t, dir, filepath.Join(metadataDir, paths.TranscriptFileName), transcript) + testutil.WriteFile(t, dir, "test.txt", "agent content") sessionUsage := &agent.TokenUsage{ InputTokens: 10_000,