mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-22 16:31:28 +02:00
[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:
17
proxy/internal/llm/vertex_model.go
Normal file
17
proxy/internal/llm/vertex_model.go
Normal 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
|
||||
}
|
||||
23
proxy/internal/llm/vertex_model_test.go
Normal file
23
proxy/internal/llm/vertex_model_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user