[proxy] normalize Vertex model ids in routing, guardrail, and token counting

Vertex requests reach the router/guardrail with the "@version" suffix
already stripped from the URL model by the request parser, but the
operator-facing config may carry the raw versioned id the Google console
documents (e.g. "claude-opus-4-6@20250514"):

- a Vertex provider registered with versioned model ids denied every
  request as model_not_routable (the Vertex analog of the Bedrock gap
  fixed in #6773), and
- a guardrail allowlist entry with a versioned id never matched, denying
  the model as model_blocked.

Both bit a customer driving Anthropic-on-Vertex through the agent
network with unversioned model ids (…/models/claude-opus-4-6:rawPredict
at the global location).

Introduce a shared llm.NormalizeVertexModel (the same "@" stripping the
parser already does) and apply it to a Vertex route's candidate models in
routeClaimsModel and to guardrail allowlist entries, so either spelling
of the same model matches. Non-Vertex routes and entries without "@"
keep exact matching.

Also resolve the real model for the Vertex token-count endpoint:
…/models/count-tokens:rawPredict carries the literal "count-tokens"
pseudo-model in the URL and the actual model in the body (Claude Code
sends one such call per request), so any configured allowlist denied all
token counting. The parser now reads the body model for that endpoint;
an unreadable body keeps the pseudo-model and fails closed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XDgMfwgU1Xerfaedu6bQse
This commit is contained in:
Maycon Santos
2026-07-22 07:32:26 +00:00
parent d820aaa8fc
commit 203efce679
8 changed files with 264 additions and 10 deletions

View File

@@ -0,0 +1,17 @@
package llm
import "strings"
// NormalizeVertexModel strips the "@version" suffix from a Google Vertex AI
// publisher model id so it matches the catalog/pricing key, e.g.
// "claude-sonnet-4-5@20250929" -> "claude-sonnet-4-5". It is the single source
// of truth shared by the request parser (which normalizes the request model
// from the URL path), the router (which normalizes the operator's registered
// Vertex model ids), and the guardrail (which normalizes allowlist entries) so
// both spellings of the same model compare equal.
func NormalizeVertexModel(modelID string) string {
if at := strings.Index(modelID, "@"); at >= 0 {
return modelID[:at]
}
return modelID
}

View File

@@ -0,0 +1,23 @@
package llm
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestNormalizeVertexModel(t *testing.T) {
cases := map[string]string{
"claude-sonnet-4-5@20250929": "claude-sonnet-4-5",
"claude-opus-4-6@20250514": "claude-opus-4-6",
// Bare ids (the form Vertex resolves to the default version) pass through.
"claude-opus-4-6": "claude-opus-4-6",
// Numeric version aliases used by Google publisher models.
"text-embedding-005@001": "text-embedding-005",
// A leading "@" yields the empty string; callers treat that as "no model".
"@cf/meta/llama-3-8b": "",
}
for in, want := range cases {
require.Equal(t, want, NormalizeVertexModel(in), "normalize %q", in)
}
}

View File

@@ -13,6 +13,7 @@ import (
"context"
"unicode/utf8"
"github.com/netbirdio/netbird/proxy/internal/llm"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
@@ -162,6 +163,15 @@ func (m *Middleware) modelInAllowlist(model string) bool {
if allowed == normalised {
return true
}
// Vertex model ids are often stored with their "@version" suffix (e.g.
// "claude-opus-4-6@20250514") — the id the Google console documents —
// while the request parser strips the suffix before the model reaches
// the guardrail. Accept an entry whose version-stripped form matches so
// both spellings of the same model pass one allowlist. Entries without
// an "@" are untouched, so non-Vertex matching stays exact.
if v := llm.NormalizeVertexModel(allowed); v != "" && v != allowed && v == normalised {
return true
}
}
return false
}

View File

@@ -70,6 +70,27 @@ func TestAllowlistMatchAllows(t *testing.T) {
assert.Equal(t, middleware.DecisionAllow, out.Decision, "model in allowlist must be allowed")
}
// A Vertex allowlist entry stored with its "@version" suffix must match the
// request model, which reaches the guardrail with the suffix already stripped
// by the request parser.
func TestAllowlistVertexVersionedEntryMatchesStrippedModel(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"claude-opus-4-6@20250514"}})
out, err := mw.Invoke(context.Background(), newInput(
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4-6"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"@version allowlist entry must match the version-stripped request model")
// The stripped form must not over-match a different model.
out, err = mw.Invoke(context.Background(), newInput(
middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-4-8"},
))
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"a different model must stay denied")
}
func TestAllowlistMissDenies(t *testing.T) {
mw := New(Config{ModelAllowlist: []string{"gpt-4o"}})
out, err := mw.Invoke(context.Background(), newInput(

View File

@@ -104,3 +104,96 @@ func TestModelAllowlist_URLRoutedProviders(t *testing.T) {
})
}
}
// TestModelAllowlist_VertexRequestShapes replays the Vertex request shapes an
// Anthropic SDK client (Claude Code) actually sends — the model travels in the
// URL path, optionally without an "@version" suffix, at the global location —
// against allowlists holding either the bare catalog id or the raw versioned
// id. The URLs mirror the shape of a customer-reported request (unversioned
// model, global location, a project id that itself contains "claude").
func TestModelAllowlist_VertexRequestShapes(t *testing.T) {
const (
opusBare = "/v1/projects/corp-gcp-it-all-claude/locations/global/publishers/anthropic/models/claude-opus-4-6:rawPredict"
opusBareSSE = "/v1/projects/corp-gcp-it-all-claude/locations/global/publishers/anthropic/models/claude-opus-4-6:streamRawPredict"
countTokens = "/v1/projects/corp-gcp-it-all-claude/locations/global/publishers/anthropic/models/count-tokens:rawPredict"
messagesBody = `{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"hi"}]}`
countOpusBody = `{"model":"claude-opus-4-6","messages":[{"role":"user","content":"hi"}]}`
)
tests := []struct {
name string
url string
body string
allowlist []string
decision middleware.Decision
denyCode string
}{
{
name: "unversioned model allowed by bare catalog entry",
url: opusBare,
body: messagesBody,
allowlist: []string{"claude-opus-4-6"},
decision: middleware.DecisionAllow,
},
{
name: "unversioned model allowed by @version allowlist entry",
url: opusBare,
body: messagesBody,
allowlist: []string{"claude-opus-4-6@20250514"},
decision: middleware.DecisionAllow,
},
{
name: "streaming action allowed the same as rawPredict",
url: opusBareSSE,
body: messagesBody,
allowlist: []string{"claude-opus-4-6"},
decision: middleware.DecisionAllow,
},
{
// The original customer report: a Sonnet-only allowlist must block
// an Opus request regardless of the id carrying a version suffix.
name: "unversioned model outside the allowlist denied",
url: opusBare,
body: messagesBody,
allowlist: []string{"claude-sonnet-4-5"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_blocked",
},
{
name: "count-tokens resolves the body model and passes when allowed",
url: countTokens,
body: countOpusBody,
allowlist: []string{"claude-opus-4-6"},
decision: middleware.DecisionAllow,
},
{
name: "count-tokens with a disallowed body model denied",
url: countTokens,
body: countOpusBody,
allowlist: []string{"claude-sonnet-4-5"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_blocked",
},
{
// No model in the body: the pseudo-model stays and fails closed.
name: "count-tokens without a body model fails closed",
url: countTokens,
body: messagesBody,
allowlist: []string{"claude-opus-4-6"},
decision: middleware.DecisionDeny,
denyCode: "llm_policy.model_blocked",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
out := runParserGuardrail(t, tt.url, []byte(tt.body), tt.allowlist)
assert.Equal(t, tt.decision, out.Decision, "unexpected decision for %s", tt.name)
if tt.decision == middleware.DecisionDeny {
require.NotNil(t, out.DenyReason, "deny reason must be set for %s", tt.name)
assert.Equal(t, 403, out.DenyStatus, "deny status must be 403 for %s", tt.name)
assert.Equal(t, tt.denyCode, out.DenyReason.Code, "deny code for %s", tt.name)
}
})
}
}

View File

@@ -253,9 +253,7 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) {
if c := strings.LastIndex(rest, ":"); c >= 0 {
model, action = rest[:c], rest[c+1:]
}
if at := strings.Index(model, "@"); at >= 0 {
model = model[:at]
}
model = llm.NormalizeVertexModel(model)
if model == "" {
return vertexRequest{}, false
}
@@ -276,24 +274,44 @@ func vertexPublisherVendor(publisher string) string {
}
}
// vertexCountTokensModel is the pseudo-model id of the Vertex token-count
// endpoint (…/models/count-tokens:rawPredict). It is the one Vertex request
// shape whose real model travels in the JSON body rather than the URL path.
const vertexCountTokensModel = "count-tokens"
// invokeVertex emits the model/vendor/session/prompt for a Vertex publisher
// request, using the publisher's parser to read the (vendor-native) body.
func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *middleware.Output {
out := &middleware.Output{Decision: middleware.DecisionAllow}
vendor := vertexPublisherVendor(vx.publisher)
md := []middleware.KV{}
if vendor != "" {
md = append(md, middleware.KV{Key: middleware.KeyLLMProvider, Value: vendor})
}
md = append(md, middleware.KV{Key: middleware.KeyLLMModel, Value: vx.model})
md = append(md, middleware.KV{Key: middleware.KeyLLMStream, Value: strconv.FormatBool(vx.stream)})
var parser llm.Parser
if vendor != "" {
parser, _ = llm.ParserByName(vendor)
}
model := vx.model
// Token counting carries the literal "count-tokens" pseudo-model in the URL
// and the real model in the body (Claude Code issues one such call per
// request). Resolve the body model so the guardrail allowlist and the
// router evaluate the actual model — otherwise every count-tokens call
// under an allowlist denies as model_blocked. A body the parser cannot
// read keeps the pseudo-model, which fails closed downstream.
if model == vertexCountTokensModel && parser != nil {
if facts, err := parser.ParseRequest(in.Body); err == nil && facts.Model != "" {
if bodyModel := llm.NormalizeVertexModel(facts.Model); bodyModel != "" {
model = bodyModel
}
}
}
md := []middleware.KV{}
if vendor != "" {
md = append(md, middleware.KV{Key: middleware.KeyLLMProvider, Value: vendor})
}
md = append(md, middleware.KV{Key: middleware.KeyLLMModel, Value: model})
md = append(md, middleware.KV{Key: middleware.KeyLLMStream, Value: strconv.FormatBool(vx.stream)})
sessionID := sessionIDFromHeaders(in.Headers)
if sessionID == "" && parser != nil {
sessionID = parser.ExtractSessionID(in.Body)

View File

@@ -564,6 +564,14 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
return true
}
// Vertex request models likewise reach the router with the "@version"
// suffix already stripped by the parser, while the operator may register
// the raw versioned id the Google console documents (e.g.
// "claude-opus-4-6@20250514"). Normalize the candidate so both spellings
// of the same model match.
if route.Vertex && llm.NormalizeVertexModel(candidate) == model {
return true
}
}
return false
}

View File

@@ -0,0 +1,64 @@
package llm_router
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/proxy/internal/middleware"
)
// TestRouteClaimsModel_VertexNormalizesCandidate guards the Vertex analog of
// the native-Bedrock routing gap: the request model reaches the router already
// normalized (the parser strips the "@version" suffix from the URL path), so a
// provider registered with the raw versioned id must still match.
func TestRouteClaimsModel_VertexNormalizesCandidate(t *testing.T) {
route := ProviderRoute{Vertex: true, Models: []string{"claude-opus-4-6@20250514"}}
assert.True(t, routeClaimsModel(route, "claude-opus-4-6"),
"raw @version Vertex model must match the normalized request model")
assert.False(t, routeClaimsModel(route, "claude-haiku-4-5"),
"a model outside the provider's list must not match")
bare := ProviderRoute{Vertex: true, Models: []string{"claude-opus-4-6"}}
assert.True(t, routeClaimsModel(bare, "claude-opus-4-6"),
"a bare catalog id keeps matching exactly")
// Non-Vertex routes keep exact matching for "@"-suffixed candidates.
direct := ProviderRoute{Models: []string{"claude-opus-4-6@20250514"}}
assert.False(t, routeClaimsModel(direct, "claude-opus-4-6"),
"non-Vertex routes must not normalize @version candidates")
}
// TestRouter_VertexUnversionedModelRoutes replays the customer-reported request
// end to end through the router: an unversioned model id at the global location
// (…/models/claude-opus-4-6:rawPredict) must route on a Vertex provider whose
// models were registered with the raw "@version" ids.
func TestRouter_VertexUnversionedModelRoutes(t *testing.T) {
route := vertexRoute()
route.Models = []string{"claude-opus-4-6@20250514", "claude-sonnet-4-5@20250929"}
mw := New(Config{Providers: []ProviderRoute{route}})
in := pathRoutedInput(
"/v1/projects/corp-gcp-it-all-claude/locations/global/publishers/anthropic/models/claude-opus-4-6:rawPredict",
"anthropic",
"claude-opus-4-6",
)
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionAllow, out.Decision,
"unversioned request model must route on a provider registered with @version ids")
denied := pathRoutedInput(
"/v1/projects/corp-gcp-it-all-claude/locations/global/publishers/anthropic/models/claude-haiku-4-5:rawPredict",
"anthropic",
"claude-haiku-4-5",
)
out, err = mw.Invoke(context.Background(), denied)
require.NoError(t, err)
assert.Equal(t, middleware.DecisionDeny, out.Decision,
"a model outside the provider's registered list must still deny")
require.NotNil(t, out.DenyReason)
assert.Equal(t, denyCodeNotRoutable, out.DenyReason.Code)
}