fix: input token overwritten by message delta (#1631)

**Description**
according to Anthropic's API message_start events contain the input
tokens (tokens consumed for the request) and message_delta events
contain incremental output tokens (tokens generated in the response).
The message_start event correctly sets tokenUsage with input tokens
(e.g., input_tokens=15). However when a message_delta event is
processed, it completely overwrites the tokenUsage variable. The
message_delta event typically has input_tokens=0 (since it's about
output deltas), which overwrites the correct input token count.

---------

Signed-off-by: Dan Sun <[email protected]>

Author: yuzisun

internal/translator/anthropic_anthropic.go

@@ -36,6 +36,7 @@ type anthropicToAnthropicTranslator struct {
 	stream                 bool
 	buffered               []byte
 	streamingResponseModel internalapi.ResponseModel
+	streamingTokenUsage    metrics.TokenUsage
 }
 
 // RequestBody implements [AnthropicMessagesTranslator.RequestBody].
@@ -83,11 +84,17 @@ func (a *anthropicToAnthropicTranslator) ResponseBody(_ map[string]string, body
 		if err != nil {
 			return nil, nil, tokenUsage, a.requestModel, fmt.Errorf("failed to read body: %w", err)
 		}
+
+		// If this is a fresh start (no buffered data), reset the streaming token usage
+		if len(a.buffered) == 0 {
+			a.streamingTokenUsage = metrics.TokenUsage{}
+		}
+
 		a.buffered = append(a.buffered, buf...)
-		tokenUsage = a.extractUsageFromBufferEvent(span)
+		a.extractUsageFromBufferEvent(span)
 		// Use stored streaming response model, fallback to request model for non-compliant backends
 		responseModel = cmp.Or(a.streamingResponseModel, a.requestModel)
-		return
+		return nil, nil, a.streamingTokenUsage, responseModel, nil
 	}
 
 	// Parse the Anthropic response to extract token usage.
@@ -110,11 +117,13 @@ func (a *anthropicToAnthropicTranslator) ResponseBody(_ map[string]string, body
 }
 
 // extractUsageFromBufferEvent extracts the token usage from the buffered event.
-// It scans complete lines and returns the latest usage found in this batch.
-func (a *anthropicToAnthropicTranslator) extractUsageFromBufferEvent(s tracing.MessageSpan) (tokenUsage metrics.TokenUsage) {
+// It scans complete lines and accumulates usage from all events in this batch.
+func (a *anthropicToAnthropicTranslator) extractUsageFromBufferEvent(s tracing.MessageSpan) {
 	for {
 		i := bytes.IndexByte(a.buffered, '\n')
 		if i == -1 {
+			// Recalculate total tokens before returning
+			a.updateTotalTokens()
 			return
 		}
 		line := a.buffered[:i]
@@ -133,32 +142,56 @@ func (a *anthropicToAnthropicTranslator) extractUsageFromBufferEvent(s tracing.M
 		switch {
 		case eventUnion.MessageStart != nil:
 			message := eventUnion.MessageStart
-			// Message only valid in message_start events.
+			// Store the response model for future batches
 			if message.Model != "" {
-				// Store the response model for future batches
 				a.streamingResponseModel = message.Model
 			}
-			// Extract usage from message_start event
+			// Extract usage from message_start event - this sets the baseline input tokens
 			if u := message.Usage; u != nil {
-				tokenUsage = metrics.ExtractTokenUsageFromAnthropic(
+				messageStartUsage := metrics.ExtractTokenUsageFromAnthropic(
 					int64(u.InputTokens),
 					int64(u.OutputTokens),
 					int64(u.CacheReadInputTokens),
 					int64(u.CacheCreationInputTokens),
 				)
+				// Override with message_start usage (contains input tokens and initial state)
+				a.streamingTokenUsage.Override(messageStartUsage)
 			}
 		case eventUnion.MessageDelta != nil:
 			u := eventUnion.MessageDelta.Usage
-			tokenUsage = metrics.ExtractTokenUsageFromAnthropic(
-				int64(u.InputTokens),
-				int64(u.OutputTokens),
-				int64(u.CacheReadInputTokens),
-				int64(u.CacheCreationInputTokens),
-			)
+			// message_delta events provide final counts for specific token types
+			// Update output tokens from message_delta (final count)
+			if u.OutputTokens >= 0 {
+				a.streamingTokenUsage.SetOutputTokens(uint32(u.OutputTokens)) //nolint:gosec
+			}
 		}
 	}
 }
 
+// updateTotalTokens recalculates and sets the total token count
+func (a *anthropicToAnthropicTranslator) updateTotalTokens() {
+	inputTokens, inputSet := a.streamingTokenUsage.InputTokens()
+	outputTokens, outputSet := a.streamingTokenUsage.OutputTokens()
+
+	// Initialize missing values to 0 if we have any token data
+	if outputSet && !inputSet {
+		a.streamingTokenUsage.SetInputTokens(0)
+		inputTokens = 0
+		inputSet = true
+	}
+
+	// Set cached tokens to 0 if not set but we have other token data
+	if outputSet {
+		if _, cachedSet := a.streamingTokenUsage.CachedInputTokens(); !cachedSet {
+			a.streamingTokenUsage.SetCachedInputTokens(0)
+		}
+	}
+
+	if inputSet && outputSet {
+		a.streamingTokenUsage.SetTotalTokens(inputTokens + outputTokens)
+	}
+}
+
 // ResponseError implements [AnthropicMessagesTranslator] for Anthropic to AWS Bedrock Anthropic translation.
 func (a *anthropicToAnthropicTranslator) ResponseError(map[string]string, io.Reader) (
 	newHeaders []internalapi.Header,

internal/translator/anthropic_anthropic_test.go

@@ -107,7 +107,7 @@ func TestAnthropicToAnthropic_ResponseBody_streaming(t *testing.T) {
 	// We split the response into two parts to simulate streaming where each part can end in the
 	// middle of an event.
 	const responseHead = `event: message_start
-data: {"type":"message_start","message":{"model":"claude-sonnet-4-5-20250929","id":"msg_01BfvfMsg2gBzwsk6PZRLtDg","type":"message","role":"assistant","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard"}}    }
+data: {"type":"message_start","message":{"model":"claude-sonnet-4-5-20250929","id":"msg_01BfvfMsg2gBzwsk6PZRLtDg","type":"message","role":"assistant","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":1,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":0,"service_tier":"standard"}}    }
 
 event: content_block_start
 data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}      }
@@ -130,7 +130,7 @@ event: content_block_stop
 data: {"type":"content_block_stop","index":0             }
 
 event: message_delta
-data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":1,"output_tokens":16}               }
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":16}               }
 
 event: message_stop
 data: {"type":"message_stop"       }`
@@ -139,7 +139,7 @@ data: {"type":"message_stop"       }`
 	require.NoError(t, err)
 	require.Nil(t, headerMutation)
 	require.Nil(t, bodyMutation)
-	expected := tokenUsageFrom(9, 0, 1, 10)
+	expected := tokenUsageFrom(10, 1, 0, 10)
 	require.Equal(t, expected, tokenUsage)
 	require.Equal(t, "claude-sonnet-4-5-20250929", responseModel)
 

internal/translator/anthropic_gcpanthropic_test.go

@@ -8,6 +8,7 @@ package translator
 import (
 	"bytes"
 	"encoding/json"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -589,3 +590,68 @@ func tokenUsageFrom(in, cachedInput, out, total int32) metrics.TokenUsage {
 	}
 	return usage
 }
+
+func TestAnthropicToGCPAnthropicTranslator_ResponseBody_StreamingFullScenario(t *testing.T) {
+	// Test to reproduce and verify fix for the input_token=0 issue in Anthropic streaming
+	// This test verifies that input_tokens from message_start are preserved when
+	// message_delta doesn't provide input_tokens (real-world scenario)
+
+	translator := NewAnthropicToGCPAnthropicTranslator("v1", "")
+
+	// Simulate request body to set stream=true
+	reqBody := anthropic.MessagesRequest{
+		Stream: true,
+		Model:  "claude-3-sonnet-20240229",
+	}
+	_, _, err := translator.RequestBody([]byte(`{"stream":true}`), &reqBody, false)
+	require.NoError(t, err)
+
+	// Sample streaming response from Anthropic with realistic flow:
+	// 1. message_start provides input_tokens=15
+	// 2. content_block events provide the actual text content
+	// 3. message_delta at the end provides output_tokens=5 but no input_tokens
+	// 4. message_stop ends the stream
+	sseStream := `event: message_start
+data: {"type": "message_start", "message": {"id": "msg_123", "type": "message", "role": "assistant", "content": [], "model": "claude-3-sonnet-20240229", "usage": {"input_tokens": 15, "output_tokens": 0}}}
+
+event: content_block_start
+data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
+
+event: content_block_delta
+data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}
+
+event: content_block_stop
+data: {"type": "content_block_stop", "index": 0}
+
+event: message_delta
+data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}}
+
+event: message_stop
+data: {"type": "message_stop"}
+
+`
+
+	// Process the streaming response
+	reader := strings.NewReader(sseStream)
+	_, _, tokenUsage, _, err := translator.ResponseBody(nil, reader, false, nil)
+	require.NoError(t, err)
+
+	// Verify token usage - this should preserve input_tokens from message_start
+	inputTokens, inputSet := tokenUsage.InputTokens()
+	outputTokens, outputSet := tokenUsage.OutputTokens()
+	totalTokens, totalSet := tokenUsage.TotalTokens()
+	cachedTokens, cachedSet := tokenUsage.CachedInputTokens()
+
+	// Assertions
+	assert.True(t, inputSet, "Input tokens should be set")
+	assert.Equal(t, uint32(15), inputTokens, "Input tokens should be preserved from message_start")
+
+	assert.True(t, outputSet, "Output tokens should be set")
+	assert.Equal(t, uint32(5), outputTokens, "Output tokens should come from message_delta")
+
+	assert.True(t, totalSet, "Total tokens should be calculated")
+	assert.Equal(t, uint32(20), totalTokens, "Total tokens should be input + output")
+
+	assert.True(t, cachedSet, "Cached tokens should be set")
+	assert.Equal(t, uint32(0), cachedTokens, "No cached tokens in this scenario")
+}