From 3364a942c87981071e0a324fbcf8160783e5919b Mon Sep 17 00:00:00 2001 From: Kristof Berger Date: Tue, 14 Jul 2026 16:40:56 +0200 Subject: [PATCH] fix(transcript): sum usage.iterations to count fallback attempts 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. --- internal/transcript/transcript.go | 50 ++++++++++++----- internal/transcript/transcript_test.go | 75 ++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 13 deletions(-) diff --git a/internal/transcript/transcript.go b/internal/transcript/transcript.go index eb2d62e..ca391ca 100644 --- a/internal/transcript/transcript.go +++ b/internal/transcript/transcript.go @@ -35,10 +35,30 @@ type messageEnvelope struct { } type usageData struct { - InputTokens int64 `json:"input_tokens"` - OutputTokens int64 `json:"output_tokens"` - CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` - CacheReadInputTokens int64 `json:"cache_read_input_tokens"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` + CacheReadInputTokens int64 `json:"cache_read_input_tokens"` + Iterations []usageData `json:"iterations"` +} + +// effective returns the token counts to attribute to this API call. When a +// request is retried on a fallback model, the top-level fields mirror only the +// final iteration while usage.iterations lists every billed attempt — in that +// case the iterations are summed. With zero or one iteration the top-level +// fields already hold the full picture and are returned unchanged. +func (u *usageData) effective() usageData { + if len(u.Iterations) <= 1 { + return *u + } + var sum usageData + for _, it := range u.Iterations { + sum.InputTokens += it.InputTokens + sum.OutputTokens += it.OutputTokens + sum.CacheCreationInputTokens += it.CacheCreationInputTokens + sum.CacheReadInputTokens += it.CacheReadInputTokens + } + return sum } // contentType is used to peek at a content block's type field without fully @@ -52,7 +72,9 @@ type contentType struct { // message). Returns nil when no usage data is found. // // Streaming duplicates (same requestId across multiple transcript entries) are -// deduplicated so usage is counted only once per API call. +// deduplicated so usage is counted only once per API call. When a call was +// retried on a fallback model, all billed iterations are summed (see +// usageData.effective). func ReadTurnUsage(transcriptPath string) (*Usage, error) { f, err := os.Open(transcriptPath) if err != nil { @@ -95,20 +117,22 @@ func ReadTurnUsage(transcriptPath string) (*Usage, error) { if entry.RequestID != "" { perReq[entry.RequestID] = u } else { - noReqUsage.InputTokens += u.InputTokens - noReqUsage.OutputTokens += u.OutputTokens - noReqUsage.CacheCreationInputTokens += u.CacheCreationInputTokens - noReqUsage.CacheReadInputTokens += u.CacheReadInputTokens + eff := u.effective() + noReqUsage.InputTokens += eff.InputTokens + noReqUsage.OutputTokens += eff.OutputTokens + noReqUsage.CacheCreationInputTokens += eff.CacheCreationInputTokens + noReqUsage.CacheReadInputTokens += eff.CacheReadInputTokens } } // Sum final usage across all API calls in the turn. usage := noReqUsage for _, u := range perReq { - usage.InputTokens += u.InputTokens - usage.OutputTokens += u.OutputTokens - usage.CacheCreationInputTokens += u.CacheCreationInputTokens - usage.CacheReadInputTokens += u.CacheReadInputTokens + eff := u.effective() + usage.InputTokens += eff.InputTokens + usage.OutputTokens += eff.OutputTokens + usage.CacheCreationInputTokens += eff.CacheCreationInputTokens + usage.CacheReadInputTokens += eff.CacheReadInputTokens } if !hasUsage { diff --git a/internal/transcript/transcript_test.go b/internal/transcript/transcript_test.go index 818a1fe..0eaab68 100644 --- a/internal/transcript/transcript_test.go +++ b/internal/transcript/transcript_test.go @@ -75,6 +75,81 @@ func TestReadTurnUsageAggregatesMultipleIterations(t *testing.T) { assert.Equal(t, int64(60), usage.CacheReadInputTokens) } +// fallbackUsage is a two-iteration usage payload modeled on a real fallback +// request (Fable attempt, then Opus fallback): the top-level fields mirror only +// the final iteration. +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}]}` + +func TestReadTurnUsageSumsFallbackIterations(t *testing.T) { + path := filepath.Join(t.TempDir(), "transcript.jsonl") + writeTranscript(t, path, []string{ + `{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`, + `{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":` + fallbackUsage + `}}`, + }) + + usage, err := ReadTurnUsage(path) + require.NoError(t, err) + require.NotNil(t, usage) + + // Both iterations are billed, so all fields are summed across them. + assert.Equal(t, int64(11214), usage.InputTokens) + assert.Equal(t, int64(700), usage.OutputTokens) + assert.Equal(t, int64(150), usage.CacheCreationInputTokens) + assert.Equal(t, int64(3000), usage.CacheReadInputTokens) +} + +func TestReadTurnUsageSingleIterationMatchesTopLevel(t *testing.T) { + path := filepath.Join(t.TempDir(), "transcript.jsonl") + writeTranscript(t, path, []string{ + `{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`, + // The single iteration duplicates the top-level fields, as observed in + // real transcripts; behavior must be identical to no iterations at all. + `{"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}]}}}`, + }) + + usage, err := ReadTurnUsage(path) + require.NoError(t, err) + require.NotNil(t, usage) + + assert.Equal(t, int64(100), usage.InputTokens) + assert.Equal(t, int64(50), usage.OutputTokens) + assert.Equal(t, int64(200), usage.CacheCreationInputTokens) + assert.Equal(t, int64(300), usage.CacheReadInputTokens) +} + +func TestReadTurnUsageDeduplicatesFallbackIterations(t *testing.T) { + path := filepath.Join(t.TempDir(), "transcript.jsonl") + writeTranscript(t, path, []string{ + `{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`, + // Streaming split: both entries repeat the same usage incl. iterations. + `{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"thinking","thinking":"..."}],"usage":` + fallbackUsage + `}}`, + `{"type":"assistant","requestId":"req_001","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"usage":` + fallbackUsage + `}}`, + }) + + usage, err := ReadTurnUsage(path) + require.NoError(t, err) + require.NotNil(t, usage) + + // Iterations are summed once per requestId, not once per entry. + assert.Equal(t, int64(11214), usage.InputTokens) + assert.Equal(t, int64(700), usage.OutputTokens) +} + +func TestReadTurnUsageEmptyIterations(t *testing.T) { + path := filepath.Join(t.TempDir(), "transcript.jsonl") + writeTranscript(t, path, []string{ + `{"type":"user","message":{"role":"user","content":[{"type":"text","text":"hello"}]}}`, + `{"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":[]}}}`, + }) + + usage, err := ReadTurnUsage(path) + require.NoError(t, err) + require.NotNil(t, usage) + + assert.Equal(t, int64(100), usage.InputTokens) + assert.Equal(t, int64(50), usage.OutputTokens) +} + func TestReadTurnUsageIgnoresPreviousTurns(t *testing.T) { path := filepath.Join(t.TempDir(), "transcript.jsonl") writeTranscript(t, path, []string{