mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-05 14:31:29 +02:00
[management,proxy] Serve guardrail allowlists of declared model ids (#7389)
After #7221, guardrail allowlists built from a path-style provider's declared model ids (Bedrock, Vertex) stopped working: the raw region/version form was compared against the parser's canonical id, so the agent config advertised an empty model list and requests for the allowlisted model were refused. Make every allowlist compare provider-aware, keyed on the destination provider's catalog id: the agent config, the policy gate, and the synthesized proxy allowlists match an entry on both its verbatim and canonical form — Bedrock's strip only under bedrock_api, Vertex's only under vertex_ai_api, verbatim everywhere else, so a plain provider's suffixed entries never widen. The router's claim compare learns the Vertex @version strip. New e2e, realstore, and unit tests reproduce both regressions and pin the fix.
This commit is contained in:
143
e2e/agentnetwork/agent_config_test.go
Normal file
143
e2e/agentnetwork/agent_config_test.go
Normal file
@@ -0,0 +1,143 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// joinGroup places the PAT's own user into the group so caller-scoped answers
|
||||
// (GET /api/agent-network/agent-config) see the policies sourced from it, and
|
||||
// restores the previous auto-groups on cleanup. Self-service updates of one's
|
||||
// own auto_groups are permitted for every role, so this needs no second user.
|
||||
func joinGroup(t *testing.T, ctx context.Context, groupID string) {
|
||||
t.Helper()
|
||||
me, err := srv.API().Users.Current(ctx)
|
||||
require.NoError(t, err, "read current user")
|
||||
before := append([]string(nil), me.AutoGroups...)
|
||||
_, err = srv.API().Users.Update(ctx, me.Id, api.PutApiUsersUserIdJSONRequestBody{
|
||||
Role: me.Role,
|
||||
IsBlocked: me.IsBlocked,
|
||||
AutoGroups: append(append([]string(nil), before...), groupID),
|
||||
})
|
||||
require.NoError(t, err, "add the caller to the policy source group")
|
||||
t.Cleanup(func() {
|
||||
_, _ = srv.API().Users.Update(context.Background(), me.Id, api.PutApiUsersUserIdJSONRequestBody{
|
||||
Role: me.Role,
|
||||
IsBlocked: me.IsBlocked,
|
||||
AutoGroups: before,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// configProvider returns the agent-config entry for the named provider, nil
|
||||
// when the answer does not offer it. The suite shares one account, so other
|
||||
// tests' fixtures may add unrelated providers to the caller's answer.
|
||||
func configProvider(cfg api.AgentNetworkAgentConfig, name string) *api.AgentNetworkAgentConfigProvider {
|
||||
for i := range cfg.Providers {
|
||||
if cfg.Providers[i].Name == name {
|
||||
return &cfg.Providers[i]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// TestAgentConfigAllowlistOfDeclaredModels reproduces the post-#7221 field
|
||||
// report: a provider carrying a declared model set plus a policy guardrail
|
||||
// whose allowlist holds those same declared ids must advertise the models on
|
||||
// GET /api/agent-network/agent-config — the guardrail was built FROM the
|
||||
// provider's model list (the dashboard's allowlist picker persists the
|
||||
// declared ids verbatim), so nothing about the setup excludes them.
|
||||
//
|
||||
// The plain case passes today. The path-style case (Bedrock; Vertex has the
|
||||
// same shape) fails: the declared id is compared through the proxy parser's
|
||||
// canonical form (region prefix and version suffix stripped) while the
|
||||
// allowlist entry is not, so the raw-vs-raw pair never intersects and the
|
||||
// caller sees an empty model list. The same one-sided normalization sits in
|
||||
// policyPermitsModel, so the proxy also denies the model at request time —
|
||||
// the guardrail meant to allow exactly this model turns it off end to end.
|
||||
func TestAgentConfigAllowlistOfDeclaredModels(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
catalogID string
|
||||
upstream string
|
||||
declared string
|
||||
}{
|
||||
{
|
||||
name: "plain-declared-id",
|
||||
catalogID: "openai_api",
|
||||
upstream: "https://api.openai.com",
|
||||
declared: "gpt-4o-mini",
|
||||
},
|
||||
{
|
||||
// The operator declares the id AWS issues — region-prefixed
|
||||
// inference profile with a version suffix — and the allowlist
|
||||
// picker copies it as-is.
|
||||
name: "bedrock-declared-id",
|
||||
catalogID: "bedrock_api",
|
||||
upstream: "https://bedrock-runtime.eu-central-1.amazonaws.com",
|
||||
declared: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
tc := tc
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-agentcfg-" + tc.name})
|
||||
require.NoError(t, err, "create source group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
joinGroup(t, ctx, grp.Id)
|
||||
|
||||
providerName := "e2e-agentcfg-" + tc.name
|
||||
prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{
|
||||
Name: providerName,
|
||||
ProviderId: tc.catalogID,
|
||||
UpstreamUrl: tc.upstream,
|
||||
ApiKey: ptr("sk-dummy-e2e-key"),
|
||||
Enabled: ptr(true),
|
||||
Models: &[]api.AgentNetworkProviderModel{{Id: tc.declared, InputPer1k: 0.001, OutputPer1k: 0.002}},
|
||||
})
|
||||
require.NoError(t, err, "create provider")
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) })
|
||||
|
||||
// Allowlist exactly the declared model, the way the dashboard
|
||||
// builds a guardrail from the provider's model list.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-agentcfg-" + tc.name
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = []string{tc.declared}
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-agentcfg-" + tc.name,
|
||||
Enabled: ptr(true),
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: []string{prov.Id},
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
cfg, err := srv.GetAgentConfig(ctx)
|
||||
require.NoError(t, err, "read the caller-scoped agent config")
|
||||
require.True(t, cfg.Configured, "the account endpoint is bootstrapped by TestMain")
|
||||
|
||||
entry := configProvider(cfg, providerName)
|
||||
require.NotNil(t, entry, "the policy authorizes the caller for the provider, so it must be offered")
|
||||
assert.False(t, entry.AllModelsAllowed, "an allowlist guardrail restricts the provider")
|
||||
assert.Equal(t, []string{tc.declared}, entry.Models,
|
||||
"the allowlist holds the provider's own declared id, so that model must be advertised")
|
||||
})
|
||||
}
|
||||
}
|
||||
132
e2e/agentnetwork/guardrail_declared_ids_test.go
Normal file
132
e2e/agentnetwork/guardrail_declared_ids_test.go
Normal file
@@ -0,0 +1,132 @@
|
||||
//go:build e2e
|
||||
|
||||
package agentnetwork
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/e2e/harness"
|
||||
"github.com/netbirdio/netbird/shared/management/http/api"
|
||||
)
|
||||
|
||||
// TestModelAllowlistOfDeclaredIDsServed drives the setup an operator actually
|
||||
// builds for a path-routed provider: the models are declared in the form the
|
||||
// vendor issues (Bedrock's region-prefixed, versioned inference-profile id;
|
||||
// Vertex's model@version), and the guardrail allowlist is built from that
|
||||
// declared list — the dashboard's allowlist picker persists the declared ids
|
||||
// verbatim. A request for the declared model must be served end to end, and a
|
||||
// model outside the allowlist must still be denied.
|
||||
//
|
||||
// TestModelAllowlistEnforced never caught this because it registers and
|
||||
// allowlists the pre-normalized catalog form (see the catalogModel comment
|
||||
// there and the one in providerRequest: "register the normalized form here or
|
||||
// routing fails as model_not_routable") — the harness encoded the
|
||||
// canonicalization workaround instead of the shape operators configure.
|
||||
func TestModelAllowlistOfDeclaredIDsServed(t *testing.T) {
|
||||
var providers []providerCase
|
||||
for _, pc := range availableProviders() {
|
||||
if pc.kind == harness.WireBedrock || pc.kind == harness.WireVertex {
|
||||
providers = append(providers, pc)
|
||||
}
|
||||
}
|
||||
if len(providers) == 0 {
|
||||
t.Skip("no path-routed provider keys set (AWS_BEARER_TOKEN_BEDROCK / GOOGLE_VERTEX_*); source ~/.llm-keys")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-declared-allowlist"})
|
||||
require.NoError(t, err, "create group")
|
||||
t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) })
|
||||
|
||||
ephemeral := false
|
||||
sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{
|
||||
Name: "e2e-declared-allowlist-client",
|
||||
Type: "reusable",
|
||||
ExpiresIn: 86400,
|
||||
UsageLimit: 0,
|
||||
AutoGroups: []string{grp.Id},
|
||||
Ephemeral: &ephemeral,
|
||||
})
|
||||
require.NoError(t, err, "mint setup key")
|
||||
t.Cleanup(func() { _ = srv.API().SetupKeys.Delete(context.Background(), sk.Id) })
|
||||
|
||||
// Providers declaring the raw vendor-issued model id — NOT the normalized
|
||||
// catalog form providerRequest would register.
|
||||
ids := make([]string, 0, len(providers))
|
||||
declared := make([]string, 0, len(providers))
|
||||
for _, pc := range providers {
|
||||
req := providerRequest(pc)
|
||||
req.Models = &[]api.AgentNetworkProviderModel{{Id: pc.model, InputPer1k: 0.001, OutputPer1k: 0.002}}
|
||||
prov, perr := srv.CreateProvider(ctx, req)
|
||||
require.NoError(t, perr, "create provider %s", pc.name)
|
||||
id := prov.Id
|
||||
ids = append(ids, id)
|
||||
declared = append(declared, pc.model)
|
||||
t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) })
|
||||
}
|
||||
|
||||
// Guardrail allowlisting the declared ids verbatim, the way the dashboard
|
||||
// builds an allowlist from the providers' model lists.
|
||||
var gr api.AgentNetworkGuardrailRequest
|
||||
gr.Name = "e2e-declared-allowlist"
|
||||
gr.Checks.ModelAllowlist.Enabled = true
|
||||
gr.Checks.ModelAllowlist.Models = declared
|
||||
guard, err := srv.CreateGuardrail(ctx, gr)
|
||||
require.NoError(t, err, "create guardrail")
|
||||
t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) })
|
||||
|
||||
enabled := true
|
||||
pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{
|
||||
Name: "e2e-declared-allowlist",
|
||||
Enabled: &enabled,
|
||||
SourceGroups: []string{grp.Id},
|
||||
DestinationProviderIds: ids,
|
||||
GuardrailIds: &[]string{guard.Id},
|
||||
})
|
||||
require.NoError(t, err, "create policy")
|
||||
t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) })
|
||||
|
||||
settings, err := srv.GetSettings(ctx)
|
||||
require.NoError(t, err, "read settings for endpoint")
|
||||
require.NotEmpty(t, settings.Endpoint, "agent-network endpoint must be assigned")
|
||||
|
||||
proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-proxy-declared-allowlist")
|
||||
require.NoError(t, err, "mint proxy token via CLI")
|
||||
px, err := harness.StartProxy(ctx, srv, proxyToken)
|
||||
require.NoError(t, err, "start proxy")
|
||||
t.Cleanup(func() { _ = px.Terminate(context.Background()) })
|
||||
|
||||
cl, err := harness.StartClient(ctx, srv, sk.Key)
|
||||
require.NoError(t, err, "start client")
|
||||
t.Cleanup(func() { _ = cl.Terminate(context.Background()) })
|
||||
|
||||
require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management")
|
||||
// Probe first: the GET resolves the endpoint (DNS error fails) and its first packet wakes the lazy proxy peer, so WaitProxyPeer sees it connected; any HTTP status counts.
|
||||
proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint)
|
||||
require.NoError(t, err, "resolve agent-network endpoint to proxy IP")
|
||||
if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil {
|
||||
t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background()))
|
||||
}
|
||||
|
||||
for _, pc := range providers {
|
||||
pc := pc
|
||||
t.Run(pc.name, func(t *testing.T) {
|
||||
// The model the operator declared and allowlisted is served end to
|
||||
// end: the route must claim it and the guardrail must permit it,
|
||||
// both through the canonicalization the parser applies at request
|
||||
// time — whatever id form the operator configured.
|
||||
assert.Equal(t, 200, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, pc.model),
|
||||
"the declared and allowlisted model must be served for %s", pc.name)
|
||||
// A model outside the allowlist stays denied.
|
||||
assert.Equal(t, 403, sendModel(ctx, t, cl, settings.Endpoint, proxyIP, pc, disallowedModel(pc)),
|
||||
"model outside the allowlist must be denied for %s", pc.name)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -100,6 +100,15 @@ func (c *Combined) SetProviderEnabled(ctx context.Context, id string, enabled bo
|
||||
return err
|
||||
}
|
||||
|
||||
// GetAgentConfig returns the caller-scoped self-service connection config —
|
||||
// the answer the dashboard's "Connect Agent" view renders for the PAT's user.
|
||||
// Providers appear only when the caller's own groups intersect an enabled
|
||||
// policy's source groups, so tests must place the PAT user into the policy's
|
||||
// source group first (via the Users API auto-groups).
|
||||
func (c *Combined) GetAgentConfig(ctx context.Context) (api.AgentNetworkAgentConfig, error) {
|
||||
return anRequest[api.AgentNetworkAgentConfig](ctx, c, http.MethodGet, "/api/agent-network/agent-config", nil)
|
||||
}
|
||||
|
||||
// CreatePolicy creates an agent-network policy.
|
||||
func (c *Combined) CreatePolicy(ctx context.Context, req api.AgentNetworkPolicyRequest) (api.AgentNetworkPolicy, error) {
|
||||
return anRequest[api.AgentNetworkPolicy](ctx, c, http.MethodPost, "/api/agent-network/policies", req)
|
||||
|
||||
@@ -179,6 +179,10 @@ func policiesForProvider(policies []*types.Policy, providerID string) []*types.P
|
||||
// only claims declared models, so an allowlisted-but-undeclared model is
|
||||
// unreachable and must not be advertised. With no declared models the
|
||||
// router claims every model, so the allowlist union stands alone.
|
||||
// Allowlist entries and declared ids both compare through the canonical
|
||||
// id the proxy's parser emits, so an allowlist may hold either form: the
|
||||
// raw declared id the dashboard's picker copies from the provider, or
|
||||
// the stripped id the parser matches at request time.
|
||||
func effectiveModelsForProvider(provider *types.Provider, policies []*types.Policy, guardrailsByID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := true
|
||||
union := make([]string, 0)
|
||||
@@ -192,7 +196,7 @@ func effectiveModelsForProvider(provider *types.Provider, policies []*types.Poli
|
||||
}
|
||||
policyRestricted = true
|
||||
for _, model := range g.Checks.ModelAllowlist.Models {
|
||||
key := normaliseModelID(model)
|
||||
key := canonicalModelKey(provider.ProviderID, model)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
@@ -221,17 +225,26 @@ func effectiveModelsForProvider(provider *types.Provider, policies []*types.Poli
|
||||
for _, id := range declared {
|
||||
// Compare through the canonical id the proxy's parser emits — a
|
||||
// Bedrock declaration may carry the region/version form
|
||||
// ("eu.anthropic.claude-...-v1:0") while the allowlist holds the
|
||||
// stripped id the parser matches at request time, and the raw
|
||||
// forms would never intersect. The declared id itself is what
|
||||
// gets advertised, matching the router's route claim.
|
||||
if _, ok := seen[normaliseModelID(normalizePricingModelID(provider.ProviderID, id))]; ok {
|
||||
// ("eu.anthropic.claude-...-v1:0") that the parser strips at
|
||||
// request time, and the raw forms would never intersect. The
|
||||
// declared id itself is what gets advertised, matching the
|
||||
// router's route claim.
|
||||
if _, ok := seen[canonicalModelKey(provider.ProviderID, id)]; ok {
|
||||
out = append(out, id)
|
||||
}
|
||||
}
|
||||
return false, out
|
||||
}
|
||||
|
||||
// canonicalModelKey builds the compare key for a model id: lowercased,
|
||||
// trimmed, and canonicalized through the provider-aware normalization the
|
||||
// proxy's parser applies. Lowercase/trim comes FIRST — the path-style
|
||||
// strippers anchor on a lowercase id's tail, so a trailing space or a
|
||||
// case-variant geography/version would otherwise survive into the key.
|
||||
func canonicalModelKey(catalogProviderID, id string) string {
|
||||
return normaliseModelID(normalizePricingModelID(catalogProviderID, normaliseModelID(id)))
|
||||
}
|
||||
|
||||
// providerModelsByID maps effective model ids (as effectiveModelsForProvider
|
||||
// returns them) back onto the operator's declared entries, keeping the
|
||||
// declared casing and prices. With no operator declaration the ids are the
|
||||
|
||||
@@ -144,6 +144,54 @@ func TestAgentConfig_RealStore_AllowlistMatchesBedrockDeclaredIDsCanonically(t *
|
||||
"the allowlisted canonical id must admit the declared region/version form, and only it")
|
||||
}
|
||||
|
||||
func TestAgentConfig_RealStore_AllowlistHoldsRawDeclaredIDs(t *testing.T) {
|
||||
// The dashboard's allowlist picker copies the provider's declared ids
|
||||
// verbatim, so for path-style providers the allowlist carries the
|
||||
// region/version form rather than the canonical id the parser emits.
|
||||
// Both forms must admit the declared model.
|
||||
cases := []struct {
|
||||
name string
|
||||
catalogID string
|
||||
declared string
|
||||
allowlist string
|
||||
}{
|
||||
{"bedrock", "bedrock_api", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", ""},
|
||||
{"vertex", "vertex_ai_api", "claude-sonnet-4-5@20250929", ""},
|
||||
// The geography/version strippers anchor on a lowercase tail, so a
|
||||
// case-variant entry must be lowercased before canonicalization or
|
||||
// the prefix and suffix survive into the compare key.
|
||||
{"bedrock-case-variant", "bedrock_api", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
" EU.Anthropic.Claude-Sonnet-4-5-20250929-V1:0 "},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
mgr, s := newAgentConfigTestMgr(t)
|
||||
ctx := context.Background()
|
||||
|
||||
allowlisted := tc.allowlist
|
||||
if allowlisted == "" {
|
||||
allowlisted = tc.declared
|
||||
}
|
||||
require.NoError(t, s.SaveAgentNetworkSettings(ctx, newSynthTestSettings()))
|
||||
provider := newSynthTestProvider()
|
||||
provider.ProviderID = tc.catalogID
|
||||
provider.Name = tc.name
|
||||
provider.Models = []types.ProviderModel{{ID: tc.declared}}
|
||||
require.NoError(t, s.SaveAgentNetworkProvider(ctx, provider))
|
||||
require.NoError(t, s.SaveAgentNetworkGuardrail(ctx, newSetupTestGuardrail("guard-1", allowlisted)))
|
||||
require.NoError(t, s.SaveAgentNetworkPolicy(ctx, newSynthTestPolicy(provider.ID, "grp-eng", "guard-1")))
|
||||
|
||||
setup, err := mgr.agentConfigForGroups(ctx, testAccountID, []string{"grp-eng"})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, setup.Providers, 1)
|
||||
p := setup.Providers[0]
|
||||
assert.False(t, p.AllModelsAllowed)
|
||||
assert.Equal(t, []string{tc.declared}, p.Models,
|
||||
"an allowlist holding the raw declared id must admit that declared model")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentConfig_RealStore_UnrestrictedPolicyWinsOverRestricted(t *testing.T) {
|
||||
mgr, s := newAgentConfigTestMgr(t)
|
||||
ctx := context.Background()
|
||||
|
||||
@@ -164,23 +164,12 @@ func (m *managerImpl) SelectPolicyForRequest(ctx context.Context, in PolicySelec
|
||||
}
|
||||
candidates := filterApplicablePolicies(policies, in)
|
||||
|
||||
// Model-allowlist gate scoped to the matched policies: keep candidates whose
|
||||
// guardrails permit the model (none enabled = unrestricted), deny when
|
||||
// policies apply but none permits it. Skip the load when none has a guardrail.
|
||||
if len(candidates) > 0 && anyPolicyHasGuardrails(candidates) {
|
||||
guardrailsByID, gErr := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if gErr != nil {
|
||||
return nil, gErr
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model)
|
||||
if len(permitted) == 0 {
|
||||
return &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
candidates = permitted
|
||||
candidates, denied, err := m.applyModelGate(ctx, in, candidates)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if denied != nil {
|
||||
return denied, nil
|
||||
}
|
||||
|
||||
// Prefetch every consumption counter the ceiling + candidate policies will
|
||||
@@ -285,6 +274,59 @@ func anyPolicyHasGuardrails(policies []*types.Policy) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// applyModelGate is the model-allowlist gate scoped to the matched policies:
|
||||
// it keeps the candidates whose guardrails permit the model (none enabled =
|
||||
// unrestricted) and returns a deny result when policies apply but none
|
||||
// permits it. The guardrail load is skipped when no candidate references a
|
||||
// guardrail, and the provider's catalog id — which picks the model-id
|
||||
// normalizer — is resolved only when a candidate actually restricts models:
|
||||
// with no enabled allowlist every candidate is unrestricted, and a
|
||||
// provider-store failure must not fail a request the gate would have waved
|
||||
// through.
|
||||
func (m *managerImpl) applyModelGate(ctx context.Context, in PolicySelectionInput, candidates []*types.Policy) ([]*types.Policy, *PolicySelectionResult, error) {
|
||||
if len(candidates) == 0 || !anyPolicyHasGuardrails(candidates) {
|
||||
return candidates, nil, nil
|
||||
}
|
||||
guardrailsByID, err := m.loadGuardrailsByID(ctx, in.AccountID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if !anyEnabledModelAllowlist(candidates, guardrailsByID) {
|
||||
return candidates, nil, nil
|
||||
}
|
||||
catalogID, err := m.providerCatalogID(ctx, in.AccountID, in.ProviderID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
permitted := filterModelPermittedPolicies(candidates, guardrailsByID, in.Model, catalogID)
|
||||
if len(permitted) == 0 {
|
||||
return nil, &PolicySelectionResult{
|
||||
Allow: false,
|
||||
DenyCode: denyCodeModelBlocked,
|
||||
DenyReason: modelBlockedReason(in.Model),
|
||||
}, nil
|
||||
}
|
||||
return permitted, nil, nil
|
||||
}
|
||||
|
||||
// anyEnabledModelAllowlist reports whether any policy references a guardrail
|
||||
// whose model allowlist is enabled — the only case the model gate restricts
|
||||
// anything. Disabled allowlists, stale guardrail references, and guardrails
|
||||
// carrying only other checks all leave every candidate unrestricted.
|
||||
func anyEnabledModelAllowlist(policies []*types.Policy, byID map[string]*types.Guardrail) bool {
|
||||
for _, p := range policies {
|
||||
if p == nil {
|
||||
continue
|
||||
}
|
||||
for _, gID := range p.GuardrailIDs {
|
||||
if g, ok := byID[gID]; ok && g != nil && g.Checks.ModelAllowlist.Enabled {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// loadGuardrailsByID loads the account's guardrails indexed by ID. Used by the
|
||||
// model-allowlist gate to resolve each candidate policy's attached guardrails.
|
||||
func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string) (map[string]*types.Guardrail, error) {
|
||||
@@ -301,12 +343,33 @@ func (m *managerImpl) loadGuardrailsByID(ctx context.Context, accountID string)
|
||||
return byID, nil
|
||||
}
|
||||
|
||||
// providerCatalogID resolves a provider record id to its catalog provider
|
||||
// id, the key the model-id normalizers are picked by. A missing provider
|
||||
// resolves to the empty catalog id — the compare then runs verbatim-only,
|
||||
// which can never widen an allowlist — while a store failure propagates
|
||||
// rather than degrading a security decision.
|
||||
func (m *managerImpl) providerCatalogID(ctx context.Context, accountID, providerID string) (string, error) {
|
||||
if providerID == "" {
|
||||
return "", nil
|
||||
}
|
||||
provider, err := m.store.GetAgentNetworkProviderByID(ctx, store.LockingStrengthNone, accountID, providerID)
|
||||
switch {
|
||||
case err == nil:
|
||||
return provider.ProviderID, nil
|
||||
case isNotFound(err):
|
||||
return "", nil
|
||||
default:
|
||||
return "", fmt.Errorf("get provider: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// filterModelPermittedPolicies returns the subset of policies whose guardrails
|
||||
// permit the model. Order is preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model string) []*types.Policy {
|
||||
// permit the model on the provider with the given catalog id. Order is
|
||||
// preserved so downstream scoring is unaffected.
|
||||
func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, model, catalogProviderID string) []*types.Policy {
|
||||
out := make([]*types.Policy, 0, len(policies))
|
||||
for _, p := range policies {
|
||||
if policyPermitsModel(p, byID, model) {
|
||||
if policyPermitsModel(p, byID, model, catalogProviderID) {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
@@ -316,8 +379,13 @@ func filterModelPermittedPolicies(policies []*types.Policy, byID map[string]*typ
|
||||
// policyPermitsModel reports whether a policy permits the model. No
|
||||
// allowlist-enabled guardrail = unrestricted (permits any, incl. empty);
|
||||
// otherwise the model must be in the union of its allowlists, so an
|
||||
// empty/undetermined model fails closed.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model string) bool {
|
||||
// empty/undetermined model fails closed. An entry matches on its own
|
||||
// normalised form or, for a path-style provider, its canonical form: the
|
||||
// parser emits the canonical id for path-routed requests, while an
|
||||
// allowlist may hold the raw declared id the dashboard's picker copies
|
||||
// from the provider. The catalog id picks the normalizer, so a plain
|
||||
// provider's entries always compare verbatim.
|
||||
func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model, catalogProviderID string) bool {
|
||||
if p == nil {
|
||||
return false
|
||||
}
|
||||
@@ -333,7 +401,7 @@ func policyPermitsModel(p *types.Policy, byID map[string]*types.Guardrail, model
|
||||
continue
|
||||
}
|
||||
for _, allowed := range g.Checks.ModelAllowlist.Models {
|
||||
if normaliseModelID(allowed) == wanted {
|
||||
if normaliseModelID(allowed) == wanted || canonicalModelKey(catalogProviderID, allowed) == wanted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,12 +6,13 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/mock/gomock"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types"
|
||||
"github.com/netbirdio/netbird/management/server/store"
|
||||
"github.com/netbirdio/netbird/shared/management/status"
|
||||
)
|
||||
|
||||
// guardedPolicy builds an enabled, uncapped policy that authorises sourceGroups
|
||||
@@ -53,6 +54,17 @@ func expectGuardrails(mockStore *store.MockStore, account string, guardrails ...
|
||||
Return(guardrails, nil)
|
||||
}
|
||||
|
||||
// expectProviderCatalog resolves the destination provider to the given
|
||||
// catalog provider id, which picks the model-id normalizer the allowlist
|
||||
// gate compares through. AnyTimes: the lookup runs only when the guardrail
|
||||
// gate is reached.
|
||||
func expectProviderCatalog(mockStore *store.MockStore, account, providerID, catalog string) {
|
||||
mockStore.EXPECT().
|
||||
GetAgentNetworkProviderByID(gomock.Any(), gomock.Any(), account, providerID).
|
||||
Return(&types.Provider{ID: providerID, AccountID: account, ProviderID: catalog}, nil).
|
||||
AnyTimes()
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ModelBlockedByAllowlist proves the authoritative allowlist
|
||||
// decision: a policy authorises the (provider, group) but restricts the model,
|
||||
// and the requested model isn't on the list, so the request is denied.
|
||||
@@ -63,6 +75,7 @@ func TestSelectPolicy_ModelBlockedByAllowlist(t *testing.T) {
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
@@ -86,6 +99,7 @@ func TestSelectPolicy_ModelAllowedByAllowlist(t *testing.T) {
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o", "claude-opus-4"))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -109,6 +123,7 @@ func TestSelectPolicy_CaseInsensitiveModelMatch(t *testing.T) {
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", " GPT-4o "))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -132,6 +147,7 @@ func TestSelectPolicy_UnguardedPolicyIsUnrestricted(t *testing.T) {
|
||||
open := guardedPolicy("pol-open", "acc-1", []string{"grp-eng"}, "prov-1") // no guardrail
|
||||
expectPolicies(mockStore, "acc-1", restricted, open)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -159,6 +175,7 @@ func TestSelectPolicy_AllowlistDoesNotLeakAcrossGroups(t *testing.T) {
|
||||
allowlistGuardrail("g-a", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-b", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
@@ -181,6 +198,7 @@ func TestSelectPolicy_UndeterminedModelFailsClosed(t *testing.T) {
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
@@ -210,6 +228,8 @@ func TestSelectPolicy_DisabledAllowlistDoesNotRestrict(t *testing.T) {
|
||||
}
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", disabled)
|
||||
// Deliberately no provider expectation: with no enabled allowlist the
|
||||
// gate must skip the catalog-id lookup entirely.
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -235,6 +255,7 @@ func TestSelectPolicy_UnionAcrossPolicyGuardrails(t *testing.T) {
|
||||
allowlistGuardrail("g-1", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-2", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -281,6 +302,8 @@ func TestSelectPolicy_MissingGuardrailReferenceTreatedAsUnrestricted(t *testing.
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-missing")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1")
|
||||
// Deliberately no provider expectation: an orphaned guardrail reference
|
||||
// restricts nothing, so the gate must skip the catalog-id lookup.
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -314,6 +337,7 @@ func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
|
||||
allowlistGuardrail("g-restrict", "acc-1", "gpt-4o"),
|
||||
allowlistGuardrail("g-permit", "acc-1", "claude-opus-4"),
|
||||
)
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
@@ -327,3 +351,159 @@ func TestSelectPolicy_PartialCandidatesPermittedAfterModelFilter(t *testing.T) {
|
||||
assert.Equal(t, "pol-small", res.SelectedPolicyID,
|
||||
"the model filter must exclude pol-big before cap scoring")
|
||||
}
|
||||
|
||||
// TestSelectPolicy_RawDeclaredAllowlistPermitsCanonicalModel proves an
|
||||
// allowlist holding the raw vendor-issued id — the form the dashboard's
|
||||
// picker copies from a provider's declared models — permits the request:
|
||||
// the parser emits the path-style canonical id, so the entry must match
|
||||
// through the same canonicalization.
|
||||
func TestSelectPolicy_RawDeclaredAllowlistPermitsCanonicalModel(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
catalog string
|
||||
entry string
|
||||
request string
|
||||
}{
|
||||
{"bedrock raw region/version form", "bedrock_api", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", "anthropic.claude-sonnet-4-5"},
|
||||
{"vertex raw @version form", "vertex_ai_api", "claude-sonnet-4-5@20250929", "claude-sonnet-4-5"},
|
||||
{"vertex raw dated @version form", "vertex_ai_api", "gpt-4o@2024-08-06", "gpt-4o"},
|
||||
{"bedrock raw form with case and whitespace", "bedrock_api", " EU.Anthropic.Claude-Sonnet-4-5-20250929-V1:0 ", "anthropic.claude-sonnet-4-5"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", tc.entry))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", tc.catalog)
|
||||
expectConsumptionBatch(mockStore, nil)
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: tc.request,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, res.Allow, "the raw declared allowlist entry must permit its canonical model")
|
||||
assert.Equal(t, "pol-A", res.SelectedPolicyID)
|
||||
})
|
||||
}
|
||||
|
||||
// A model outside the allowlist stays denied under the same entry shape.
|
||||
t.Run("unrelated canonical model stays denied", func(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "bedrock_api")
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anthropic.claude-opus-4-8",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a model the allowlist never names must stay denied")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
})
|
||||
}
|
||||
|
||||
// TestSelectPolicy_PlainProviderEntriesStayVerbatim proves the canonical-form
|
||||
// compare never relaxes an allowlist on a body-routed provider: its catalog
|
||||
// id selects no normalizer, so a suffix that would be stripped under Bedrock
|
||||
// ("-v2") or Vertex ("@...") stays part of the entry and must NOT also admit
|
||||
// the stripped id — on this provider that is a different model.
|
||||
func TestSelectPolicy_PlainProviderEntriesStayVerbatim(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
entry string
|
||||
request string
|
||||
}{
|
||||
{"a -vN suffix is not a Bedrock version tag here", "claude-3-5-sonnet-v2", "claude-3-5-sonnet"},
|
||||
{"an @word suffix is not a Vertex version tag here", "custom-model@team", "custom-model"},
|
||||
{"an @digits suffix is not a Vertex version tag here", "custom-model@2024", "custom-model"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", tc.entry))
|
||||
expectProviderCatalog(mockStore, "acc-1", "prov-1", "openai_api")
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: tc.request,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "a plain provider's allowlist entry must not widen to its stripped form")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestSelectPolicy_MissingProviderRecordComparesVerbatim proves a provider the
|
||||
// store no longer holds degrades to the verbatim-only compare — the raw entry
|
||||
// still matches itself, and nothing widens — rather than erroring or guessing
|
||||
// a normalizer.
|
||||
func TestSelectPolicy_MissingProviderRecordComparesVerbatim(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"))
|
||||
mockStore.EXPECT().
|
||||
GetAgentNetworkProviderByID(gomock.Any(), gomock.Any(), "acc-1", "prov-1").
|
||||
Return(nil, status.Errorf(status.NotFound, "provider not found")).
|
||||
AnyTimes()
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "anthropic.claude-sonnet-4-5",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, res.Allow, "without the provider record the compare runs verbatim and must not widen")
|
||||
assert.Equal(t, denyCodeModelBlocked, res.DenyCode)
|
||||
}
|
||||
|
||||
// TestSelectPolicy_ProviderLookupErrorPropagates proves a store failure while
|
||||
// resolving the provider's catalog id surfaces as an error — the model gate is
|
||||
// a security decision and must not silently degrade.
|
||||
func TestSelectPolicy_ProviderLookupErrorPropagates(t *testing.T) {
|
||||
ctrl := gomock.NewController(t)
|
||||
mgr, mockStore := newSelectorMgr(t, ctrl)
|
||||
|
||||
policy := guardedPolicy("pol-A", "acc-1", []string{"grp-eng"}, "prov-1", "g-1")
|
||||
expectPolicies(mockStore, "acc-1", policy)
|
||||
expectGuardrails(mockStore, "acc-1", allowlistGuardrail("g-1", "acc-1", "gpt-4o"))
|
||||
mockStore.EXPECT().
|
||||
GetAgentNetworkProviderByID(gomock.Any(), gomock.Any(), "acc-1", "prov-1").
|
||||
Return(nil, errors.New("store unavailable"))
|
||||
|
||||
res, err := mgr.SelectPolicyForRequest(context.Background(), PolicySelectionInput{
|
||||
AccountID: "acc-1",
|
||||
UserID: "user-1",
|
||||
GroupIDs: []string{"grp-eng"},
|
||||
ProviderID: "prov-1",
|
||||
Model: "gpt-4o",
|
||||
})
|
||||
require.Error(t, err, "a provider-lookup failure must surface as an error")
|
||||
assert.Nil(t, res)
|
||||
}
|
||||
|
||||
@@ -211,18 +211,19 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([
|
||||
}
|
||||
|
||||
groupIndex := indexProviderGroups(enabledPolicies)
|
||||
catalogByProvider := catalogIDsByProvider(enabledProviders)
|
||||
|
||||
// The proxy guardrail is a per-provider fail-closed backstop; the
|
||||
// authoritative per-policy/group decision is management's
|
||||
// SelectPolicyForRequest. A provider lands in that map only when every
|
||||
// authorising policy restricts models.
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID)
|
||||
providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID, catalogByProvider)
|
||||
|
||||
// Discovery gets the finer view: per policy rather than flattened per
|
||||
// provider, so a listing can be bounded to what the calling groups may
|
||||
// actually use instead of the union across everyone who reaches the
|
||||
// provider.
|
||||
modelPolicies := buildModelPolicies(enabledPolicies, guardrailsByID)
|
||||
modelPolicies := buildModelPolicies(enabledPolicies, guardrailsByID, catalogByProvider)
|
||||
|
||||
routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex, modelPolicies)
|
||||
if err != nil {
|
||||
@@ -907,7 +908,9 @@ func marshalGuardrailConfig(providerAllowlists map[string][]string, capture Merg
|
||||
// buildProviderAllowlists returns the proxy's per-provider backstop: a provider
|
||||
// is included only when every authorising policy restricts models (their union);
|
||||
// if any leaves it unrestricted it is omitted, so management decides per group.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]string {
|
||||
// Entries carry their provider-specific canonical form alongside the verbatim
|
||||
// one, resolved through catalogByProvider.
|
||||
func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Guardrail, catalogByProvider map[string]string) map[string][]string {
|
||||
type providerAcc struct {
|
||||
models map[string]struct{}
|
||||
anyUnrestricted bool
|
||||
@@ -931,7 +934,7 @@ func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Gu
|
||||
acc.anyUnrestricted = true
|
||||
continue
|
||||
}
|
||||
for _, m := range models {
|
||||
for _, m := range expandModelsForProvider(models, catalogByProvider[providerID]) {
|
||||
acc.models[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
@@ -952,8 +955,10 @@ func buildProviderAllowlists(policies []*types.Policy, byID map[string]*types.Gu
|
||||
}
|
||||
|
||||
// policyModelAllowlist reports whether a policy restricts models (has an
|
||||
// allowlist-enabled guardrail) and the union of allowed models. Models are
|
||||
// verbatim; the proxy factory lowercases/trims them at decode time.
|
||||
// allowlist-enabled guardrail) and the union of allowed models, verbatim.
|
||||
// Consumers expand the entries per destination provider with
|
||||
// expandModelsForProvider — the canonical form is provider-specific — and
|
||||
// the proxy factory lowercases/trims them at decode time.
|
||||
func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bool, []string) {
|
||||
restricted := false
|
||||
var models []string
|
||||
@@ -972,6 +977,45 @@ func policyModelAllowlist(p *types.Policy, byID map[string]*types.Guardrail) (bo
|
||||
return restricted, models
|
||||
}
|
||||
|
||||
// expandModelsForProvider returns the allowlist entries for one destination
|
||||
// provider: each entry verbatim plus, when it differs, its canonical form
|
||||
// under that provider's catalog id — the id the proxy's parser emits at
|
||||
// request time — deduplicated. The proxy-side compares (guardrail backstop,
|
||||
// per-group router rules) then admit an allowlist however the operator
|
||||
// wrote it, raw declared id or canonical, while a plain provider's entries
|
||||
// stay verbatim and can never widen.
|
||||
func expandModelsForProvider(models []string, catalogProviderID string) []string {
|
||||
out := make([]string, 0, len(models))
|
||||
seen := make(map[string]struct{}, len(models))
|
||||
add := func(m string) {
|
||||
if m == "" {
|
||||
return
|
||||
}
|
||||
if _, dup := seen[m]; dup {
|
||||
return
|
||||
}
|
||||
seen[m] = struct{}{}
|
||||
out = append(out, m)
|
||||
}
|
||||
for _, m := range models {
|
||||
add(m)
|
||||
add(canonicalModelKey(catalogProviderID, m))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// catalogIDsByProvider indexes providers' catalog ids by provider record id,
|
||||
// the lookup the per-provider allowlist expansion keys the normalizer on.
|
||||
func catalogIDsByProvider(providers []*types.Provider) map[string]string {
|
||||
out := make(map[string]string, len(providers))
|
||||
for _, p := range providers {
|
||||
if p != nil {
|
||||
out[p.ID] = p.ProviderID
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// buildAccountService composes the per-account gateway Service. The
|
||||
// target carries the noop placeholder URL — the router middleware
|
||||
// rewrites every request to the matched provider's upstream before the
|
||||
@@ -1180,23 +1224,25 @@ type routerModelPolicy struct {
|
||||
// models — a picker full of entries the next request refuses. Keeping the
|
||||
// source groups alongside the models lets the router answer it at request time,
|
||||
// where it knows the caller's groups.
|
||||
func buildModelPolicies(policies []*types.Policy, byID map[string]*types.Guardrail) map[string][]routerModelPolicy {
|
||||
func buildModelPolicies(policies []*types.Policy, byID map[string]*types.Guardrail, catalogByProvider map[string]string) map[string][]routerModelPolicy {
|
||||
out := make(map[string][]routerModelPolicy)
|
||||
for _, p := range policies {
|
||||
if p == nil || len(p.SourceGroups) == 0 {
|
||||
continue
|
||||
}
|
||||
restricted, models := policyModelAllowlist(p, byID)
|
||||
rule := routerModelPolicy{GroupIDs: append([]string(nil), p.SourceGroups...)}
|
||||
if restricted {
|
||||
// Never nil when restricted: an allowlist permitting nothing must
|
||||
// stay distinguishable from no allowlist at all.
|
||||
rule.Models = append([]string{}, models...)
|
||||
}
|
||||
for _, providerID := range p.DestinationProviderIDs {
|
||||
if providerID == "" {
|
||||
continue
|
||||
}
|
||||
rule := routerModelPolicy{GroupIDs: append([]string(nil), p.SourceGroups...)}
|
||||
if restricted {
|
||||
// Never nil when restricted: an allowlist permitting nothing
|
||||
// must stay distinguishable from no allowlist at all. The
|
||||
// expansion is per provider — the canonical form of an entry
|
||||
// depends on the destination's catalog id.
|
||||
rule.Models = append([]string{}, expandModelsForProvider(models, catalogByProvider[providerID])...)
|
||||
}
|
||||
out[providerID] = append(out[providerID], rule)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {"claude-opus-4", "gpt-4o"}}, got,
|
||||
"a provider every policy restricts carries the sorted union of their models")
|
||||
})
|
||||
@@ -43,7 +43,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", nil, "prov-x"), // no guardrail
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a provider reachable by an un-guardrailed policy must be omitted so the proxy treats it as unrestricted")
|
||||
})
|
||||
@@ -52,7 +52,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-disabled"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.NotContains(t, got, "prov-x",
|
||||
"a policy whose only guardrail has a disabled allowlist is unrestricted")
|
||||
})
|
||||
@@ -62,7 +62,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x"),
|
||||
policyForProviders("p2", []string{"g-opus"}, "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"], "prov-x keeps only its own model")
|
||||
assert.Equal(t, []string{"claude-opus-4"}, got["prov-y"], "prov-y keeps only its own model")
|
||||
})
|
||||
@@ -71,7 +71,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o"}, "prov-x", "prov-y"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-x"])
|
||||
assert.Equal(t, []string{"gpt-4o"}, got["prov-y"])
|
||||
})
|
||||
@@ -80,7 +80,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
policies := []*types.Policy{
|
||||
policyForProviders("p1", []string{"g-4o", "g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildProviderAllowlists(policies, byID)
|
||||
got := buildProviderAllowlists(policies, byID, nil)
|
||||
assert.ElementsMatch(t, []string{"claude-opus-4", "gpt-4o"}, got["prov-x"],
|
||||
"a policy's own multiple allowlist guardrails union together")
|
||||
})
|
||||
@@ -89,7 +89,7 @@ func TestBuildProviderAllowlists(t *testing.T) {
|
||||
empty := map[string]*types.Guardrail{"g-empty": allowlistGuardrail("g-empty", "acc-1")}
|
||||
got := buildProviderAllowlists([]*types.Policy{
|
||||
policyForProviders("p1", []string{"g-empty"}, "prov-x"),
|
||||
}, empty)
|
||||
}, empty, nil)
|
||||
assert.Equal(t, map[string][]string{"prov-x": {}}, got,
|
||||
"an enabled-but-empty allowlist is restricted with an empty set, not unrestricted")
|
||||
})
|
||||
@@ -124,7 +124,7 @@ func TestBuildModelPolicies(t *testing.T) {
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-sales"}, []string{"g-opus"}, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
got := buildModelPolicies(policies, byID, nil)
|
||||
assert.Equal(t, []routerModelPolicy{
|
||||
{GroupIDs: []string{"grp-eng"}, Models: []string{"gpt-4o"}},
|
||||
{GroupIDs: []string{"grp-sales"}, Models: []string{"claude-opus-4"}},
|
||||
@@ -137,14 +137,14 @@ func TestBuildModelPolicies(t *testing.T) {
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"),
|
||||
policyForGroups("p2", []string{"grp-admin"}, nil, "prov-x"),
|
||||
}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
got := buildModelPolicies(policies, byID, nil)
|
||||
assert.Nil(t, got["prov-x"][1].Models,
|
||||
"no allowlist must reach the router as nil, which lifts the restriction for its groups")
|
||||
})
|
||||
|
||||
t.Run("a disabled allowlist is not a restriction", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-disabled"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byID)
|
||||
got := buildModelPolicies(policies, byID, nil)
|
||||
assert.Nil(t, got["prov-x"][0].Models,
|
||||
"a guardrail with the allowlist check off restricts nothing")
|
||||
})
|
||||
@@ -154,7 +154,7 @@ func TestBuildModelPolicies(t *testing.T) {
|
||||
"g-empty": {ID: "g-empty", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: true}}},
|
||||
}
|
||||
policies := []*types.Policy{policyForGroups("p1", []string{"grp-eng"}, []string{"g-empty"}, "prov-x")}
|
||||
got := buildModelPolicies(policies, byIDEmpty)
|
||||
got := buildModelPolicies(policies, byIDEmpty, nil)
|
||||
require.NotNil(t, got["prov-x"][0].Models,
|
||||
"an empty allowlist must not arrive as nil — that would read as unrestricted")
|
||||
assert.Empty(t, got["prov-x"][0].Models)
|
||||
@@ -162,7 +162,71 @@ func TestBuildModelPolicies(t *testing.T) {
|
||||
|
||||
t.Run("a policy binding no groups is skipped", func(t *testing.T) {
|
||||
policies := []*types.Policy{policyForGroups("p1", nil, []string{"g-4o"}, "prov-x")}
|
||||
assert.Empty(t, buildModelPolicies(policies, byID),
|
||||
assert.Empty(t, buildModelPolicies(policies, byID, nil),
|
||||
"a policy with no source groups authorises nobody, so it bounds nobody's listing")
|
||||
})
|
||||
}
|
||||
|
||||
// TestSynthesizedAllowlists_ExpandDeclaredIDsPerProvider proves the
|
||||
// synthesized allowlists carry the canonical form alongside a raw declared
|
||||
// entry — under the destination provider's own catalog id, never another's —
|
||||
// so the proxy-side compares (guardrail backstop, per-group router rules)
|
||||
// admit the allowlist however the operator wrote it, while a plain provider's
|
||||
// "-vN"- or "@"-suffixed entries stay verbatim and cannot widen.
|
||||
func TestSynthesizedAllowlists_ExpandDeclaredIDsPerProvider(t *testing.T) {
|
||||
byID := map[string]*types.Guardrail{
|
||||
"g-raw": allowlistGuardrail("g-raw", "acc-1",
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"claude-sonnet-4-5@20250929",
|
||||
"gpt-4o"),
|
||||
}
|
||||
catalogByProvider := map[string]string{
|
||||
"prov-bedrock": "bedrock_api",
|
||||
"prov-vertex": "vertex_ai_api",
|
||||
"prov-plain": "openai_api",
|
||||
}
|
||||
policies := []*types.Policy{
|
||||
policyForGroups("p1", []string{"grp-eng"}, []string{"g-raw"},
|
||||
"prov-bedrock", "prov-vertex", "prov-plain"),
|
||||
}
|
||||
|
||||
t.Run("guardrail backstop expands under each provider's own normalizer", func(t *testing.T) {
|
||||
got := buildProviderAllowlists(policies, byID, catalogByProvider)
|
||||
assert.ElementsMatch(t, []string{
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"anthropic.claude-sonnet-4-5",
|
||||
"claude-sonnet-4-5@20250929",
|
||||
"gpt-4o",
|
||||
}, got["prov-bedrock"],
|
||||
"the Bedrock destination strips geography/version, but must not apply Vertex's @-strip")
|
||||
assert.ElementsMatch(t, []string{
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"claude-sonnet-4-5@20250929",
|
||||
"claude-sonnet-4-5",
|
||||
"gpt-4o",
|
||||
}, got["prov-vertex"],
|
||||
"the Vertex destination strips @version, but must not apply Bedrock's suffix strip")
|
||||
assert.ElementsMatch(t, []string{
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"claude-sonnet-4-5@20250929",
|
||||
"gpt-4o",
|
||||
}, got["prov-plain"],
|
||||
"a body-routed provider keeps every entry verbatim — no alternate can widen it")
|
||||
})
|
||||
|
||||
t.Run("router model rules expand the same way", func(t *testing.T) {
|
||||
got := buildModelPolicies(policies, byID, catalogByProvider)
|
||||
require.Len(t, got["prov-bedrock"], 1)
|
||||
assert.Contains(t, got["prov-bedrock"][0].Models, "anthropic.claude-sonnet-4-5")
|
||||
assert.NotContains(t, got["prov-bedrock"][0].Models, "claude-sonnet-4-5")
|
||||
require.Len(t, got["prov-vertex"], 1)
|
||||
assert.Contains(t, got["prov-vertex"][0].Models, "claude-sonnet-4-5")
|
||||
assert.NotContains(t, got["prov-vertex"][0].Models, "anthropic.claude-sonnet-4-5")
|
||||
require.Len(t, got["prov-plain"], 1)
|
||||
assert.ElementsMatch(t, []string{
|
||||
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"claude-sonnet-4-5@20250929",
|
||||
"gpt-4o",
|
||||
}, got["prov-plain"][0].Models)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -115,3 +115,20 @@ func TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix(t *testing.T) {
|
||||
assert.Equal(t, "/bedrock", out.Mutations.RewriteUpstream.StripPathPrefix,
|
||||
"the namespace prefix must not reach the real Bedrock endpoint")
|
||||
}
|
||||
|
||||
// TestRouteClaimsModel_VertexNormalizesCandidate is the Vertex counterpart of
|
||||
// the Bedrock case above: the parser strips the "@version" suffix from the
|
||||
// path model, so a provider registered with the versioned form must still
|
||||
// match the normalized request model.
|
||||
func TestRouteClaimsModel_VertexNormalizesCandidate(t *testing.T) {
|
||||
route := ProviderRoute{Vertex: true, Models: []string{"claude-sonnet-4-5@20250929"}}
|
||||
assert.True(t, routeClaimsModel(route, "claude-sonnet-4-5"),
|
||||
"raw @version Vertex model must match the normalized request model")
|
||||
assert.False(t, routeClaimsModel(route, "claude-opus-4-8"),
|
||||
"a model outside the provider's list must not match")
|
||||
|
||||
// Non-Vertex routes keep exact matching (no @version stripping).
|
||||
openai := ProviderRoute{Models: []string{"gpt-4o@2024"}}
|
||||
assert.False(t, routeClaimsModel(openai, "gpt-4o"),
|
||||
"non-Vertex routes must not strip an @version suffix")
|
||||
}
|
||||
|
||||
@@ -331,6 +331,11 @@ func discoverableModels(route ProviderRoute, userGroups []string) ([]string, boo
|
||||
intersection[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
if route.Vertex {
|
||||
if _, ok := permitted[llm.NormalizeVertexModel(m)]; ok {
|
||||
intersection[m] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
return sortedModels(intersection), true
|
||||
}
|
||||
@@ -869,6 +874,11 @@ func routeClaimsModel(route ProviderRoute, model string) bool {
|
||||
if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model {
|
||||
return true
|
||||
}
|
||||
// Vertex likewise: the parser strips the "@version" suffix from the
|
||||
// path model, while the operator may register the versioned form.
|
||||
if route.Vertex && llm.NormalizeVertexModel(candidate) == model {
|
||||
return true
|
||||
}
|
||||
// A client may pin a dated Anthropic id ("claude-sonnet-4-5-20250929")
|
||||
// where the operator registered the undated one. Only an undated
|
||||
// registration absorbs a dated request: normalising both sides would
|
||||
|
||||
Reference in New Issue
Block a user