Preserve session-wide token totals during condensation

cmd/entire/cli/strategy/manual_commit_condensation.go

@@ -1178,19 +1178,19 @@ 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
 		}
 		logging.Debug(ctx, "copilot-cli: full-session token read produced no data, falling back to checkpoint usage")
 	}
 
-	if agentType == agent.AgentTypeCopilotCLI && hasTokenUsageData(checkpointUsage) {
-		return checkpointUsage
-	}
-
-	if checkpointUsage != nil && checkpointUsage.InputTokens > 0 {
+	if hasTokenUsageData(checkpointUsage) {
 		return checkpointUsage
 	}
 

cmd/entire/cli/strategy/manual_commit_condensation_test.go

@@ -258,6 +258,97 @@ 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 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)
+
+	s := &ManualCommitStrategy{}
+	sessionID := "non-copilot-token-usage"
+	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"
+	testutil.WriteFile(t, dir, filepath.Join(metadataDir, paths.TranscriptFileName), transcript)
+	testutil.WriteFile(t, dir, "test.txt", "agent content")
+
+	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 +375,52 @@ 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 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()
 

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