[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-22 21:14:05 +02:00
committed by GitHub
parent 1949a234e7
commit d0340a9e2f
4 changed files with 80 additions and 10 deletions
@@ -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")
}
@@ -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