From 9d7aebcb2210cb344d232f89d893ff40e7b7492f Mon Sep 17 00:00:00 2001 From: fullsend-code <278716306+fullsend-ai-coder[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:23:39 +0000 Subject: [PATCH] fix(#6905): capture token telemetry on cancelled Claude runs The Claude runtime only populated RunMetrics token fields from the terminal ResultEvent, which is never emitted when the process is killed (SIGTERM on cancellation). This left InputTokens, OutputTokens, and cache token fields at zero for every cancelled run, corrupting budget tracking and usage analytics. The fix has three parts: 1. Cumulative token tracking in parseClaudeStream: per-message token counters are folded into cumulative counters on each message_start, so TokensEvent now carries run-wide totals instead of single-message values. 2. TokensEvent handler in the metrics handler (claude.go and progressParser): incrementally updates RunMetrics from the cumulative TokensEvent. ResultEvent still overwrites with authoritative totals on successful runs. 3. Deferred final TokensEvent at EOF: when the stream ends without a ResultEvent (cancelled run), a deferred function emits one last TokensEvent with the cumulative snapshot, ensuring the metrics handler captures all available data. Cost (TotalCostUSD) remains zero for cancelled runs because the Claude stream does not expose per-API-call cost; a separate strategy is needed (noted in #6905). Closes #6905 --- internal/runtime/claude.go | 9 ++ internal/runtime/claude_progress.go | 56 ++++++- internal/runtime/claude_progress_test.go | 193 ++++++++++++++++++++++- 3 files changed, 251 insertions(+), 7 deletions(-) diff --git a/internal/runtime/claude.go b/internal/runtime/claude.go index c5caafa204..868427f507 100644 --- a/internal/runtime/claude.go +++ b/internal/runtime/claude.go @@ -145,7 +145,16 @@ func (ClaudeRuntime) Run(ctx context.Context, params RunParams, printer *ui.Prin if metrics.Model == "" { metrics.Model = e.Model } + case TokensEvent: + // Capture cumulative token usage from the stream so cancelled + // runs (no ResultEvent) retain non-zero telemetry (#6905). + metrics.InputTokens = e.InputTokens + metrics.OutputTokens = e.OutputTokens + metrics.CacheReadInputTokens = e.CacheRead + metrics.CacheCreationInputTokens = e.CacheWrite case ResultEvent: + // Authoritative totals from the terminal result event overwrite + // the incremental snapshot. metrics.NumTurns = e.NumTurns metrics.TotalCostUSD = e.TotalCostUSD metrics.InputTokens = e.InputTokens diff --git a/internal/runtime/claude_progress.go b/internal/runtime/claude_progress.go index bfab70e6ff..72142d01bf 100644 --- a/internal/runtime/claude_progress.go +++ b/internal/runtime/claude_progress.go @@ -115,14 +115,44 @@ func parseClaudeStream(r io.Reader, onEvent func(AgentEvent)) error { seenStreamEvent bool currentToolName string toolInputJSON strings.Builder - // token tracking for throttled TokensEvent + // per-message token tracking for throttled TokensEvent totalInput int totalOutput int totalCacheRead int totalCacheWrite int lastEmittedTotal int + // cumulative token tracking across all API calls so cancelled + // runs retain the best-effort total instead of zeros (#6905). + cumulativeInput int + cumulativeOutput int + cumulativeCacheRead int + cumulativeCacheWrite int + seenResult bool ) + // Emit a final cumulative TokensEvent when the stream ends without + // a ResultEvent (cancelled/killed run) and the cumulative total + // exceeds the last emitted snapshot. The deferred call covers + // every exit path: EOF, read error, and normal return. + defer func() { + if seenResult { + return + } + finalInput := cumulativeInput + totalInput + finalOutput := cumulativeOutput + totalOutput + finalCacheRead := cumulativeCacheRead + totalCacheRead + finalCacheWrite := cumulativeCacheWrite + totalCacheWrite + total := finalInput + finalOutput + finalCacheRead + finalCacheWrite + if total > lastEmittedTotal { + onEvent(TokensEvent{ + InputTokens: finalInput, + OutputTokens: finalOutput, + CacheRead: finalCacheRead, + CacheWrite: finalCacheWrite, + }) + } + }() + for { line, isPrefix, err := br.ReadLine() if err == io.EOF { @@ -236,6 +266,13 @@ func parseClaudeStream(r io.Reader, onEvent func(AgentEvent)) error { } `json:"message"` } if err := json.Unmarshal(wrapper.Event, &msg); err == nil { + // Fold the previous message's final tokens into the + // cumulative counters before resetting for the new message. + cumulativeInput += totalInput + cumulativeOutput += totalOutput + cumulativeCacheRead += totalCacheRead + cumulativeCacheWrite += totalCacheWrite + totalInput = msg.Message.Usage.InputTokens totalOutput = 0 totalCacheRead = msg.Message.Usage.CacheReadInputTokens @@ -250,20 +287,22 @@ func parseClaudeStream(r io.Reader, onEvent func(AgentEvent)) error { } if err := json.Unmarshal(wrapper.Event, &md); err == nil && md.Usage.OutputTokens > 0 { totalOutput = md.Usage.OutputTokens - total := totalInput + totalOutput + totalCacheRead + totalCacheWrite + total := cumulativeInput + totalInput + cumulativeOutput + totalOutput + + cumulativeCacheRead + totalCacheRead + cumulativeCacheWrite + totalCacheWrite if total-lastEmittedTotal >= tokenThreshold { lastEmittedTotal = total onEvent(TokensEvent{ - InputTokens: totalInput, - OutputTokens: totalOutput, - CacheRead: totalCacheRead, - CacheWrite: totalCacheWrite, + InputTokens: cumulativeInput + totalInput, + OutputTokens: cumulativeOutput + totalOutput, + CacheRead: cumulativeCacheRead + totalCacheRead, + CacheWrite: cumulativeCacheWrite + totalCacheWrite, }) } } } case "result": + seenResult = true var re resultEvent if err := json.Unmarshal(line, &re); err != nil { continue @@ -340,6 +379,11 @@ func progressParser(r io.Reader, printer *ui.Printer, metrics *RunMetrics) error if metrics.Model == "" { metrics.Model = e.Model } + case TokensEvent: + metrics.InputTokens = e.InputTokens + metrics.OutputTokens = e.OutputTokens + metrics.CacheReadInputTokens = e.CacheRead + metrics.CacheCreationInputTokens = e.CacheWrite case ResultEvent: metrics.NumTurns = e.NumTurns metrics.TotalCostUSD = e.TotalCostUSD diff --git a/internal/runtime/claude_progress_test.go b/internal/runtime/claude_progress_test.go index ce2521b46b..3aed0969ae 100644 --- a/internal/runtime/claude_progress_test.go +++ b/internal/runtime/claude_progress_test.go @@ -946,6 +946,9 @@ func TestParseClaudeStreamTokensEventThrottled(t *testing.T) { lines := []string{ `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":4000}}}}`, `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":200}}}`, + // ResultEvent prevents the deferred EOF emission so this test + // isolates the in-stream throttle behavior. + `{"type":"result","num_turns":1,"total_cost_usd":0.01,"usage":{"input_tokens":4000,"output_tokens":200}}`, } events := collectEvents(t, strings.Join(lines, "\n")) @@ -955,7 +958,7 @@ func TestParseClaudeStreamTokensEventThrottled(t *testing.T) { tokens = append(tokens, te) } } - // Total = 4200, below 5k threshold + // Total = 4200, below 5k threshold — no in-stream TokensEvent emitted. if len(tokens) != 0 { t.Fatalf("expected 0 tokens events (below threshold), got %d", len(tokens)) } @@ -974,3 +977,191 @@ func TestParseClaudeStreamAssistantSuppressedWhenStreaming(t *testing.T) { } } } + +// TestProgressParserCancelledRunCapturesTokens verifies that cancelled runs +// (stream with TokensEvents but no ResultEvent) produce non-zero token +// metrics. This is the core regression test for #6905. +func TestProgressParserCancelledRunCapturesTokens(t *testing.T) { + lines := []string{ + `{"type":"system","subtype":"init","model":"claude-opus-4-6"}`, + // Two tool calls (already works before fix) + `{"type":"stream_event","event":{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","name":"Read"}}}`, + `{"type":"stream_event","event":{"type":"content_block_stop","index":0}}`, + `{"type":"stream_event","event":{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","name":"Bash"}}}`, + `{"type":"stream_event","event":{"type":"content_block_stop","index":1}}`, + // One API call with token usage that crosses the threshold + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":4000,"cache_read_input_tokens":500,"cache_creation_input_tokens":200}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":2000}}}`, + // No result event — simulates SIGTERM cancellation. + } + + input := strings.NewReader(strings.Join(lines, "\n")) + var buf bytes.Buffer + printer := ui.New(&buf) + metrics := &RunMetrics{} + + if err := progressParser(input, printer, metrics); err != nil { + t.Fatalf("progressParser returned error: %v", err) + } + + // Tool calls are already tracked incrementally (not affected by this bug). + if got := metrics.ToolCalls.Load(); got != 2 { + t.Errorf("expected 2 tool calls, got %d", got) + } + + // Token metrics must be non-zero — this was the #6905 regression. + if metrics.InputTokens == 0 { + t.Error("expected non-zero InputTokens on cancelled run") + } + if metrics.OutputTokens == 0 { + t.Error("expected non-zero OutputTokens on cancelled run") + } + if metrics.InputTokens != 4000 { + t.Errorf("expected 4000 input tokens, got %d", metrics.InputTokens) + } + if metrics.OutputTokens != 2000 { + t.Errorf("expected 2000 output tokens, got %d", metrics.OutputTokens) + } + if metrics.CacheReadInputTokens != 500 { + t.Errorf("expected 500 cache read tokens, got %d", metrics.CacheReadInputTokens) + } + if metrics.CacheCreationInputTokens != 200 { + t.Errorf("expected 200 cache creation tokens, got %d", metrics.CacheCreationInputTokens) + } +} + +// TestProgressParserResultOverwritesIncrementalTokens verifies that when a +// ResultEvent is present (successful run), its authoritative totals overwrite +// the incremental TokensEvent snapshot. +func TestProgressParserResultOverwritesIncrementalTokens(t *testing.T) { + lines := []string{ + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":4000,"cache_read_input_tokens":500,"cache_creation_input_tokens":200}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":2000}}}`, + // ResultEvent with authoritative totals (different from stream values). + `{"type":"result","num_turns":8,"total_cost_usd":0.42,"usage":{"input_tokens":12000,"output_tokens":3400,"cache_creation_input_tokens":8000,"cache_read_input_tokens":5000}}`, + } + + input := strings.NewReader(strings.Join(lines, "\n")) + var buf bytes.Buffer + printer := ui.New(&buf) + metrics := &RunMetrics{} + + if err := progressParser(input, printer, metrics); err != nil { + t.Fatalf("progressParser returned error: %v", err) + } + + // ResultEvent values must win over the incremental TokensEvent snapshot. + if metrics.InputTokens != 12000 { + t.Errorf("expected 12000 input tokens from ResultEvent, got %d", metrics.InputTokens) + } + if metrics.OutputTokens != 3400 { + t.Errorf("expected 3400 output tokens from ResultEvent, got %d", metrics.OutputTokens) + } + if metrics.TotalCostUSD != 0.42 { + t.Errorf("expected cost 0.42 from ResultEvent, got %f", metrics.TotalCostUSD) + } + if metrics.NumTurns != 8 { + t.Errorf("expected 8 turns from ResultEvent, got %d", metrics.NumTurns) + } +} + +// TestParseClaudeStreamCumulativeTokensAcrossMessages verifies that +// TokensEvent carries cumulative token counts across multiple API calls, +// not just the current message's counts. +func TestParseClaudeStreamCumulativeTokensAcrossMessages(t *testing.T) { + lines := []string{ + // First API call: input=3000, output=2000, cache_r=500, cache_w=200 + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":3000,"cache_read_input_tokens":500,"cache_creation_input_tokens":200}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":2000}}}`, + // Second API call: input=4000, output=3000, cache_r=600, cache_w=100 + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":4000,"cache_read_input_tokens":600,"cache_creation_input_tokens":100}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":3000}}}`, + // No result event — cancelled. + } + + events := collectEvents(t, strings.Join(lines, "\n")) + + var lastTokens TokensEvent + for _, e := range events { + if te, ok := e.(TokensEvent); ok { + lastTokens = te + } + } + + // The final TokensEvent (emitted at EOF for cancelled runs) should + // contain cumulative totals: input=3000+4000=7000, output=2000+3000=5000, + // cache_r=500+600=1100, cache_w=200+100=300. + if lastTokens.InputTokens != 7000 { + t.Errorf("expected cumulative input 7000, got %d", lastTokens.InputTokens) + } + if lastTokens.OutputTokens != 5000 { + t.Errorf("expected cumulative output 5000, got %d", lastTokens.OutputTokens) + } + if lastTokens.CacheRead != 1100 { + t.Errorf("expected cumulative cache read 1100, got %d", lastTokens.CacheRead) + } + if lastTokens.CacheWrite != 300 { + t.Errorf("expected cumulative cache write 300, got %d", lastTokens.CacheWrite) + } +} + +// TestParseClaudeStreamNoFinalTokensEventAfterResult verifies that when +// a ResultEvent is present, no extra TokensEvent is emitted at EOF — the +// authoritative ResultEvent data must not be overwritten. +func TestParseClaudeStreamNoFinalTokensEventAfterResult(t *testing.T) { + lines := []string{ + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":4000,"cache_read_input_tokens":500,"cache_creation_input_tokens":200}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":2000}}}`, + `{"type":"result","num_turns":8,"total_cost_usd":0.42,"usage":{"input_tokens":12000,"output_tokens":3400,"cache_creation_input_tokens":8000,"cache_read_input_tokens":5000}}`, + } + + events := collectEvents(t, strings.Join(lines, "\n")) + + // Find the ResultEvent and check no TokensEvent comes after it. + seenResult := false + for _, e := range events { + switch e.(type) { + case ResultEvent: + seenResult = true + case TokensEvent: + if seenResult { + t.Error("TokensEvent must not be emitted after ResultEvent") + } + } + } + if !seenResult { + t.Error("expected ResultEvent in the stream") + } +} + +// TestParseClaudeStreamFinalTokensEventOnCancel verifies that a deferred +// TokensEvent is emitted at EOF when the stream ends without a ResultEvent, +// even when the per-message total is below the throttle threshold. +func TestParseClaudeStreamFinalTokensEventOnCancel(t *testing.T) { + // Total per-message tokens: 2000+500+100+50 = 2650, below the 5000 threshold. + // Without the deferred emission, no TokensEvent would be emitted at all. + lines := []string{ + `{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":2000,"cache_read_input_tokens":100,"cache_creation_input_tokens":50}}}}`, + `{"type":"stream_event","event":{"type":"message_delta","usage":{"output_tokens":500}}}`, + // No result event — cancelled. + } + + events := collectEvents(t, strings.Join(lines, "\n")) + + var tokens []TokensEvent + for _, e := range events { + if te, ok := e.(TokensEvent); ok { + tokens = append(tokens, te) + } + } + + if len(tokens) != 1 { + t.Fatalf("expected 1 deferred TokensEvent at EOF, got %d", len(tokens)) + } + if tokens[0].InputTokens != 2000 { + t.Errorf("expected 2000 input tokens, got %d", tokens[0].InputTokens) + } + if tokens[0].OutputTokens != 500 { + t.Errorf("expected 500 output tokens, got %d", tokens[0].OutputTokens) + } +}