fix(translator): merge Anthropic message_delta usage at field level

Signed-off-by: pjdurden <prajjwalchittori1@gmail.com>

Author: pjdurden

internal/translator/anthropic_anthropic.go

@@ -178,24 +178,44 @@ func (a *anthropicToAnthropicTranslator) reflectStreamingEvent(eventUnion *anthr
 		}
 	case eventUnion.MessageDelta != nil:
 		u := eventUnion.MessageDelta.Usage
-		// message_delta events provide final counts for specific token types.
-		// Standard Anthropic only reports output_tokens here, but some Anthropic-compatible
-		// backends report the final input/cache token counts on message_delta instead of
-		// message_start. When those fields are present, merge them so they are not dropped.
-		// They are only merged when non-zero to avoid clobbering message_start values with a
-		// delta that carries output tokens only. See https://github.com/envoyproxy/ai-gateway/issues/2290.
-		if u.InputTokens > 0 || u.CacheReadInputTokens > 0 || u.CacheCreationInputTokens > 0 {
-			messageDeltaUsage := metrics.ExtractTokenUsageFromExplicitCaching(
-				int64(u.InputTokens),
-				int64(u.OutputTokens),
-				ptr.To(int64(u.CacheReadInputTokens)),
-				ptr.To(int64(u.CacheCreationInputTokens)),
-			)
-			a.streamingTokenUsage.Override(messageDeltaUsage)
-		} else if u.OutputTokens >= 0 {
-			// Update output tokens from message_delta (final count).
+		// message_delta carries the final counts. Standard Anthropic only reports output_tokens
+		// here, but some Anthropic-compatible backends report the final input/cache counts on
+		// message_delta instead of message_start. See https://github.com/envoyproxy/ai-gateway/issues/2290.
+		//
+		// output_tokens is always the final value on message_delta, so take it unconditionally.
+		if u.OutputTokens >= 0 {
 			a.streamingTokenUsage.SetOutputTokens(uint32(u.OutputTokens)) //nolint:gosec
 		}
+		// Merge the input/cache counts per field rather than replacing the whole usage snapshot:
+		// a delta may report only the fields that apply (the rest arrive as zero), so overwriting
+		// every field would clobber values already set on message_start. We can only treat a field
+		// as "present" when it is non-zero, since the upstream usage fields are not pointers.
+		if u.InputTokens > 0 || u.CacheReadInputTokens > 0 || u.CacheCreationInputTokens > 0 {
+			// The unified input_tokens is the sum of raw input + cache-read + cache-creation, so
+			// recover the latest known value of each component and overwrite only the ones present
+			// on this delta before recomputing the sum.
+			cacheRead, _ := a.streamingTokenUsage.CachedInputTokens()
+			cacheCreation, _ := a.streamingTokenUsage.CacheCreationInputTokens()
+			unifiedInput, _ := a.streamingTokenUsage.InputTokens()
+			var rawInput uint32
+			if unifiedInput >= cacheRead+cacheCreation {
+				rawInput = unifiedInput - cacheRead - cacheCreation
+			}
+
+			if u.InputTokens > 0 {
+				rawInput = uint32(u.InputTokens) //nolint:gosec
+			}
+			if u.CacheReadInputTokens > 0 {
+				cacheRead = uint32(u.CacheReadInputTokens) //nolint:gosec
+			}
+			if u.CacheCreationInputTokens > 0 {
+				cacheCreation = uint32(u.CacheCreationInputTokens) //nolint:gosec
+			}
+
+			a.streamingTokenUsage.SetCachedInputTokens(cacheRead)
+			a.streamingTokenUsage.SetCacheCreationInputTokens(cacheCreation)
+			a.streamingTokenUsage.SetInputTokens(rawInput + cacheRead + cacheCreation)
+		}
 	}
 }
 

internal/translator/anthropic_anthropic_test.go

@@ -238,6 +238,30 @@ data: {"type":"message_stop"}`
 		// Total input = input_tokens(4522) + cache_read_input_tokens(4511) = 9033; total = 9033 + 5 = 9038.
 		require.Equal(t, tokenUsageFrom(9033, 4511, 0, 5, 9038, -1), tokenUsage)
 	})
+
+	t.Run("partial fields on message_delta do not clobber message_start", func(t *testing.T) {
+		// A backend may set input_tokens on message_start and only report the final cache_read on
+		// message_delta. Merging must be per field: the delta omits input_tokens (reported as 0),
+		// which must not zero out the value already recorded from message_start.
+		translator := NewAnthropicToAnthropicTranslator("", "")
+		require.NotNil(t, translator)
+		translator.(*anthropicToAnthropicTranslator).stream = true
+
+		const responseBody = `event: message_start
+data: {"type":"message_start","message":{"model":"claude-sonnet-4-5-20250929","id":"msg_x","type":"message","role":"assistant","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":1000,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":0}}}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":10,"cache_read_input_tokens":500}}
+
+event: message_stop
+data: {"type":"message_stop"}`
+
+		_, _, tokenUsage, _, err := translator.ResponseBody(nil, strings.NewReader(responseBody), false, nil)
+		require.NoError(t, err)
+		// message_start raw input(1000) is preserved; delta adds cache_read(500).
+		// Total input = 1000 + 500 = 1500; total = 1500 + 10 = 1510.
+		require.Equal(t, tokenUsageFrom(1500, 500, 0, 10, 1510, -1), tokenUsage)
+	})
 }
 
 func TestAnthropicToAnthropic_ResponseError(t *testing.T) {