[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.
This commit is contained in:
mlsmaycon
2026-08-11 02:56:00 +00:00
parent 789d416215
commit 03e02c86ce
4 changed files with 80 additions and 10 deletions
@@ -1,9 +1,13 @@
package llm_request_parser package llm_request_parser
import ( import (
"context"
"testing" "testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
) )
func TestParseBedrockPath(t *testing.T) { 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")
}
@@ -350,7 +350,9 @@ func trimBedrockNamespace(reqPath string) string {
// //
// /model/{modelId}/{action} // /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 // 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 // prefix and a version suffix; normalizeBedrockModel strips both so the model
// matches catalog pricing. // matches catalog pricing.
@@ -374,7 +376,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) {
return bedrockRequest{}, false return bedrockRequest{}, false
} }
switch action { switch action {
case "invoke", "converse": case "invoke", "converse", "count-tokens":
return bedrockRequest{model: model}, true return bedrockRequest{model: model}, true
case "invoke-with-response-stream", "converse-stream": case "invoke-with-response-stream", "converse-stream":
return bedrockRequest{model: model, stream: true}, true return bedrockRequest{model: model, stream: true}, true
@@ -1,9 +1,13 @@
package llm_router package llm_router
import ( import (
"context"
"testing" "testing"
"github.com/stretchr/testify/assert" "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 // 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"), assert.False(t, routeClaimsModel(openai, "us.gpt-4o"),
"non-Bedrock routes must not strip a us. prefix") "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)
}
@@ -337,20 +337,33 @@ func splitBedrockNamespace(reqPath string) (string, bool) {
return reqPath, false 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 // isBedrockPath reports whether reqPath is an AWS Bedrock runtime model
// endpoint: /model/{modelId}/{action} where action is invoke, // endpoint: /model/{modelId}/{action} — optionally behind a "/bedrock"
// invoke-with-response-stream, converse, or converse-stream — optionally behind // gateway-namespace prefix. The model lives in the path, so these requests
// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these // are routed by path to the Bedrock provider.
// requests are routed by path to the Bedrock provider.
func isBedrockPath(reqPath string) bool { func isBedrockPath(reqPath string) bool {
native, _ := splitBedrockNamespace(reqPath) native, _ := splitBedrockNamespace(reqPath)
if !strings.HasPrefix(native, "/model/") { if !strings.HasPrefix(native, "/model/") {
return false return false
} }
return strings.HasSuffix(native, "/invoke") || for _, action := range bedrockActions {
strings.HasSuffix(native, "/invoke-with-response-stream") || if strings.HasSuffix(native, action) {
strings.HasSuffix(native, "/converse") || return true
strings.HasSuffix(native, "/converse-stream") }
}
return false
} }
// matchVertex selects the Vertex provider authorised for the caller's groups // matchVertex selects the Vertex provider authorised for the caller's groups