diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 012f6a9c99..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,11 +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) { - return checkpointUsage - } - - if checkpointUsage != nil && checkpointUsage.InputTokens > 0 { + 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 7915b10af7..f3308fed72 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/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() 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