diff --git a/e2e/agentnetwork/chat_test.go b/e2e/agentnetwork/chat_test.go index 90c3766ec..29d0c7c6c 100644 --- a/e2e/agentnetwork/chat_test.go +++ b/e2e/agentnetwork/chat_test.go @@ -9,12 +9,83 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/netbirdio/netbird/e2e/harness" "github.com/netbirdio/netbird/shared/management/http/api" ) +// publishedPer1k carries the vendors' PUBLISHED per-1k-token USD rates for the +// models the live matrix can drive, keyed by the normalized model id the proxy +// stamps on the access-log row. These are intentionally hardcoded from the +// public price lists (not read from the proxy's pricing table) so the e2e run +// cross-checks the whole billing pipeline against an independent source: a +// wrong embedded rate, a broken model normalization, or a per-token-instead-of +// -per-1k-chunk regression (a 1000× blowup) all fail the assertion. +var publishedPer1k = map[string]struct{ in, out float64 }{ + "gpt-4o-mini": {0.00015, 0.0006}, // $0.15 / $0.60 per MTok + "gpt-4o": {0.0025, 0.01}, // $2.50 / $10 per MTok + "claude-haiku-4-5": {0.001, 0.005}, // $1 / $5 per MTok + "claude-sonnet-4-5": {0.003, 0.015}, // $3 / $15 per MTok + "claude-sonnet-4-6": {0.003, 0.015}, + "kimi-k3": {0.003, 0.015}, + "anthropic.claude-haiku-4-5": {0.001, 0.005}, // Bedrock mirrors first-party rates + "anthropic.claude-sonnet-4-5": {0.003, 0.015}, + "anthropic.claude-sonnet-4-6": {0.003, 0.015}, +} + +// validateAccessLogCost recomputes the expected USD cost of a live access-log +// row from the published per-1k rates and asserts the stored cost_usd matches. +// +// Anthropic-shape providers report prompt-cache buckets ADDITIVELY: they are +// billed (read ≈0.1×, write ≈1.25× the input rate) and folded into +// total_tokens, but not into input_tokens/output_tokens — so when +// total > in + out the expectation widens to the [all-read, all-write] band +// for the extra tokens instead of failing on correct cache billing. +// +// Models the proxy deliberately does not price (gateway-prefixed ids like +// "openai/gpt-4o-mini") must store cost 0, never a guessed rate. +func validateAccessLogCost(t *testing.T, pc providerCase, row api.AgentNetworkAccessLog) { + t.Helper() + model := catalogModel(pc) + rates, known := publishedPer1k[model] + if !known { + if strings.Contains(model, "/") { + assert.Zerof(t, row.CostUsd, + "gateway-prefixed model %q is not in the pricing table so the cost meter must skip (cost 0), got %v", model, row.CostUsd) + return + } + t.Logf("no published rate on file for model %q (env-overridden?); skipping cost validation", model) + return + } + + require.Positive(t, row.InputTokens, "priced row must carry input tokens") + require.Positive(t, row.OutputTokens, "priced row must carry output tokens") + + base := float64(row.InputTokens)/1000*rates.in + float64(row.OutputTokens)/1000*rates.out + + // Cache buckets ride total_tokens only (additive Anthropic/Bedrock shape). + cacheTokens := row.TotalTokens - row.InputTokens - row.OutputTokens + if cacheTokens < 0 { + cacheTokens = 0 + } + + if cacheTokens == 0 { + assert.InDeltaf(t, base, row.CostUsd, 1e-6, + "cost for %s (%s): %d in × $%v/1k + %d out × $%v/1k must equal the stored cost", + pc.name, model, row.InputTokens, rates.in, row.OutputTokens, rates.out) + return + } + + lo := base + float64(cacheTokens)/1000*rates.in*0.1 // whole bucket read from cache + hi := base + float64(cacheTokens)/1000*rates.in*1.25 // whole bucket written to cache + assert.GreaterOrEqualf(t, row.CostUsd, lo-1e-6, + "cost for %s (%s) below the all-cache-read floor (base %v, %d cache tokens)", pc.name, model, base, cacheTokens) + assert.LessOrEqualf(t, row.CostUsd, hi+1e-6, + "cost for %s (%s) above the all-cache-write ceiling (base %v, %d cache tokens)", pc.name, model, base, cacheTokens) +} + // providerCase is one entry in the live provider matrix. The same scenario runs // for every available provider; availability is keyed off env vars so the suite // covers whatever credentials are present (source ~/.llm-keys locally / set the @@ -290,6 +361,7 @@ func TestProvidersMatrix(t *testing.T) { // The session id sent as x-session-id must round-trip into the // access-log row for this provider. + var row api.AgentNetworkAccessLog require.Eventually(t, func() bool { logs, lerr := srv.ListAccessLogs(ctx) if lerr != nil { @@ -297,11 +369,17 @@ func TestProvidersMatrix(t *testing.T) { } for _, r := range logs.Data { if r.SessionId != nil && *r.SessionId == sessionID { + row = r return true } } return false }, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row for %s", sessionID, pc.name) + + // The stored cost must match the vendor's published per-1k rates + // applied to the row's token counts (cache-aware for the additive + // Anthropic/Bedrock buckets, zero for unpriced gateway model ids). + validateAccessLogCost(t, pc, row) }) } diff --git a/proxy/internal/llm/bedrock.go b/proxy/internal/llm/bedrock.go index f7802beb2..54bdf40ef 100644 --- a/proxy/internal/llm/bedrock.go +++ b/proxy/internal/llm/bedrock.go @@ -56,10 +56,14 @@ 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, + // mirroring the InvokeModel snake_case fields above (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 +87,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..bf916c702 100644 --- a/proxy/internal/llm/bedrock_test.go +++ b/proxy/internal/llm/bedrock_test.go @@ -26,6 +26,21 @@ func TestBedrockParser_ParseResponse_Converse(t *testing.T) { require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total") } +// TestBedrockParser_ParseResponse_ConverseCacheBuckets proves the Converse +// camelCase cache fields (cacheReadInputTokens / cacheWriteInputTokens) land +// in the same Usage buckets as the InvokeModel snake_case fields — the cost +// meter bills them, so dropping them silently under-counts cached requests. +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/llm/pricing/catalog_sync_test.go b/proxy/internal/llm/pricing/catalog_sync_test.go new file mode 100644 index 000000000..8bfbd6e84 --- /dev/null +++ b/proxy/internal/llm/pricing/catalog_sync_test.go @@ -0,0 +1,62 @@ +package pricing + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog" +) + +// catalogPricingProviders maps each metered catalog provider id to the +// pricing-table provider key(s) its requests are billed under (the +// llm.Parser surface the proxy meters that provider with). A catalog +// provider whose requests can hit more than one parser surface (Kimi +// serves both the OpenAI and Anthropic body shapes off one host) must +// carry matching entries in every listed table. +// +// Gateway/custom catalog entries (LiteLLM, Portkey, vLLM, …) publish no +// model list, so they never appear here. +var catalogPricingProviders = map[string][]string{ + "openai_api": {"openai"}, + "azure_openai_api": {"openai"}, + "mistral_api": {"openai"}, + "anthropic_api": {"anthropic"}, + "vertex_ai_api": {"anthropic"}, // Anthropic-on-Vertex: bare claude-* ids, anthropic parser + "bedrock_api": {"bedrock"}, + "kimi_api": {"openai", "anthropic"}, +} + +// TestDefaultPricing_MatchesCatalog pins the management catalog (the prices +// the dashboard displays) to the proxy's embedded default pricing table (the +// prices the cost meter bills). A drift between the two makes correct costs +// look wrong — the dashboard advertises one rate while the proxy meters +// another — so every catalog model must resolve to a pricing entry with +// byte-identical input/output per-1k rates. +func TestDefaultPricing_MatchesCatalog(t *testing.T) { + table := DefaultTable() + + for _, p := range catalog.All() { + if len(p.Models) == 0 { + continue // gateways / custom endpoints publish no models + } + keys, metered := catalogPricingProviders[p.ID] + require.Truef(t, metered, + "catalog provider %q publishes models but has no pricing-table mapping — add it to catalogPricingProviders and defaults_pricing.yaml", + p.ID) + + for _, m := range p.Models { + for _, key := range keys { + entry, ok := table.entries[key][m.ID] + require.Truef(t, ok, + "catalog model %s/%s must have a %q pricing entry or the cost meter silently skips it (cost.skipped=unknown_model)", + p.ID, m.ID, key) + assert.Equalf(t, m.InputPer1k, entry.InputPer1K, + "input rate drift for %s/%s: dashboard shows %v, proxy bills %v", p.ID, m.ID, m.InputPer1k, entry.InputPer1K) + assert.Equalf(t, m.OutputPer1k, entry.OutputPer1K, + "output rate drift for %s/%s: dashboard shows %v, proxy bills %v", p.ID, m.ID, m.OutputPer1k, entry.OutputPer1K) + } + } + } +} diff --git a/proxy/internal/middleware/builtin/cost_calculation_matrix_test.go b/proxy/internal/middleware/builtin/cost_calculation_matrix_test.go new file mode 100644 index 000000000..1ee3f7724 --- /dev/null +++ b/proxy/internal/middleware/builtin/cost_calculation_matrix_test.go @@ -0,0 +1,359 @@ +package builtin_test + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "strconv" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser" + "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser" +) + +// TestCostCalculation_ProviderMatrix drives the proxy's real cost pipeline — +// llm_request_parser (provider/model detection + Bedrock/Vertex path +// normalization) → llm_response_parser (usage extraction from realistic wire +// bodies) → cost_meter (pricing) — against the EMBEDDED DEFAULT pricing table +// (no operator override file), and asserts the exact USD amount for every +// metered provider surface. +// +// Expected values are hardcoded dollar amounts derived from the providers' +// published per-million-token prices, NOT recomputed from the pricing table, +// so the test cross-checks three things at once: +// +// 1. the embedded rates match the published prices, +// 2. the per-1k rates are applied to 1k CHUNKS (tokens/1000 × rate) — a +// missing ÷1000 would inflate every expectation by exactly 1000×, +// 3. the per-provider cache-bucket semantics (OpenAI cached-subset discount +// vs Anthropic/Bedrock additive read/write buckets) bill correctly. +// +// It includes the field-reported scenario: Bedrock + Claude Sonnet 4.6 with +// 3 input / 1514 output tokens costs $0.022719 bare, and $0.137199 when the +// first request of a session also writes a 30,528-token prompt cache +// (cache_creation_input_tokens — no previous request needed; the write IS the +// first request). +func TestCostCalculation_ProviderMatrix(t *testing.T) { + // Empty data dir → cost_meter runs on the embedded defaults, exactly like + // a proxy with no operator pricing override. + builtin.Configure(context.Background(), t.TempDir(), nil, nil, nil) + + reqMW, err := llm_request_parser.Factory{}.New(nil) + require.NoError(t, err, "build llm_request_parser") + respMW, err := llm_response_parser.Factory{}.New(nil) + require.NoError(t, err, "build llm_response_parser") + costMW, err := cost_meter.Factory{}.New(nil) + require.NoError(t, err, "build cost_meter") + t.Cleanup(func() { _ = costMW.Close() }) + + const jsonCT = "application/json" + const sseCT = "text/event-stream" + const awsCT = "application/vnd.amazon.eventstream" + + cases := []struct { + name string + url string + reqBody []byte + respCT string + respBody []byte + + wantProvider string + wantModel string + wantCost float64 // exact expected USD; ignored when wantSkip is set + wantSkip string // expected cost.skipped reason, "" when priced + }{ + { + // OpenAI Chat Completions, $0.15/M in + $0.60/M out (gpt-4o-mini): + // 1000×0.15/1M + 500×0.60/1M. + name: "openai chat completions", + url: "https://api.openai.com/v1/chat/completions", + reqBody: []byte(`{"model":"gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"choices":[{"message":{"content":"pong"}}],"usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`), + wantProvider: "openai", + wantModel: "gpt-4o-mini", + wantCost: 0.00045, + }, + { + // OpenAI cached prompt tokens are a SUBSET of prompt_tokens and + // bill at the discount rate; gpt-4o at $2.50/$10 per MTok with + // $1.25/M cached: 250×2.5/1M + 750×1.25/1M + 500×10/1M. + name: "openai cached subset discount", + url: "https://api.openai.com/v1/chat/completions", + reqBody: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500,"prompt_tokens_details":{"cached_tokens":750}}}`), + wantProvider: "openai", + wantModel: "gpt-4o", + wantCost: 0.0065625, + }, + { + // OpenAI streaming: usage rides the final SSE frame. + name: "openai chat SSE stream", + url: "https://api.openai.com/v1/chat/completions", + reqBody: []byte(`{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}`), + respCT: sseCT, + respBody: sseBody(`{"choices":[{"delta":{"content":"po"}}]}`, `{"choices":[{"delta":{"content":"ng"}}]}`, `{"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":500}}`, "[DONE]"), + wantProvider: "openai", + wantModel: "gpt-4o-mini", + wantCost: 0.00045, + }, + { + // Mistral speaks the OpenAI shape and is priced under the openai + // table: mistral-large-latest at $0.50/$1.50 per MTok. + name: "mistral via openai shape", + url: "https://api.mistral.ai/v1/chat/completions", + reqBody: []byte(`{"model":"mistral-large-latest","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":1000}}`), + wantProvider: "openai", + wantModel: "mistral-large-latest", + wantCost: 0.002, + }, + { + // The field report, minus caching: Bedrock Claude Sonnet 4.6 at + // $3/M in + $15/M out. 3×3/1M + 1514×15/1M = $0.022719. Also + // covers inference-profile normalization: the URL carries the + // full region-prefixed versioned id. + name: "bedrock invoke — reported scenario, no cache", + url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke", + reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":3,"output_tokens":1514}}`), + wantProvider: "bedrock", + wantModel: "anthropic.claude-sonnet-4-6", + wantCost: 0.022719, + }, + { + // The field report as observed: same request whose FIRST call + // also wrote a 30,528-token prompt cache. Cache writes bill at + // 1.25× input ($3.75/M): 0.022719 + 30528×3.75/1M = $0.137199, + // which renders as the reported $0.1372. + name: "bedrock invoke — reported scenario with cache write", + url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke", + reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"input_tokens":3,"output_tokens":1514,"cache_creation_input_tokens":30528,"cache_read_input_tokens":0}}`), + wantProvider: "bedrock", + wantModel: "anthropic.claude-sonnet-4-6", + wantCost: 0.137199, + }, + { + // Same numbers over the InvokeModel event-stream: message_start + // carries input + cache buckets, message_delta the output count. + name: "bedrock invoke stream with cache write", + url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-sonnet-4-6-20260115-v1:0/invoke-with-response-stream", + reqBody: []byte(`{"messages":[{"role":"user","content":"hi"}]}`), + respCT: awsCT, + respBody: bedrockInvokeStream(t, `{"type":"message_start","message":{"usage":{"input_tokens":3,"output_tokens":1,"cache_creation_input_tokens":30528}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":1514}}`), + wantProvider: "bedrock", + wantModel: "anthropic.claude-sonnet-4-6", + wantCost: 0.137199, + }, + { + // Bedrock Converse reports usage camelCase, cache buckets + // included (cacheRead/cacheWriteInputTokens). Haiku 4.5 at + // $1/$5 per MTok, cache read $0.10/M, cache write $1.25/M: + // 50×1/1M + 100×5/1M + 2000×0.1/1M + 1000×1.25/1M = $0.002. + name: "bedrock converse with cache buckets", + url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse", + reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`), + respCT: jsonCT, + respBody: []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`), + wantProvider: "bedrock", + wantModel: "anthropic.claude-haiku-4-5", + wantCost: 0.002, + }, + { + // Same numbers over converse-stream: usage rides the trailing + // metadata frame. + name: "bedrock converse stream with cache buckets", + url: "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse-stream", + reqBody: []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`), + respCT: awsCT, + respBody: bedrockConverseStream(t, + `{"delta":{"text":"pong"}}`, + `{"usage":{"inputTokens":50,"outputTokens":100,"totalTokens":3150,"cacheReadInputTokens":2000,"cacheWriteInputTokens":1000}}`, + ), + wantProvider: "bedrock", + wantModel: "anthropic.claude-haiku-4-5", + wantCost: 0.002, + }, + { + // First-party Anthropic Messages API, cache buckets additive: + // Sonnet 4.6 with 256 in + 200 out + 768 cache read ($0.30/M) + // + 512 cache write ($3.75/M): + // 256×3/1M + 200×15/1M + 768×0.3/1M + 512×3.75/1M. + name: "anthropic messages with cache buckets", + url: "https://api.anthropic.com/v1/messages", + reqBody: []byte(`{"model":"claude-sonnet-4-6","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"content":[{"type":"text","text":"pong"}],"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`), + wantProvider: "anthropic", + wantModel: "claude-sonnet-4-6", + wantCost: 0.0059184, + }, + { + // Anthropic SSE stream: input + cache from message_start, + // output from message_delta. Haiku 4.5, 1000 in + 2000 out: + // 1000×1/1M + 2000×5/1M = $0.011. + name: "anthropic SSE stream", + url: "https://api.anthropic.com/v1/messages", + reqBody: []byte(`{"model":"claude-haiku-4-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`), + respCT: sseCT, + respBody: sseBody(`{"type":"message_start","message":{"usage":{"input_tokens":1000,"output_tokens":2}}}`, `{"type":"content_block_delta","delta":{"type":"text_delta","text":"pong"}}`, `{"type":"message_delta","usage":{"output_tokens":2000}}`, `{"type":"message_stop"}`), + wantProvider: "anthropic", + wantModel: "claude-haiku-4-5", + wantCost: 0.011, + }, + { + // Kimi's Anthropic-compatible endpoint (the Claude Code setup): + // kimi-k3 at $3/$15 per MTok under the anthropic table. + name: "kimi anthropic shape", + url: "https://api.moonshot.ai/anthropic/v1/messages", + reqBody: []byte(`{"model":"kimi-k3","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"input_tokens":1000,"output_tokens":1000}}`), + wantProvider: "anthropic", + wantModel: "kimi-k3", + wantCost: 0.018, + }, + { + // Vertex carries publisher + model in the URL path; the + // "@version" suffix is stripped and Anthropic-on-Vertex is + // priced under the anthropic table: 200×3/1M + 100×15/1M. + name: "vertex anthropic path-routed", + url: "https://aiplatform.googleapis.com/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-4-6@20260115:rawPredict", + reqBody: []byte(`{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"input_tokens":200,"output_tokens":100}}`), + wantProvider: "anthropic", + wantModel: "claude-sonnet-4-6", + wantCost: 0.0021, + }, + { + // Gateway-prefixed model ids (Vercel / OpenRouter style) are not + // in the pricing table: the meter must SKIP, never guess a rate. + name: "gateway-prefixed model skips pricing", + url: "https://gateway.example.com/v1/chat/completions", + reqBody: []byte(`{"model":"openai/gpt-4o-mini","messages":[{"role":"user","content":"hi"}]}`), + respCT: jsonCT, + respBody: []byte(`{"usage":{"prompt_tokens":1000,"completion_tokens":500}}`), + wantProvider: "openai", + wantModel: "openai/gpt-4o-mini", + wantSkip: "unknown_model", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + in := &middleware.Input{ + Method: "POST", + URL: tc.url, + Headers: []middleware.KV{{Key: "Content-Type", Value: "application/json"}}, + Body: tc.reqBody, + } + + reqOut, err := reqMW.Invoke(context.Background(), in) + require.NoError(t, err, "request parser") + in.Metadata = append(in.Metadata, reqOut.Metadata...) + + require.Equal(t, tc.wantProvider, metaKV(in.Metadata, middleware.KeyLLMProvider), "detected provider") + require.Equal(t, tc.wantModel, metaKV(in.Metadata, middleware.KeyLLMModel), "detected (normalized) model") + + in.Status = 200 + in.RespHeaders = []middleware.KV{{Key: "Content-Type", Value: tc.respCT}} + in.RespBody = tc.respBody + + respOut, err := respMW.Invoke(context.Background(), in) + require.NoError(t, err, "response parser") + in.Metadata = append(in.Metadata, respOut.Metadata...) + + costOut, err := costMW.Invoke(context.Background(), in) + require.NoError(t, err, "cost meter") + + if tc.wantSkip != "" { + assert.Equal(t, tc.wantSkip, metaKV(costOut.Metadata, middleware.KeyCostSkipped), "expected cost skip reason") + assert.Empty(t, metaKV(costOut.Metadata, middleware.KeyCostUSDTotal), "no cost may be emitted on skip") + return + } + + raw := metaKV(costOut.Metadata, middleware.KeyCostUSDTotal) + require.NotEmpty(t, raw, "cost.usd_total must be emitted; skip=%q", metaKV(costOut.Metadata, middleware.KeyCostSkipped)) + got, err := strconv.ParseFloat(raw, 64) + require.NoError(t, err, "cost must be a float") + // cost.usd_total is rendered with %.6f, so allow half of the + // last printed digit on top of float error. + assert.InDelta(t, tc.wantCost, got, 5.1e-7, "USD cost for %s", tc.name) + }) + } +} + +// metaKV returns the value for key in kvs, or "" when absent. +func metaKV(kvs []middleware.KV, key string) string { + for _, kv := range kvs { + if kv.Key == key { + return kv.Value + } + } + return "" +} + +// sseBody renders data frames as a text/event-stream body. +func sseBody(frames ...string) []byte { + var b bytes.Buffer + for _, f := range frames { + b.WriteString("data: ") + b.WriteString(f) + b.WriteString("\n\n") + } + return b.Bytes() +} + +// awsFrame encodes one AWS event-stream frame with the given :event-type. +func awsFrame(t *testing.T, eventType string, payload []byte) []byte { + t.Helper() + var buf bytes.Buffer + enc := eventstream.NewEncoder() + require.NoError(t, enc.Encode(&buf, eventstream.Message{ + Headers: eventstream.Headers{{Name: ":event-type", Value: eventstream.StringValue(eventType)}}, + Payload: payload, + }), "encode event-stream frame") + return buf.Bytes() +} + +// bedrockInvokeStream builds an invoke-with-response-stream body: each +// "chunk" frame wraps a base64-encoded Anthropic stream event. +func bedrockInvokeStream(t *testing.T, events ...string) []byte { + t.Helper() + var body bytes.Buffer + for _, ev := range events { + wrap, err := json.Marshal(map[string]string{"bytes": base64.StdEncoding.EncodeToString([]byte(ev))}) + require.NoError(t, err) + body.Write(awsFrame(t, "chunk", wrap)) + } + return body.Bytes() +} + +// bedrockConverseStream builds a converse-stream body: N contentBlockDelta +// frames followed by the trailing metadata frame carrying usage. +func bedrockConverseStream(t *testing.T, deltas ...string) []byte { + t.Helper() + var body bytes.Buffer + for i, ev := range deltas { + eventType := "contentBlockDelta" + if i == len(deltas)-1 { + eventType = "metadata" + } + body.Write(awsFrame(t, eventType, []byte(ev))) + } + return body.Bytes() +} 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..ea74bb191 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,19 @@ 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). The cache +// buckets are additive to inputTokens, same as the InvokeModel snake_case +// shape (AWS names the 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 +109,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..84f47613c 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,26 @@ func TestAccumulateBedrockStream_Converse(t *testing.T) { require.Equal(t, "pong", completion, "converse text deltas concatenated") } +// TestAccumulateBedrockStream_ConverseCacheBuckets proves the converse-stream +// metadata frame's camelCase cache fields reach the billed cache buckets, same +// as the InvokeModel wrapped Anthropic events. +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}}))