From 03e02c86ceb98faed9fc5ea2cb30c2367289d442 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 11 Aug 2026 02:56:00 +0000 Subject: [PATCH] [proxy] Route the Bedrock count-tokens action Both the request parser and the router enumerated Bedrock actions without count-tokens, so the path carried no model and the request denied as not-routable. Nothing breaks outright, because the client falls back to counting context through the inference endpoint, but that fallback is billable and the dedicated endpoint exists to avoid exactly that. The action carries a model in the path and returns no usage, so it routes like any other Bedrock action and meters to zero. --- .../llm_request_parser/bedrock_test.go | 26 +++++++++++++++++ .../builtin/llm_request_parser/middleware.go | 6 ++-- .../builtin/llm_router/bedrock_route_test.go | 29 +++++++++++++++++++ .../builtin/llm_router/middleware.go | 29 ++++++++++++++----- 4 files changed, 80 insertions(+), 10 deletions(-) diff --git a/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go index d8cd81437..82f44cb50 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go @@ -1,9 +1,13 @@ package llm_request_parser import ( + "context" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" ) func TestParseBedrockPath(t *testing.T) { @@ -36,3 +40,25 @@ func TestParseBedrockPath(t *testing.T) { } } } + +// TestInvoke_BedrockCountTokens covers the dedicated token-counting +// endpoint. Denying it does not break the client, it just pushes context +// counting back onto the inference endpoint, which is billable. +func TestInvoke_BedrockCountTokens(t *testing.T) { + mw := newMiddleware(t) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens", + Body: []byte(`{"input":{"converse":{"messages":[]}}}`), + }) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + + model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel) + require.True(t, ok, "count-tokens carries a model in the path and must emit it") + assert.Equal(t, "anthropic.claude-sonnet-4-5", model, "model must be normalized like any other action") + + stream, _ := metaValue(t, out.Metadata, middleware.KeyLLMStream) + assert.Equal(t, "false", stream, "count-tokens never streams") +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index 7f6e80a73..051e7949b 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -350,7 +350,9 @@ func trimBedrockNamespace(reqPath string) string { // // /model/{modelId}/{action} // -// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}. +// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream, +// count-tokens}. Token counting carries a model and no usage, so it routes +// like any other action and meters to zero. // The modelId may be URL-encoded and may carry a cross-region inference-profile // prefix and a version suffix; normalizeBedrockModel strips both so the model // matches catalog pricing. @@ -374,7 +376,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) { return bedrockRequest{}, false } switch action { - case "invoke", "converse": + case "invoke", "converse", "count-tokens": return bedrockRequest{model: model}, true case "invoke-with-response-stream", "converse-stream": return bedrockRequest{model: model, stream: true}, true diff --git a/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go b/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go index 40cbcb6bd..8c7810724 100644 --- a/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go +++ b/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go @@ -1,9 +1,13 @@ package llm_router import ( + "context" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" ) // TestRouteClaimsModel_BedrockNormalizesCandidate guards the fix for the native @@ -28,3 +32,28 @@ func TestRouteClaimsModel_BedrockNormalizesCandidate(t *testing.T) { assert.False(t, routeClaimsModel(openai, "us.gpt-4o"), "non-Bedrock routes must not strip a us. prefix") } + +// TestRouter_BedrockCountTokensRoutes pins that the token-counting action +// reaches the Bedrock route instead of denying as not-routable. +func TestRouter_BedrockCountTokensRoutes(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "bedrock-prod", + Bedrock: true, + Models: []string{"anthropic.claude-sonnet-4-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + }}}) + + in := newInputWithModelAndURL("anthropic.claude-sonnet-4-5", + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "bedrock"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "count-tokens must route, not deny") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host) +} diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 904e8f1fd..b6eba92f4 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -337,20 +337,33 @@ func splitBedrockNamespace(reqPath string) (string, bool) { return reqPath, false } +// bedrockActions are the runtime actions that follow the model id in a +// Bedrock path. count-tokens is here so a client can price its context +// against the dedicated endpoint; denying it pushes that work back onto +// the inference endpoint, which bills for it. +var bedrockActions = []string{ + "/invoke", + "/invoke-with-response-stream", + "/converse", + "/converse-stream", + "/count-tokens", +} + // isBedrockPath reports whether reqPath is an AWS Bedrock runtime model -// endpoint: /model/{modelId}/{action} where action is invoke, -// invoke-with-response-stream, converse, or converse-stream — optionally behind -// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these -// requests are routed by path to the Bedrock provider. +// endpoint: /model/{modelId}/{action} — optionally behind a "/bedrock" +// gateway-namespace prefix. The model lives in the path, so these requests +// are routed by path to the Bedrock provider. func isBedrockPath(reqPath string) bool { native, _ := splitBedrockNamespace(reqPath) if !strings.HasPrefix(native, "/model/") { return false } - return strings.HasSuffix(native, "/invoke") || - strings.HasSuffix(native, "/invoke-with-response-stream") || - strings.HasSuffix(native, "/converse") || - strings.HasSuffix(native, "/converse-stream") + for _, action := range bedrockActions { + if strings.HasSuffix(native, action) { + return true + } + } + return false } // matchVertex selects the Vertex provider authorised for the caller's groups