From a69fe9ba8cc246d33c8686b769d846215e5daa7c Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Wed, 22 Jul 2026 07:32:26 +0000 Subject: [PATCH] [proxy] normalize Vertex model ids in routing, guardrail, and token counting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- proxy/internal/llm/vertex_model.go | 14 +++ proxy/internal/llm/vertex_model_test.go | 20 +++++ .../builtin/llm_guardrail/middleware.go | 6 ++ .../builtin/llm_guardrail/middleware_test.go | 19 ++++ .../guardrail_allowlist_test.go | 90 +++++++++++++++++++ .../builtin/llm_request_parser/middleware.go | 34 ++++--- .../builtin/llm_router/middleware.go | 12 +-- .../builtin/llm_router/vertex_route_test.go | 56 ++++++++++++ 8 files changed, 236 insertions(+), 15 deletions(-) create mode 100644 proxy/internal/llm/vertex_model.go create mode 100644 proxy/internal/llm/vertex_model_test.go create mode 100644 proxy/internal/middleware/builtin/llm_router/vertex_route_test.go diff --git a/proxy/internal/llm/vertex_model.go b/proxy/internal/llm/vertex_model.go new file mode 100644 index 000000000..369a1213b --- /dev/null +++ b/proxy/internal/llm/vertex_model.go @@ -0,0 +1,14 @@ +package llm + +import "strings" + +// NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id +// so it matches the catalog key, e.g. "claude-sonnet-4-5@20250929" -> +// "claude-sonnet-4-5". Shared by the parser, router, and guardrail so both +// spellings compare equal. +func NormalizeVertexModel(modelID string) string { + if at := strings.Index(modelID, "@"); at >= 0 { + return modelID[:at] + } + return modelID +} diff --git a/proxy/internal/llm/vertex_model_test.go b/proxy/internal/llm/vertex_model_test.go new file mode 100644 index 000000000..83f5f8b5e --- /dev/null +++ b/proxy/internal/llm/vertex_model_test.go @@ -0,0 +1,20 @@ +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", + "claude-opus-4-6": "claude-opus-4-6", // bare id passes through + "text-embedding-005@001": "text-embedding-005", + "@cf/meta/llama-3-8b": "", // leading "@" -> empty (treated as no model) + } + for in, want := range cases { + require.Equal(t, want, NormalizeVertexModel(in), "normalize %q", in) + } +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go index eded877ac..5605ac93b 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go @@ -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,11 @@ func (m *Middleware) modelInAllowlist(model string) bool { if allowed == normalised { return true } + // Accept a Vertex entry stored with its "@version" suffix against the + // suffix-stripped request model. Entries without "@" stay exact. + if v := llm.NormalizeVertexModel(allowed); v != "" && v != allowed && v == normalised { + return true + } } return false } diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go index cd7e256dd..b936e7a2a 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -70,6 +70,25 @@ func TestAllowlistMatchAllows(t *testing.T) { assert.Equal(t, middleware.DecisionAllow, out.Decision, "model in allowlist must be allowed") } +// A Vertex "@version" allowlist entry must match the version-stripped request +// model the parser emits. +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") + + 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( diff --git a/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go b/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go index 0074411cc..eab51d6fe 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/guardrail_allowlist_test.go @@ -104,3 +104,93 @@ func TestModelAllowlist_URLRoutedProviders(t *testing.T) { }) } } + +// TestModelAllowlist_VertexRequestShapes replays the Vertex request shapes an +// Anthropic SDK client sends (model in the URL path, optionally unversioned, +// plus the count-tokens body-model endpoint) against bare and "@version" +// allowlists. URLs mirror a customer-reported request. +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 customer report: a Sonnet-only allowlist must block Opus. + 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 body model: 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) + } + }) + } +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index 64ca04e6a..207387aba 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -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,40 @@ func vertexPublisherVendor(publisher string) string { } } +// vertexCountTokensModel is the pseudo-model of the Vertex token-count endpoint, +// the one Vertex shape whose real model lives in the body, not 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 + // count-tokens carries its real model in the body; resolve it so the + // guardrail and router evaluate the actual model. An unreadable body keeps + // the pseudo-model and 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) diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 2d987eef6..2dc300f25 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -556,14 +556,16 @@ func routeClaimsModel(route ProviderRoute, model string) bool { if candidate == model { return true } - // Bedrock request models reach the router already normalized (the parser - // strips the region / inference-profile prefix and version suffix), but - // the operator may register the raw inference-profile id (e.g. - // "us.anthropic.claude-haiku-4-5"). Normalize the candidate so both sides - // compare equal; otherwise a native Bedrock request denies as not-routable. + // The request model is already normalized by the parser, but the operator + // may register the raw id (Bedrock inference-profile, Vertex "@version"). + // Normalize the candidate so both spellings match; else it denies as + // not-routable. if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model { return true } + if route.Vertex && llm.NormalizeVertexModel(candidate) == model { + return true + } } return false } diff --git a/proxy/internal/middleware/builtin/llm_router/vertex_route_test.go b/proxy/internal/middleware/builtin/llm_router/vertex_route_test.go new file mode 100644 index 000000000..24866b9b9 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/vertex_route_test.go @@ -0,0 +1,56 @@ +package llm_router + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// A Vertex route registered with the raw "@version" id must match the +// version-stripped request model; non-Vertex routes stay exact. +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")) + assert.False(t, routeClaimsModel(route, "claude-haiku-4-5")) + + bare := ProviderRoute{Vertex: true, Models: []string{"claude-opus-4-6"}} + assert.True(t, routeClaimsModel(bare, "claude-opus-4-6")) + + 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: +// an unversioned model id must route on a provider registered with "@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 registered list must still deny") + require.NotNil(t, out.DenyReason) + assert.Equal(t, denyCodeNotRoutable, out.DenyReason.Code) +}