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