From 99ad5642b9c8d840dc177cce2b7e6f7d15a4b8c8 Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Sat, 25 Jul 2026 19:15:58 +0000 Subject: [PATCH] [proxy] parse Bedrock Converse camelCase cache token fields Converse responses report prompt-cache usage as cacheReadInputTokens / cacheWriteInputTokens. They were not decoded, so cached Converse requests were billed without their cache buckets. Map them into the same Usage fields as the InvokeModel snake_case variants, for both buffered and streaming responses. --- proxy/internal/llm/bedrock.go | 18 +++++++++++------- proxy/internal/llm/bedrock_test.go | 12 ++++++++++++ .../llm_response_parser/streaming_bedrock.go | 17 +++++++++++++---- .../streaming_bedrock_test.go | 18 ++++++++++++++++++ 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/proxy/internal/llm/bedrock.go b/proxy/internal/llm/bedrock.go index f7802beb2..eb64167a1 100644 --- a/proxy/internal/llm/bedrock.go +++ b/proxy/internal/llm/bedrock.go @@ -56,10 +56,12 @@ type bedrockResponse struct { OutputTokens int64 `json:"output_tokens"` CacheReadInputTokens int64 `json:"cache_read_input_tokens"` CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"` - // Converse — camelCase. - InputTokensCamel int64 `json:"inputTokens"` - OutputTokensCamel int64 `json:"outputTokens"` - TotalTokensCamel int64 `json:"totalTokens"` + // Converse — camelCase; cache buckets are additive to inputTokens (AWS names the write bucket cacheWriteInputTokens). + InputTokensCamel int64 `json:"inputTokens"` + OutputTokensCamel int64 `json:"outputTokens"` + TotalTokensCamel int64 `json:"totalTokens"` + CacheReadTokensCamel int64 `json:"cacheReadInputTokens"` + CacheWriteTokensCamel int64 `json:"cacheWriteInputTokens"` } `json:"usage"` } @@ -83,16 +85,18 @@ func (BedrockParser) ParseResponse(status int, contentType string, body []byte) } inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel) outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel) + cacheRead := firstNonZero(resp.Usage.CacheReadInputTokens, resp.Usage.CacheReadTokensCamel) + cacheWrite := firstNonZero(resp.Usage.CacheCreationInputTokens, resp.Usage.CacheWriteTokensCamel) total := resp.Usage.TotalTokensCamel if total == 0 { - total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens + total = inTok + outTok + cacheRead + cacheWrite } return Usage{ InputTokens: inTok, OutputTokens: outTok, TotalTokens: total, - CachedInputTokens: resp.Usage.CacheReadInputTokens, - CacheCreationTokens: resp.Usage.CacheCreationInputTokens, + CachedInputTokens: cacheRead, + CacheCreationTokens: cacheWrite, }, nil } diff --git a/proxy/internal/llm/bedrock_test.go b/proxy/internal/llm/bedrock_test.go index ca6f092f3..e99ee55df 100644 --- a/proxy/internal/llm/bedrock_test.go +++ b/proxy/internal/llm/bedrock_test.go @@ -26,6 +26,18 @@ func TestBedrockParser_ParseResponse_Converse(t *testing.T) { require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total") } +// Converse camelCase cache fields must land in the billed Usage buckets, same as the InvokeModel snake_case fields. +func TestBedrockParser_ParseResponse_ConverseCacheBuckets(t *testing.T) { + body := []byte(`{"usage":{"inputTokens":11,"outputTokens":3,"cacheReadInputTokens":7,"cacheWriteInputTokens":9}}`) + u, err := BedrockParser{}.ParseResponse(200, "application/json", body) + require.NoError(t, err) + require.Equal(t, int64(11), u.InputTokens, "converse input tokens") + require.Equal(t, int64(3), u.OutputTokens, "converse output tokens") + require.Equal(t, int64(7), u.CachedInputTokens, "converse cache-read tokens") + require.Equal(t, int64(9), u.CacheCreationTokens, "converse cache-write tokens") + require.Equal(t, int64(11+3+7+9), u.TotalTokens, "total backfill is additive when the provider omits totalTokens") +} + func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) { _, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary")) require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator") diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go index a82a9cdbc..30809b46e 100644 --- a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock.go @@ -69,15 +69,18 @@ func applyBedrockInvokeChunk(payload []byte, usage *llm.Usage, completion *strin } // converseStreamEvent captures the Converse stream frames carrying completion -// text (contentBlockDelta) and the final token usage (metadata). +// text (contentBlockDelta) and the final token usage (metadata). Cache buckets +// are additive to inputTokens (AWS write bucket: cacheWriteInputTokens). type converseStreamEvent struct { Delta *struct { Text string `json:"text"` } `json:"delta"` Usage *struct { - InputTokens int64 `json:"inputTokens"` - OutputTokens int64 `json:"outputTokens"` - TotalTokens int64 `json:"totalTokens"` + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` + TotalTokens int64 `json:"totalTokens"` + CacheReadTokens int64 `json:"cacheReadInputTokens"` + CacheWriteTokens int64 `json:"cacheWriteInputTokens"` } `json:"usage"` } @@ -105,6 +108,12 @@ func applyConverseStreamEvent(eventType string, payload []byte, usage *llm.Usage if ev.Usage.TotalTokens > 0 { usage.TotalTokens = ev.Usage.TotalTokens } + if ev.Usage.CacheReadTokens > 0 { + usage.CachedInputTokens = ev.Usage.CacheReadTokens + } + if ev.Usage.CacheWriteTokens > 0 { + usage.CacheCreationTokens = ev.Usage.CacheWriteTokens + } } } } diff --git a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go index f93505882..f66347591 100644 --- a/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go +++ b/proxy/internal/middleware/builtin/llm_response_parser/streaming_bedrock_test.go @@ -66,6 +66,24 @@ func TestAccumulateBedrockStream_Converse(t *testing.T) { require.Equal(t, "pong", completion, "converse text deltas concatenated") } +// The converse-stream metadata frame's camelCase cache fields must reach the billed cache buckets. +func TestAccumulateBedrockStream_ConverseCacheBuckets(t *testing.T) { + var body bytes.Buffer + body.Write(bedrockFrame(t, "contentBlockDelta", mustJSON(t, map[string]any{"delta": map[string]any{"text": "pong"}}))) + body.Write(bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{ + "inputTokens": 11, "outputTokens": 3, "totalTokens": 30, + "cacheReadInputTokens": 7, "cacheWriteInputTokens": 9, + }}))) + + usage, completion := accumulateBedrockStream(body.Bytes()) + require.Equal(t, int64(11), usage.InputTokens, "input tokens from metadata frame") + require.Equal(t, int64(3), usage.OutputTokens, "output tokens from metadata frame") + require.Equal(t, int64(7), usage.CachedInputTokens, "cache-read tokens from metadata frame") + require.Equal(t, int64(9), usage.CacheCreationTokens, "cache-write tokens from metadata frame") + require.Equal(t, int64(30), usage.TotalTokens, "provider-reported total wins") + require.Equal(t, "pong", completion) +} + func TestAccumulateBedrockStream_Truncated(t *testing.T) { // A body cut mid-frame must not panic; partial usage is returned. full := bedrockFrame(t, "metadata", mustJSON(t, map[string]any{"usage": map[string]any{"inputTokens": 11, "outputTokens": 3}}))