diff --git a/proxy/internal/llm/vertex_model.go b/proxy/internal/llm/vertex_model.go new file mode 100644 index 000000000..44eb59b48 --- /dev/null +++ b/proxy/internal/llm/vertex_model.go @@ -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 +} diff --git a/proxy/internal/llm/vertex_model_test.go b/proxy/internal/llm/vertex_model_test.go new file mode 100644 index 000000000..58f500de4 --- /dev/null +++ b/proxy/internal/llm/vertex_model_test.go @@ -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) + } +} diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go index eded877ac..21999b7f9 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,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 } diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go index cd7e256dd..701369499 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -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( 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..95e07b049 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,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) + } + }) + } +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index 64ca04e6a..52fd5e982 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,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) diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 2d987eef6..1f22d757b 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -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 } 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..1a7312e3a --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/vertex_route_test.go @@ -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) +}