Skip to content

Commit 77c9130

Browse files
k-bergerk-berger
authored
fix(transcript): sum usage.iterations to count fallback attempts (#170)
When a request is retried on a fallback model, the transcript's top-level usage only mirrors the final iteration; every billed attempt is listed in usage.iterations. Sum the iterations when there is more than one so token counts (and cost estimates) include the failed attempt as well. Single- or no-iteration usage objects are unaffected.
1 parent a54a1d9 commit 77c9130

2 files changed

Lines changed: 112 additions & 13 deletions

File tree

‎internal/transcript/transcript.go‎

Lines changed: 37 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,30 @@ type messageEnvelope struct {
3535
}
3636

3737
type usageData struct {
38-
InputTokens int64 `json:"input_tokens"`
39-
OutputTokens int64 `json:"output_tokens"`
40-
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
41-
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
38+
InputTokens int64 `json:"input_tokens"`
39+
OutputTokens int64 `json:"output_tokens"`
40+
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
41+
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
42+
Iterations []usageData `json:"iterations"`
43+
}
44+
45+
// effective returns the token counts to attribute to this API call. When a
46+
// request is retried on a fallback model, the top-level fields mirror only the
47+
// final iteration while usage.iterations lists every billed attempt — in that
48+
// case the iterations are summed. With zero or one iteration the top-level
49+
// fields already hold the full picture and are returned unchanged.
50+
func (u *usageData) effective() usageData {
51+
if len(u.Iterations) <= 1 {
52+
return *u
53+
}
54+
var sum usageData
55+
for _, it := range u.Iterations {
56+
sum.InputTokens += it.InputTokens
57+
sum.OutputTokens += it.OutputTokens
58+
sum.CacheCreationInputTokens += it.CacheCreationInputTokens
59+
sum.CacheReadInputTokens += it.CacheReadInputTokens
60+
}
61+
return sum
4262
}
4363

4464
// contentType is used to peek at a content block's type field without fully
@@ -52,7 +72,9 @@ type contentType struct {
5272
// message). Returns nil when no usage data is found.
5373
//
5474
// Streaming duplicates (same requestId across multiple transcript entries) are
55-
// deduplicated so usage is counted only once per API call.
75+
// deduplicated so usage is counted only once per API call. When a call was
76+
// retried on a fallback model, all billed iterations are summed (see
77+
// usageData.effective).
5678
func ReadTurnUsage(transcriptPath string) (*Usage, error) {
5779
f, err := os.Open(transcriptPath)
5880
if err != nil {
@@ -95,20 +117,22 @@ func ReadTurnUsage(transcriptPath string) (*Usage, error) {
95117
if entry.RequestID != "" {
96118
perReq[entry.RequestID] = u
97119
} else {
98-
noReqUsage.InputTokens += u.InputTokens
99-
noReqUsage.OutputTokens += u.OutputTokens
100-
noReqUsage.CacheCreationInputTokens += u.CacheCreationInputTokens
101-
noReqUsage.CacheReadInputTokens += u.CacheReadInputTokens
120+
eff := u.effective()
121+
noReqUsage.InputTokens += eff.InputTokens
122+
noReqUsage.OutputTokens += eff.OutputTokens
123+
noReqUsage.CacheCreationInputTokens += eff.CacheCreationInputTokens
124+
noReqUsage.CacheReadInputTokens += eff.CacheReadInputTokens
102125
}
103126
}
104127

105128
// Sum final usage across all API calls in the turn.
106129
usage := noReqUsage
107130
for _, u := range perReq {
108-
usage.InputTokens += u.InputTokens
109-
usage.OutputTokens += u.OutputTokens
110-
usage.CacheCreationInputTokens += u.CacheCreationInputTokens
111-
usage.CacheReadInputTokens += u.CacheReadInputTokens
131+
eff := u.effective()
132+
usage.InputTokens += eff.InputTokens
133+
usage.OutputTokens += eff.OutputTokens
134+
usage.CacheCreationInputTokens += eff.CacheCreationInputTokens
135+
usage.CacheReadInputTokens += eff.CacheReadInputTokens
112136
}
113137

114138
if !hasUsage {

‎internal/transcript/transcript_test.go‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,81 @@ func TestReadTurnUsageAggregatesMultipleIterations(t *testing.T) {
7575
assert.Equal(t, int64(60), usage.CacheReadInputTokens)
7676
}
7777

78+
// fallbackUsage is a two-iteration usage payload modeled on a real fallback
79+
// request (Fable attempt, then Opus fallback): the top-level fields mirror only
80+
// the final iteration.
81+
const fallbackUsage = `{"input_tokens":5607,"output_tokens":698,"cache_creation_input_tokens":100,"cache_read_input_tokens":2000,"cache_creation":{"ephemeral_1h_input_tokens":100,"ephemeral_5m_input_tokens":0},"iterations":[{"type":"message","model":"claude-fable-5","input_tokens":5607,"output_tokens":2,"cache_creation_input_tokens":50,"cache_read_input_tokens":1000},{"type":"fallback_message","model":"claude-opus-4-8","input_tokens":5607,"output_tokens":698,"cache_creation_input_tokens":100,"cache_read_input_tokens":2000}]}`
82+
83+
func TestReadTurnUsageSumsFallbackIterations(t *testing.T) {
84+
path := filepath.Join(t.TempDir(), "transcript.jsonl")
85+
writeTranscript(t, path, []string{
86+
`{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`,
87+
`{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":` + fallbackUsage + `}}`,
88+
})
89+
90+
usage, err := ReadTurnUsage(path)
91+
require.NoError(t, err)
92+
require.NotNil(t, usage)
93+
94+
// Both iterations are billed, so all fields are summed across them.
95+
assert.Equal(t, int64(11214), usage.InputTokens)
96+
assert.Equal(t, int64(700), usage.OutputTokens)
97+
assert.Equal(t, int64(150), usage.CacheCreationInputTokens)
98+
assert.Equal(t, int64(3000), usage.CacheReadInputTokens)
99+
}
100+
101+
func TestReadTurnUsageSingleIterationMatchesTopLevel(t *testing.T) {
102+
path := filepath.Join(t.TempDir(), "transcript.jsonl")
103+
writeTranscript(t, path, []string{
104+
`{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`,
105+
// The single iteration duplicates the top-level fields, as observed in
106+
// real transcripts; behavior must be identical to no iterations at all.
107+
`{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":{"input_tokens":100,"output_tokens":50,"cache_creation_input_tokens":200,"cache_read_input_tokens":300,"iterations":[{"type":"message","input_tokens":100,"output_tokens":50,"cache_creation_input_tokens":200,"cache_read_input_tokens":300}]}}}`,
108+
})
109+
110+
usage, err := ReadTurnUsage(path)
111+
require.NoError(t, err)
112+
require.NotNil(t, usage)
113+
114+
assert.Equal(t, int64(100), usage.InputTokens)
115+
assert.Equal(t, int64(50), usage.OutputTokens)
116+
assert.Equal(t, int64(200), usage.CacheCreationInputTokens)
117+
assert.Equal(t, int64(300), usage.CacheReadInputTokens)
118+
}
119+
120+
func TestReadTurnUsageDeduplicatesFallbackIterations(t *testing.T) {
121+
path := filepath.Join(t.TempDir(), "transcript.jsonl")
122+
writeTranscript(t, path, []string{
123+
`{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`,
124+
// Streaming split: both entries repeat the same usage incl. iterations.
125+
`{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"thinking","thinking":"..."}],"usage":` + fallbackUsage + `}}`,
126+
`{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":` + fallbackUsage + `}}`,
127+
})
128+
129+
usage, err := ReadTurnUsage(path)
130+
require.NoError(t, err)
131+
require.NotNil(t, usage)
132+
133+
// Iterations are summed once per requestId, not once per entry.
134+
assert.Equal(t, int64(11214), usage.InputTokens)
135+
assert.Equal(t, int64(700), usage.OutputTokens)
136+
}
137+
138+
func TestReadTurnUsageEmptyIterations(t *testing.T) {
139+
path := filepath.Join(t.TempDir(), "transcript.jsonl")
140+
writeTranscript(t, path, []string{
141+
`{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`,
142+
`{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":{"input_tokens":100,"output_tokens":50,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"iterations":[]}}}`,
143+
})
144+
145+
usage, err := ReadTurnUsage(path)
146+
require.NoError(t, err)
147+
require.NotNil(t, usage)
148+
149+
assert.Equal(t, int64(100), usage.InputTokens)
150+
assert.Equal(t, int64(50), usage.OutputTokens)
151+
}
152+
78153
func TestReadTurnUsageIgnoresPreviousTurns(t *testing.T) {
79154
path := filepath.Join(t.TempDir(), "transcript.jsonl")
80155
writeTranscript(t, path, []string{

0 commit comments

Comments
 (0)