Merge pull request #1710 from entireio/fix/329-subagent-extract

fix(attribution): scan full transcript for subagents spawned before checkpoint

Author: peyton-alt

cmd/entire/cli/agent/agent.go

@@ -361,6 +361,24 @@ type SubagentAwareExtractor interface {
 	ExtractAllModifiedFiles(transcriptData []byte, fromOffset int, subagentsDir string) ([]string, error)
 
 	// CalculateTotalTokenUsage computes token usage including all spawned subagents.
-	// The subagentsDir parameter specifies where subagent transcripts are stored.
+	// The subagentsDir parameter specifies where subagent transcripts are stored
+	// (an empty subagentsDir skips subagent accounting and leaves SubagentTokens nil).
+	//
+	// CONTRACT — the returned SubagentTokens is a CUMULATIVE-SINCE-SESSION-START
+	// snapshot, NOT a delta scoped to fromOffset like the main-agent fields
+	// (InputTokens/OutputTokens/...). Implementations MUST discover spawned agent
+	// IDs from the FULL transcript prefix [0,end) — so a subagent spawned before
+	// fromOffset is still found (#329) — and re-read each subagent transcript from
+	// line 0 on every call. Consequently a subagent's full total repeats on every
+	// call after it is first discovered.
+	//
+	// Callers that accumulate across checkpoints/turns therefore MUST NOT sum
+	// SubagentTokens across calls: replace the running total with the latest
+	// snapshot, and rescope any window delta by subtracting a previously captured
+	// baseline (see accumulateTokenUsage / resetCheckpointWindow and
+	// session.State.SubagentTokensBaseline in cmd/entire/cli/strategy, and
+	// rescopeSubagentTokensToDeltas in cmd/entire/cli/agentimport for the import
+	// path). An implementation that instead returned per-window deltas would
+	// silently break that accounting with no compile-time or test signal.
 	CalculateTotalTokenUsage(transcriptData []byte, fromOffset int, subagentsDir string) (*TokenUsage, error)
 }

cmd/entire/cli/agent/claudecode/transcript.go

@@ -395,11 +395,37 @@ func (c *ClaudeCodeAgent) CalculateTotalTokenUsage(transcriptData []byte, startL
 	// Calculate token usage from parsed transcript
 	mainUsage := CalculateTokenUsage(parsed)
 
-	// Extract spawned agent IDs from the same parsed transcript
-	agentIDs := ExtractSpawnedAgentIDs(parsed)
+	if subagentsDir == "" {
+		return mainUsage, nil
+	}
+
+	// Extract spawned agent IDs from the FULL transcript (startLine=0), not the
+	// sliced portion. A subagent spawned before this checkpoint's startLine can
+	// keep writing to its transcript in later turns; scanning only the slice
+	// would miss it and undercount subagent token usage (#329).
+	//
+	// PERF (considered, retained deliberately): this re-parses the full
+	// transcript in addition to the sliced parse above — two JSONL parses per
+	// call, growing with session length. A single-pass version was rejected as
+	// not worth the risk: ParseFromBytes silently drops malformed lines, so a
+	// parsed-entry index does not correspond to a raw line number and naively
+	// slicing the full parse at startLine would misattribute main-agent usage;
+	// doing it safely would mean threading raw-line numbers through the shared
+	// transcript parser used by every agent. A cheap line scan for the Task
+	// marker instead of a full parse would duplicate ExtractSpawnedAgentIDs'
+	// nested tool_result decoding. The common no-subagent case already avoids
+	// this cost entirely via the subagentsDir == "" short-circuit above.
+	fullParsed, err := transcript.ParseFromBytes(transcriptData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse full transcript: %w", err)
+	}
+	agentIDs := ExtractSpawnedAgentIDs(fullParsed)
 
-	// Calculate subagent token usage (skip when subagentsDir is empty to avoid reading from cwd)
-	if len(agentIDs) > 0 && subagentsDir != "" {
+	// Calculate subagent token usage. This re-reads each subagent transcript from
+	// line 0 on every call, so mainUsage.SubagentTokens is a cumulative-since-
+	// session-start snapshot — see the CalculateTotalTokenUsage interface contract
+	// in cmd/entire/cli/agent for how callers must accumulate it.
+	if len(agentIDs) > 0 {
 		subagentUsage := &agent.TokenUsage{}
 		for agentID := range agentIDs {
 			agentPath := filepath.Join(subagentsDir, fmt.Sprintf("agent-%s.jsonl", agentID))
@@ -449,11 +475,23 @@ func (c *ClaudeCodeAgent) ExtractAllModifiedFiles(transcriptData []byte, startLi
 		}
 	}
 
-	// Find spawned subagents and collect their modified files (skip when subagentsDir is empty to avoid reading from cwd)
-	agentIDs := ExtractSpawnedAgentIDs(parsed)
 	if subagentsDir == "" {
 		return files, nil
 	}
+
+	// Find spawned subagents from the FULL transcript (startLine=0): a subagent
+	// spawned before this checkpoint's startLine may keep modifying files in
+	// later turns, and scanning only the slice would miss it (#329). Main-agent
+	// file extraction above stays scoped to the slice.
+	//
+	// PERF: the second full-transcript parse is retained deliberately for the
+	// same reasons documented on CalculateTotalTokenUsage above; the common
+	// no-subagent case is short-circuited by the subagentsDir == "" guard.
+	fullParsed, err := transcript.ParseFromBytes(transcriptData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse full transcript: %w", err)
+	}
+	agentIDs := ExtractSpawnedAgentIDs(fullParsed)
 	for agentID := range agentIDs {
 		agentPath := filepath.Join(subagentsDir, fmt.Sprintf("agent-%s.jsonl", agentID))
 		agentLines, agentErr := transcript.ParseFromFileAtLine(agentPath, 0)

cmd/entire/cli/agent/claudecode/transcript_test.go

@@ -887,3 +887,78 @@ func TestExtractAllModifiedFiles_SubagentOnlyChanges(t *testing.T) {
 		t.Errorf("missing expected file %q", f)
 	}
 }
+
+// Regression for #329: a subagent spawned BEFORE the checkpoint's startLine
+// must still be discovered, because it can keep modifying files in later turns.
+// The Task spawn/result live in lines before startLine; only the full transcript
+// scan finds them.
+func TestExtractAllModifiedFiles_FindsSubagentSpawnedBeforeStartLine(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := t.TempDir()
+	subagentsDir := tmpDir + "/tasks/toolu_task1"
+	c := &ClaudeCodeAgent{}
+	if err := os.MkdirAll(subagentsDir, 0o755); err != nil {
+		t.Fatalf("failed to create subagents dir: %v", err)
+	}
+
+	transcriptData := buildJSONL(
+		makeTaskToolUseLine(t, "a1", "toolu_taskA"),        // line 0 (before startLine)
+		makeTaskResultLine(t, "uA", "toolu_taskA", "subA"), // line 1 (before startLine)
+		makeWriteToolLine(t, "a2", "/repo/main.go"),        // line 2 (>= startLine)
+	)
+	writeJSONLFile(t, subagentsDir+"/agent-subA.jsonl",
+		makeWriteToolLine(t, "sa1", "/repo/helper.go"),
+	)
+
+	files, err := c.ExtractAllModifiedFiles(transcriptData, 2, subagentsDir)
+	if err != nil {
+		t.Fatalf("ExtractAllModifiedFiles() error: %v", err)
+	}
+
+	got := make(map[string]bool, len(files))
+	for _, f := range files {
+		got[f] = true
+	}
+	if !got["/repo/main.go"] {
+		t.Errorf("missing main-agent file /repo/main.go: %v", files)
+	}
+	if !got["/repo/helper.go"] {
+		t.Errorf("subagent spawned before startLine was not discovered; missing /repo/helper.go: %v", files)
+	}
+}
+
+// Regression for #329: subagent token usage must be counted even when the
+// subagent was spawned before the checkpoint's startLine.
+func TestCalculateTotalTokenUsage_CountsSubagentSpawnedBeforeStartLine(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := t.TempDir()
+	subagentsDir := tmpDir + "/tasks/toolu_task1"
+	c := &ClaudeCodeAgent{}
+	if err := os.MkdirAll(subagentsDir, 0o755); err != nil {
+		t.Fatalf("failed to create subagents dir: %v", err)
+	}
+
+	// Subagent spawned in lines 0-1 (before startLine=2); main usage on line 2.
+	transcriptData := buildJSONL(
+		makeTaskToolUseLine(t, "a1", "toolu_taskB"),
+		makeTaskResultLine(t, "uB", "toolu_taskB", "subB"),
+		`{"type":"assistant","uuid":"a2","message":{"id":"m2","usage":{"input_tokens":300,"output_tokens":150}}}`,
+	)
+	writeJSONLFile(t, subagentsDir+"/agent-subB.jsonl",
+		`{"type":"assistant","uuid":"sa1","message":{"id":"sm1","usage":{"input_tokens":50,"output_tokens":25}}}`,
+	)
+
+	usage, err := c.CalculateTotalTokenUsage(transcriptData, 2, subagentsDir)
+	if err != nil {
+		t.Fatalf("CalculateTotalTokenUsage() error: %v", err)
+	}
+	if usage.SubagentTokens == nil {
+		t.Fatal("subagent spawned before startLine was not counted (SubagentTokens is nil)")
+	}
+	if usage.SubagentTokens.InputTokens != 50 || usage.SubagentTokens.OutputTokens != 25 {
+		t.Errorf("subagent tokens = input %d output %d, want input 50 output 25",
+			usage.SubagentTokens.InputTokens, usage.SubagentTokens.OutputTokens)
+	}
+}

cmd/entire/cli/agent/factoryaidroid/transcript.go

@@ -366,8 +366,32 @@ func CalculateTotalTokenUsageFromBytes(data []byte, startLine int, subagentsDir
 
 	mainUsage := CalculateTokenUsage(parsed)
 
-	agentIDs := ExtractSpawnedAgentIDs(parsed)
-	if len(agentIDs) > 0 && subagentsDir != "" {
+	if subagentsDir == "" {
+		return mainUsage, nil
+	}
+
+	// Extract spawned agent IDs from the FULL transcript (startLine=0): a
+	// subagent spawned before this checkpoint's startLine can keep writing to
+	// its transcript, so scanning only the slice would undercount it (#329).
+	//
+	// PERF (considered, retained deliberately): this re-parses the full
+	// transcript in addition to the sliced parse above — two JSONL parses per
+	// call, growing with session length. A single-pass version was rejected:
+	// the Droid parser drops non-message / malformed lines, so a parsed-entry
+	// index does not map to a raw line number and naively slicing the full parse
+	// at startLine would misattribute main-agent usage; doing it safely would
+	// mean threading raw-line numbers through the shared parser. The common
+	// no-subagent case already avoids this via the subagentsDir == "" guard.
+	fullParsed, _, err := ParseDroidTranscriptFromBytes(data, 0)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse full transcript: %w", err)
+	}
+	agentIDs := ExtractSpawnedAgentIDs(fullParsed)
+	// This re-reads each subagent transcript from line 0 on every call below, so
+	// mainUsage.SubagentTokens ends up cumulative-since-session-start — see the
+	// CalculateTotalTokenUsage interface contract in cmd/entire/cli/agent for how
+	// callers must accumulate it (shared with Claude Code).
+	if len(agentIDs) > 0 {
 		subagentUsage := &agent.TokenUsage{}
 		for agentID := range agentIDs {
 			agentPath := filepath.Join(subagentsDir, fmt.Sprintf("agent-%s.jsonl", agentID))
@@ -409,10 +433,23 @@ func ExtractAllModifiedFilesFromBytes(data []byte, startLine int, subagentsDir s
 		fileSet[f] = true
 	}
 
-	agentIDs := ExtractSpawnedAgentIDs(parsed)
 	if subagentsDir == "" {
 		return files, nil
 	}
+
+	// Find spawned subagents from the FULL transcript (startLine=0): a subagent
+	// spawned before this checkpoint's startLine may keep modifying files in
+	// later turns, and scanning only the slice would miss it (#329). Main-agent
+	// file extraction above stays scoped to the slice.
+	//
+	// PERF: the second full-transcript parse is retained deliberately for the
+	// same reasons documented on CalculateTotalTokenUsageFromBytes above; the
+	// common no-subagent case is short-circuited by the subagentsDir == "" guard.
+	fullParsed, _, err := ParseDroidTranscriptFromBytes(data, 0)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse full transcript: %w", err)
+	}
+	agentIDs := ExtractSpawnedAgentIDs(fullParsed)
 	for agentID := range agentIDs {
 		agentPath := filepath.Join(subagentsDir, fmt.Sprintf("agent-%s.jsonl", agentID))
 		agentLines, _, agentErr := ParseDroidTranscript(agentPath, 0)

cmd/entire/cli/agent/factoryaidroid/transcript_test.go

@@ -1224,3 +1224,73 @@ func TestExtractAllModifiedFilesFromBytes_SubagentOnlyChanges(t *testing.T) {
 		t.Errorf("missing expected file %q", f)
 	}
 }
+
+// Regression for #329: a subagent spawned BEFORE the checkpoint's startLine must
+// still be discovered for file extraction (it can keep modifying files later).
+func TestExtractAllModifiedFilesFromBytes_FindsSubagentSpawnedBeforeStartLine(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := t.TempDir()
+	subagentsDir := tmpDir + "/tasks/toolu_task1"
+	if err := os.MkdirAll(subagentsDir, 0o755); err != nil {
+		t.Fatalf("failed to create subagents dir: %v", err)
+	}
+
+	data := joinJSONL(
+		makeTaskToolUseLine(t, "a1", "toolu_taskC"),        // line 0 (before startLine)
+		makeTaskResultLine(t, "uC", "toolu_taskC", "sub1"), // line 1 (before startLine)
+		makeWriteToolLine(t, "a2", "/repo/main.go"),        // line 2 (>= startLine)
+	)
+	writeJSONLFile(t, subagentsDir+"/agent-sub1.jsonl",
+		makeWriteToolLine(t, "sa1", "/repo/helper.go"),
+	)
+
+	files, err := ExtractAllModifiedFilesFromBytes(data, 2, subagentsDir)
+	if err != nil {
+		t.Fatalf("ExtractAllModifiedFilesFromBytes() error: %v", err)
+	}
+
+	got := make(map[string]bool, len(files))
+	for _, f := range files {
+		got[f] = true
+	}
+	if !got["/repo/main.go"] {
+		t.Errorf("missing main-agent file /repo/main.go: %v", files)
+	}
+	if !got["/repo/helper.go"] {
+		t.Errorf("subagent spawned before startLine was not discovered; missing /repo/helper.go: %v", files)
+	}
+}
+
+// Regression for #329: subagent token usage must be counted even when the
+// subagent was spawned before the checkpoint's startLine.
+func TestCalculateTotalTokenUsageFromBytes_CountsSubagentSpawnedBeforeStartLine(t *testing.T) {
+	t.Parallel()
+
+	tmpDir := t.TempDir()
+	subagentsDir := tmpDir + "/tasks/toolu_task1"
+	if err := os.MkdirAll(subagentsDir, 0o755); err != nil {
+		t.Fatalf("failed to create subagents dir: %v", err)
+	}
+
+	data := joinJSONL(
+		makeTaskToolUseLine(t, "a1", "toolu_taskD"),           // line 0 (before startLine)
+		makeTaskResultLine(t, "uD", "toolu_taskD", "sub1"),    // line 1 (before startLine)
+		makeAssistantTokenLine(t, "a2", "msg_main", 300, 150), // line 2
+	)
+	writeJSONLFile(t, subagentsDir+"/agent-sub1.jsonl",
+		makeAssistantTokenLine(t, "sa1", "msg_sub", 50, 25),
+	)
+
+	usage, err := CalculateTotalTokenUsageFromBytes(data, 2, subagentsDir)
+	if err != nil {
+		t.Fatalf("CalculateTotalTokenUsageFromBytes() error: %v", err)
+	}
+	if usage.SubagentTokens == nil {
+		t.Fatal("subagent spawned before startLine was not counted (SubagentTokens is nil)")
+	}
+	if usage.SubagentTokens.InputTokens != 50 || usage.SubagentTokens.OutputTokens != 25 {
+		t.Errorf("subagent tokens = input %d output %d, want input 50 output 25",
+			usage.SubagentTokens.InputTokens, usage.SubagentTokens.OutputTokens)
+	}
+}

cmd/entire/cli/agent/types/token_usage.go

@@ -45,3 +45,35 @@ func AddTokenUsage(a, b *TokenUsage) *TokenUsage {
 	sum.SubagentTokens = AddTokenUsage(aSub, bSub)
 	return sum
 }
+
+// SubtractTokenUsage returns a-b, recursing into subagent usage and clamping
+// every field at zero (a nil operand is treated as zero). Neither input is
+// mutated. Used to rescope a cumulative-since-session-start snapshot (e.g.
+// subagent token usage, which is always re-read from the start of each
+// subagent transcript) down to a delta since a previously captured baseline.
+func SubtractTokenUsage(a, b *TokenUsage) *TokenUsage {
+	if a == nil {
+		return nil
+	}
+	if b == nil {
+		b = &TokenUsage{}
+	}
+	diff := &TokenUsage{
+		InputTokens:         clampSubtract(a.InputTokens, b.InputTokens),
+		CacheCreationTokens: clampSubtract(a.CacheCreationTokens, b.CacheCreationTokens),
+		CacheReadTokens:     clampSubtract(a.CacheReadTokens, b.CacheReadTokens),
+		OutputTokens:        clampSubtract(a.OutputTokens, b.OutputTokens),
+		APICallCount:        clampSubtract(a.APICallCount, b.APICallCount),
+	}
+	diff.SubagentTokens = SubtractTokenUsage(a.SubagentTokens, b.SubagentTokens)
+	return diff
+}
+
+// clampSubtract returns a-b, floored at zero so a stale or racy baseline
+// never produces a negative delta.
+func clampSubtract(a, b int) int {
+	if a < b {
+		return 0
+	}
+	return a - b
+}

cmd/entire/cli/agentimport/agentimport.go

@@ -44,7 +44,16 @@ type Turn struct {
 	UUID               string
 	Prompt, Model      string
 	CreatedAt          time.Time
-	Tokens             *types.TokenUsage
+	// Tokens is this turn's token usage. Every field is a per-turn delta:
+	// main-agent fields are scoped to the turn's [LineStart, LineEnd) slice by
+	// the token helpers, and SubagentTokens is rescoped from the cumulative
+	// snapshot those helpers return to a per-turn increment by
+	// rescopeSubagentTokensToDeltas (see linesplit.go). That invariant lets
+	// callers sum turns freely: writeSessionState sums them for the session
+	// total and each imported checkpoint stores its own turn's delta, so a
+	// subagent's tokens are counted exactly once rather than re-added on every
+	// turn after it is discovered.
+	Tokens *types.TokenUsage
 }
 
 // Importer is the per-agent seam: it locates an agent's transcripts for a repo
@@ -215,6 +224,13 @@ func writeSessionState(ctx context.Context, imp Importer, sf SessionFile, turns
 		if turn.Model != "" {
 			model = turn.Model
 		}
+		// turn.Tokens holds per-turn deltas for every field, including
+		// SubagentTokens (rescoped from a cumulative snapshot in
+		// rescopeSubagentTokensToDeltas — see the Turn.Tokens doc). Summing
+		// them therefore yields the correct session total: main-agent fields
+		// add up, and the subagent deltas sum back to the final cumulative
+		// subagent snapshot exactly once instead of being multiplied by the
+		// number of turns after each subagent was first discovered.
 		tokens = types.AddTokenUsage(tokens, turn.Tokens)
 	}
 	if started.IsZero() {

cmd/entire/cli/agentimport/claude.go

@@ -32,10 +32,12 @@ func (claudeImporter) Discover(repoRoot, overridePath string, now time.Time, ses
 	return discoverSessionFiles(dir, now, sessionFilter, jsonlSessionResolver(".jsonl", identitySessionID))
 }
 
-// SplitTurns produces one Turn per user-prompt line. Token usage for each turn
-// is computed on the slice [LineStart, LineEnd) so turns don't double-count
-// later turns. tool_result lines (Type == "user" but no text content) do not
-// start a turn.
+// SplitTurns produces one Turn per user-prompt line. Main-agent token usage for
+// each turn is computed on the slice [LineStart, LineEnd) so turns don't
+// double-count later turns; subagent token usage is discovered from the full
+// prefix and rescoped to a per-turn delta by splitLineTurns (see
+// rescopeSubagentTokensToDeltas). tool_result lines (Type == "user" but no text
+// content) do not start a turn.
 func (claudeImporter) SplitTurns(sf SessionFile, full []byte) ([]Turn, error) {
 	subagentsDir := filepath.Join(filepath.Dir(sf.Path), sf.SessionID, "subagents")
 	ag := &claudecode.ClaudeCodeAgent{}