fix(agentimport): stop summing cumulative subagent tokens across imported turns · Entire
fix(agentimport): stop summing cumulative subagent tokens across imported turns
440c0ad→main·
The subagent-aware importers (Claude Code, Factory AI Droid) compute each turn's token usage via CalculateTotalTokenUsage over the full transcript prefix [0,end): main-agent fields are scoped to the turn's [start,end) slice\ (genuine per-turn deltas), but subagent IDs are discovered from the whole prefix (#329) and each agent-
Both import consumers then summed that snapshot: writeSessionState folds turns with AddTokenUsage for the session total, and every imported checkpoint stores its turn's TokenUsage on entire/checkpoints/v1 (downstream sums those). The result multiplied a subagent's tokens by the number of turns after it was first discovered, inflating both the imported session total and per-checkpoint usage.
Rescope each turn's SubagentTokens from the cumulative snapshot to a per-turn delta (this turn's cumulative minus the previous turn's) in splitLineTurns via types.SubtractTokenUsage. Each imported checkpoint then carries only the subagent usage attributable to its turn, and summing the deltas reconstructs the final cumulative total exactly once. This mirrors the live path, which keeps the latest cumulative snapshot in state.TokenUsage and rescopes each checkpoint window to "cumulative minus a captured baseline" (accumulateTokenUsage / SessionState.SubagentTokensBaseline); here the baseline for turn k is turn k-1's cumulative. The cumulative is monotonic non-decreasing, so the deltas sum back to the final snapshot exactly once. Turns with no discovered subagent keep a nil SubagentTokens, so this is a no-op for the non-subagent-aware importers.
Add regression tests for both importers: a subagent spawned in an early turn plus later prompt turns, asserting the subagent total is counted exactly once in both the summed per-turn/per-checkpoint usage and the imported session total, while main-agent totals still sum. Reverting the rescope makes both assertions fail with a 3x multiplication.
Changes
5
cmd/entire/cli/agentimport
Magentimport.go+17/-1
Mclaude.go+6/-4
Mfactory.go+4/-2
Mlinesplit.go+53/-1
Asubagent_tokens_test.go+190
43 unmodified lines
44
45
46
47
47
48
49
50
51
52
53
54
55
56
57
58
59
164 unmodified lines
224
225
226
227
228
229
230
231
232
233
234
235
236
43 unmodified lines
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
164 unmodified lines
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() {
Mcmd/entire/cli/agentimport/agentimport.go+17/-1
31 unmodified lines
32
33
34
35
36
37
38
35
36
37
38
39
40
41
42
43
31 unmodified lines
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{}
Mcmd/entire/cli/agentimport/claude.go+6/-4
33 unmodified lines
34
35
36
37
38
37
38
39
40
41
42
43
33 unmodified lines
}
// SplitTurns produces one Turn per user-prompt envelope, bounded by the next.
// Token usage (including spawned subagents) is delegated to the Factory agent;
// the model is read once from the session's adjacent settings file. Droid
// Token usage is delegated to the Factory agent; spawned-subagent usage comes
// back as a cumulative snapshot and is rescoped to a per-turn delta by
// splitLineTurns (see rescopeSubagentTokensToDeltas). The model is read once
// from the session's adjacent settings file. Droid
// envelopes carry no per-message timestamp (the agent stamps events with
// time.Now() at hook time), so every turn falls back to the transcript file's
// modtime — the same fallback the Gemini importer uses.
Mcmd/entire/cli/agentimport/factory.go+4/-2
1
2
3
3
4
5
6
7
8
9
10
45 unmodified lines
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package agentimport
import "time"
import (
"time"
"github.com/entireio/cli/cmd/entire/cli/agent/types"
)
// parseTimestamp parses an RFC3339 timestamp, returning the zero time when the
// string is empty or unparseable. Shared by the importers that read a per-turn
45 unmodified lines
turn.LineStart, turn.LineEnd = start, end
turns = append(turns, *turn)
}
rescopeSubagentTokensToDeltas(turns)
return turns, nil
}
// rescopeSubagentTokensToDeltas converts each turn's SubagentTokens from the
// cumulative-since-session-start snapshot the token helpers return into the
// per-turn increment (this turn's cumulative minus the previous turn's).
//
// The subagent-aware token helpers (claudecode/factoryaidroid
// CalculateTotalTokenUsage) discover spawned agent IDs from the full transcript
// prefix [0,end) — so a subagent spawned before the current turn is still found
// (#329) — and re-read each agent-<id>.jsonl from line 0. That makes a turn's
// SubagentTokens a cumulative snapshot that repeats every already-discovered
// subagent's full total on every later turn, unlike the main-agent fields
// (InputTokens/OutputTokens/...), which are scoped to the turn's own
// [start,end) slice and are genuine per-turn deltas.
//
// Both import consumers sum per-turn token usage: writeSessionState folds the
// turns together with AddTokenUsage for the session total, and every imported
// checkpoint stores its turn's TokenUsage (downstream consumers sum those).
// Summing the cumulative snapshot multiplies a subagent's tokens by the number
// of turns after it was first discovered (trail finding 019f5ea3). Rescoping to
// per-turn deltas fixes both without special-casing either: each checkpoint
// carries only the subagent usage attributable to its turn, and summing the
// deltas reconstructs the final cumulative total exactly once.
//
// This mirrors the live path, which keeps the latest cumulative snapshot in
// state.TokenUsage (accumulateTokenUsage replaces rather than adds
// SubagentTokens) and rescopes each checkpoint window to "cumulative minus a
// captured baseline" via types.SubtractTokenUsage and
// SessionState.SubagentTokensBaseline (see cmd/entire/cli/strategy). Here the
// baseline for turn k is turn k-1's cumulative snapshot. The cumulative is
// monotonic non-decreasing across turns (discovered-agent set only grows and
// each subagent file total is fixed), so the clamped subtraction is exact and
// the deltas sum back to the final snapshot.
//
// Turns without a discovered subagent have a nil SubagentTokens and are left
// untouched, so this is a no-op for the non-subagent-aware importers that route
// through splitLineTurns (cursor/pi/codex/copilot).
func rescopeSubagentTokensToDeltas(turns []Turn) {
var prevCumulative *types.TokenUsage
for i := range turns {
if turns[i].Tokens == nil {
continue
}
cumulative := turns[i].Tokens.SubagentTokens
turns[i].Tokens.SubagentTokens = types.SubtractTokenUsage(cumulative, prevCumulative)
prevCumulative = cumulative
}
}
Mcmd/entire/cli/agentimport/linesplit.go+53/-1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
package agentimport
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"github.com/entireio/cli/cmd/entire/cli/agent/types"
)
// These regression tests pin the trail-817 fix: the subagent-aware importers
// (Claude Code, Factory AI Droid) get their per-turn SubagentTokens as a
// cumulative-since-session-start snapshot (agent IDs are discovered from the
// full transcript prefix and each subagent transcript is re-read from line 0),
// so a subagent spawned in an early turn repeats its full total on every later
// turn. Both import consumers sum per-turn usage — writeSessionState for the
// session total and the per-imported-checkpoint TokenUsage that downstream sums
// — so before the fix a subagent's tokens were multiplied by the number of
// turns after it was first discovered. rescopeSubagentTokensToDeltas
// (linesplit.go) rescopes those snapshots to per-turn deltas so the total is
// counted exactly once. Reverting that call makes both assertions below fail.
const (
// The single spawned subagent's on-disk transcript totals, asserted to be
// counted exactly once across a 3-turn session.
wantSubagentInput = 50
wantSubagentOutput = 25
wantSubagentCalls = 1
// Main-agent totals summed across the three turns (per-slice deltas), kept
// intact by the fix.
wantMainInput = 600 // 100 + 200 + 300
wantMainOutput = 300 // 50 + 100 + 150
wantMainCalls = 3
)
// sumTurnSubagentTokens sums each turn's SubagentTokens the way both the session
// total and any downstream sum of per-checkpoint TokenUsage would. With the fix
// the turns hold per-turn deltas, so this reconstructs the subagent total once;
// without it each turn holds the cumulative snapshot and this multiplies.
func sumTurnSubagentTokens(turns []Turn) types.TokenUsage {
var sum types.TokenUsage
for _, tr := range turns {
if tr.Tokens == nil || tr.Tokens.SubagentTokens == nil {
continue
}
s := tr.Tokens.SubagentTokens
sum.InputTokens += s.InputTokens
sum.CacheCreationTokens += s.CacheCreationTokens
sum.CacheReadTokens += s.CacheReadTokens
sum.OutputTokens += s.OutputTokens
sum.APICallCount += s.APICallCount
}
return sum
}
func assertSubagentCountedOnce(t *testing.T, label string, got *types.TokenUsage) {
t.Helper()
if got == nil {
t.Fatalf("%s: SubagentTokens is nil, want input=%d output=%d calls=%d",
label, wantSubagentInput, wantSubagentOutput, wantSubagentCalls)
}
if got.InputTokens != wantSubagentInput || got.OutputTokens != wantSubagentOutput ||
got.APICallCount != wantSubagentCalls {
t.Errorf("%s: subagent tokens counted more than once: got input=%d output=%d calls=%d, "+
"want input=%d output=%d calls=%d (cumulative snapshot summed across turns)",
label, got.InputTokens, got.OutputTokens, got.APICallCount,
wantSubagentInput, wantSubagentOutput, wantSubagentCalls)
}
}
func writeSubagentTranscript(t *testing.T, sf SessionFile, agentID, line string) {
t.Helper()
subagentsDir := filepath.Join(filepath.Dir(sf.Path), sf.SessionID, "subagents")
if err := os.MkdirAll(subagentsDir, 0o755); err != nil {
t.Fatalf("mkdir subagents dir: %v", err)
}
agentPath := filepath.Join(subagentsDir, "agent-"+agentID+".jsonl")
if err := os.WriteFile(agentPath, []byte(line+"\n"), 0o600); err != nil {
t.Fatalf("write subagent transcript: %v", err)
}
}
// TestImport_ClaudeSubagentTokensCountedOnceAcrossTurns builds a Claude session
// where a subagent is spawned in the first turn and two more user-prompt turns
// follow, then asserts the subagent's tokens are counted exactly once both in
// the summed per-turn/per-checkpoint usage and in the imported session total.
func TestImport_ClaudeSubagentTokensCountedOnceAcrossTurns(t *testing.T) {
importRepo(t) // chdir into a repo for session-state storage; no t.Parallel (t.Chdir)
dir := t.TempDir()
sf := SessionFile{Path: filepath.Join(dir, "s.jsonl"), SessionID: "s"}
// Turn 1 spawns subagent "subX" (Task tool_use + tool_result carrying the
// agentId), then three user-prompt turns each with their own assistant
// usage. The tool_result line is type "user" but has no text, so it does
// not start a turn.
full := []byte(strings.Join([]string{
`{"type":"user","uuid":"u1","message":{"role":"user","content":"first"}}`,
`{"type":"assistant","uuid":"a0","message":{"content":[{"type":"tool_use","id":"toolu_task1","name":"Task","input":{"prompt":"go"}}]}}`,
`{"type":"user","uuid":"r1","message":{"content":[{"type":"tool_result","tool_use_id":"toolu_task1","content":"agentId: subX"}]}}`,
`{"type":"assistant","uuid":"a1","message":{"id":"m1","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":100,"output_tokens":50}}}`,
`{"type":"user","uuid":"u2","message":{"role":"user","content":"second"}}`,
`{"type":"assistant","uuid":"a2","message":{"id":"m2","content":[{"type":"text","text":"ok2"}],"usage":{"input_tokens":200,"output_tokens":100}}}`,
`{"type":"user","uuid":"u3","message":{"role":"user","content":"third"}}`,
`{"type":"assistant","uuid":"a3","message":{"id":"m3","content":[{"type":"text","text":"ok3"}],"usage":{"input_tokens":300,"output_tokens":150}}}`,
}, "\n") + "\n")
if err := os.WriteFile(sf.Path, full, 0o600); err != nil {
t.Fatalf("write transcript: %v", err)
}
writeSubagentTranscript(t, sf, "subX",
`{"type":"assistant","uuid":"sa1","message":{"id":"sm1","content":[{"type":"text","text":"sub"}],"usage":{"input_tokens":50,"output_tokens":25}}}`)
turns, err := claudeImporter{}.SplitTurns(sf, full)
if err != nil {
t.Fatalf("SplitTurns: %v", err)
}
assertSubagentTurns(t, claudeImporter{}, sf, turns)
}
// TestImport_FactorySubagentTokensCountedOnceAcrossTurns is the Factory AI Droid
// analogue: Droid envelopes, subagent spawned in the first turn, three prompt
// turns, subagent tokens counted exactly once.
func TestImport_FactorySubagentTokensCountedOnceAcrossTurns(t *testing.T) {
importRepo(t)
dir := t.TempDir()
sf := SessionFile{Path: filepath.Join(dir, "s.jsonl"), SessionID: "s"}
full := []byte(strings.Join([]string{
`{"type":"message","id":"u1","message":{"role":"user","content":"first"}}`,
`{"type":"message","id":"a0","message":{"role":"assistant","content":[{"type":"tool_use","id":"toolu_task1","name":"Task","input":{"prompt":"go"}}]}}`,
`{"type":"message","id":"r1","message":{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_task1","content":"agentId: subX"}]}}`,
`{"type":"message","id":"a1","message":{"role":"assistant","id":"m1","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":100,"output_tokens":50}}}`,
`{"type":"message","id":"u2","message":{"role":"user","content":"second"}}`,
`{"type":"message","id":"a2","message":{"role":"assistant","id":"m2","content":[{"type":"text","text":"ok2"}],"usage":{"input_tokens":200,"output_tokens":100}}}`,
`{"type":"message","id":"u3","message":{"role":"user","content":"third"}}`,
`{"type":"message","id":"a3","message":{"role":"assistant","id":"m3","content":[{"type":"text","text":"ok3"}],"usage":{"input_tokens":300,"output_tokens":150}}}`,
}, "\n") + "\n")
if err := os.WriteFile(sf.Path, full, 0o600); err != nil {
t.Fatalf("write transcript: %v", err)
}
writeSubagentTranscript(t, sf, "subX",
`{"type":"message","id":"se1","message":{"role":"assistant","id":"sm1","content":[{"type":"text","text":"sub"}],"usage":{"input_tokens":50,"output_tokens":25}}}`)
turns, err := factoryImporter{}.SplitTurns(sf, full)
if err != nil {
t.Fatalf("SplitTurns: %v", err)
}
assertSubagentTurns(t, factoryImporter{}, sf, turns)
}
// assertSubagentTurns runs the shared assertions for a 3-turn session with one
// spawned subagent: the summed per-turn (== per-checkpoint) subagent tokens and
// the imported session total each count the subagent exactly once, while the
// main-agent totals still sum.
func assertSubagentTurns(t *testing.T, imp Importer, sf SessionFile, turns []Turn) {
t.Helper()
ctx := context.Background()
if len(turns) != 3 {
t.Fatalf("want 3 turns, got %d", len(turns))
}
// Per-checkpoint proof: summing each turn's stored TokenUsage.SubagentTokens
// (which is exactly what writeTurn persists per imported checkpoint) must
// reconstruct the subagent total once, not 3x.
perCheckpoint := sumTurnSubagentTokens(turns)
assertSubagentCountedOnce(t, "sum of per-turn SubagentTokens", &perCheckpoint)
// Session-total proof: the imported session.State.TokenUsage folds the
// turns via writeSessionState the same way production Run does.
if err := writeSessionState(ctx, imp, sf, turns); err != nil {
t.Fatalf("writeSessionState: %v", err)
}
st := loadState(t, sf.SessionID)
if st == nil || st.TokenUsage == nil {
t.Fatalf("no imported session token usage written: %+v", st)
}
assertSubagentCountedOnce(t, "session total SubagentTokens", st.TokenUsage.SubagentTokens)
// The main-agent fields are genuine per-slice deltas and must still sum.
if st.TokenUsage.InputTokens != wantMainInput || st.TokenUsage.OutputTokens != wantMainOutput ||
st.TokenUsage.APICallCount != wantMainCalls {
t.Errorf("main-agent totals = input=%d output=%d calls=%d, want input=%d output=%d calls=%d",
st.TokenUsage.InputTokens, st.TokenUsage.OutputTokens, st.TokenUsage.APICallCount,
wantMainInput, wantMainOutput, wantMainCalls)
}
}