From f8a55ac8d41dcc1687ca790fe895bd10c077835d Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Tue, 18 Aug 2026 12:06:41 +0000 Subject: [PATCH] [proxy,management] Bound model discovery to the caller's own policies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The listing was narrowed by the provider record's enumerated models, which is the right bound only while one policy reaches a provider. Where two teams share a provider under different allowlists, every caller was offered the union: each model outside their own policy is a request the guardrail refuses a moment later, which is the empty-or-wrong picker this endpoint exists to avoid, moved one level up. A gateway record enumerating nothing was worse still — it offered the upstream's entire catalogue however narrow the policy. The synthesiser already knows which policies authorise a provider and which groups each binds, so the router can answer this at request time where it knows the caller's groups. Each route now carries one rule per authorising policy — its source groups and the models it permits — and the listing is bounded to the union across the rules matching the caller, intersected with what the provider serves. This is deliberately finer than the guardrail's own per-provider allowlist, which stays as it is: that list is a fail-closed backstop and cannot tell who is asking, so discovery is now narrower than the backstop rather than wider. A policy setting no allowlist lifts the restriction for the groups it binds, so nil and empty model lists stay distinct end to end — collapsing them would let a listing that should offer nothing fall open to everything. --- .../discovery_multipolicy_test.go | 122 +++++++++++++++ .../modules/agentnetwork/synthesizer.go | 70 ++++++++- .../synthesizer_provider_allowlist_test.go | 73 +++++++++ .../middleware/builtin/llm_router/factory.go | 18 +++ .../builtin/llm_router/middleware.go | 103 ++++++++++++- .../builtin/llm_router/middleware_test.go | 141 ++++++++++++++++++ 6 files changed, 514 insertions(+), 13 deletions(-) create mode 100644 e2e/agentnetwork/discovery_multipolicy_test.go diff --git a/e2e/agentnetwork/discovery_multipolicy_test.go b/e2e/agentnetwork/discovery_multipolicy_test.go new file mode 100644 index 000000000..f2a587d8a --- /dev/null +++ b/e2e/agentnetwork/discovery_multipolicy_test.go @@ -0,0 +1,122 @@ +//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" +) + +// TestDiscoveryBoundToCallersPolicies covers a model listing on a provider two +// teams reach under different allowlists. +// +// Bounding the listing by the provider's enumerated models alone is not enough +// once more than one policy is in play: the caller would be offered every model +// any team may use, and each one outside their own policy is a request the +// guardrail refuses a moment later — the empty-or-wrong picker this endpoint +// exists to avoid, just moved one level up. +// +// The client joins the main group only. Both models are enumerated by the same +// provider and both are advertised by the upstream, so a listing that leaked +// the other team's model would visibly contain it. +func TestDiscoveryBoundToCallersPolicies(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grpMain, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-main"}) + require.NoError(t, err, "create main group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpMain.Id) }) + + grpOther, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-mp-other"}) + require.NoError(t, err, "create other group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grpOther.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-disc-mp-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grpMain.Id}, // the client joins grpMain only + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + // One provider enumerating both models the upstream advertises, so the + // listing is narrowed by policy rather than by what the provider serves. + staticKey := "static-e2e-token" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-disc-mp", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &staticKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{ + {Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.001}, + {Id: harness.VLLMUnlistedModel, InputPer1k: 0.001, OutputPer1k: 0.001}, + }, + }) + require.NoError(t, err, "create provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + mkGuardrail := func(name, model string) api.AgentNetworkGuardrail { + var gr api.AgentNetworkGuardrailRequest + gr.Name = name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{model} + g, gerr := srv.CreateGuardrail(ctx, gr) + require.NoError(t, gerr, "create guardrail %s", name) + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) }) + return g + } + gMain := mkGuardrail("e2e-disc-mp-main", harness.VLLMModel) + gOther := mkGuardrail("e2e-disc-mp-other", harness.VLLMUnlistedModel) + + enabled := true + polMain, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-disc-mp-main", + Enabled: &enabled, + SourceGroups: []string{grpMain.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gMain.Id}, + }) + require.NoError(t, err, "create main policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polMain.Id) }) + + // The other team's policy, on the same provider, permitting the model the + // client must never be offered. + polOther, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-disc-mp-other", + Enabled: &enabled, + SourceGroups: []string{grpOther.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{gOther.Id}, + }) + require.NoError(t, err, "create other policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), polOther.Id) }) + + endpoint, proxyIP, cl, px := connectClient(t, ctx, "disc-mp", sk.Key) + _ = px + + code, body := callUntil(t, func() (int, string, error) { + return cl.Get(ctx, endpoint, proxyIP, "/v1/models?limit=1000", nil) + }, 200) + require.Equal(t, 200, code, "discovery must be served; body: %s", body) + + assert.Contains(t, body, harness.VLLMModel, + "the model the caller's own policy permits must reach the picker") + assert.NotContains(t, body, harness.VLLMUnlistedModel, + "a model only another group's policy permits must not be offered to this caller") +} diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index 3fd92be96..76944698e 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -211,7 +211,19 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([ groupIndex := indexProviderGroups(enabledPolicies) - routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex) + // 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) + + // 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) + + routerCfgJSON, err := buildRouterConfigJSON(enabledProviders, groupIndex, modelPolicies) if err != nil { return nil, err } @@ -228,11 +240,6 @@ func SynthesizeServices(ctx context.Context, s store.Store, accountID string) ([ mergedGuardrails := mergeGuardrails(enabledPolicies, guardrailsByID) applyAccountCollectionControls(&mergedGuardrails, settings) - // The proxy guardrail is a per-provider fail-closed backstop; the - // authoritative per-policy/group decision is management's - // SelectPolicyForRequest. A provider lands in this map only when every - // authorising policy restricts models. - providerAllowlists := buildProviderAllowlists(enabledPolicies, guardrailsByID) guardrailJSON, err := marshalGuardrailConfig(providerAllowlists, mergedGuardrails.PromptCapture) if err != nil { return nil, err @@ -351,6 +358,11 @@ type routerProviderRoute struct { AuthHeaderName string `json:"auth_header_name"` AuthHeaderValue string `json:"auth_header_value"` AllowedGroupIDs []string `json:"allowed_group_ids,omitempty"` + // ModelPolicies is one entry per enabled policy authorising this provider, + // carrying that policy's source groups and the models it permits. The + // router bounds a model listing with it, so a provider two groups reach + // under different allowlists offers each only its own. + ModelPolicies []routerModelPolicy `json:"model_policies,omitempty"` // Vertex marks a Google Vertex AI provider, whose requests carry the // model in the URL path. The router selects it by path, bypassing the // model/vendor table. @@ -422,7 +434,7 @@ func indexProviderGroups(policies []*types.Policy) map[string][]string { // path-prefix tiebreak. Providers no enabled policy authorises // (orphans) are intentionally OMITTED so the router never observes a // route with an empty ACL. -func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string) ([]byte, error) { +func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][]string, modelPolicies map[string][]routerModelPolicy) ([]byte, error) { cfg := routerConfig{Providers: make([]routerProviderRoute, 0, len(providers))} for _, p := range providers { groups, hasPolicy := groupIndex[p.ID] @@ -449,6 +461,7 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][] AuthHeaderName: headerName, AuthHeaderValue: headerValue, AllowedGroupIDs: groups, + ModelPolicies: modelPolicies[p.ID], Vertex: catalog.IsVertexPathStyle(p.ProviderID), Bedrock: catalog.IsBedrockPathStyle(p.ProviderID), GCPServiceAccountKeyB64: gcpSAKeyB64, @@ -1098,3 +1111,46 @@ func mergeGuardrail(g *types.Guardrail, merged *MergedGuardrails) { } } } + +// routerModelPolicy mirrors the router's ModelPolicyRule: one authorising +// policy's source groups plus the models it permits. Models is nil for a +// policy that sets no model allowlist, which lifts the restriction for the +// groups it binds — so nil and empty must survive the round trip distinctly. +type routerModelPolicy struct { + GroupIDs []string `json:"group_ids"` + Models []string `json:"models"` +} + +// buildModelPolicies indexes, per provider, one rule for each enabled policy +// authorising it: the policy's source groups and the models its guardrail +// permits. +// +// This is deliberately finer than buildProviderAllowlists, which flattens the +// same inputs into one list per provider for the proxy's fail-closed guardrail. +// A flattened list cannot answer "what may THIS caller see", so a provider two +// teams reach under different allowlists would offer each team the other's +// 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 { + 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 + } + out[providerID] = append(out[providerID], rule) + } + } + return out +} diff --git a/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go b/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go index 2cfc0db8c..a27cd2ae4 100644 --- a/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_provider_allowlist_test.go @@ -4,6 +4,7 @@ import ( "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" ) @@ -93,3 +94,75 @@ func TestBuildProviderAllowlists(t *testing.T) { "an enabled-but-empty allowlist is restricted with an empty set, not unrestricted") }) } + +// policyForGroups builds an enabled policy binding the given source groups to +// the given providers under an optional guardrail. +func policyForGroups(id string, groups []string, guardrailIDs []string, providerIDs ...string) *types.Policy { + return &types.Policy{ + ID: id, + Enabled: true, + SourceGroups: groups, + DestinationProviderIDs: providerIDs, + GuardrailIDs: guardrailIDs, + } +} + +// TestBuildModelPolicies covers the finer index discovery needs. Where +// buildProviderAllowlists flattens every authorising policy into one list per +// provider — enough for a fail-closed backstop, but blind to who is asking — +// this keeps each policy's source groups beside its models so the router can +// bound a listing to the calling groups. +func TestBuildModelPolicies(t *testing.T) { + byID := map[string]*types.Guardrail{ + "g-4o": allowlistGuardrail("g-4o", "acc-1", "gpt-4o"), + "g-opus": allowlistGuardrail("g-opus", "acc-1", "claude-opus-4"), + "g-disabled": {ID: "g-disabled", Checks: types.GuardrailChecks{ModelAllowlist: types.GuardrailModelAllowlist{Enabled: false, Models: []string{"gpt-4o"}}}}, + } + + t.Run("each policy keeps its own groups and models", func(t *testing.T) { + policies := []*types.Policy{ + policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"), + policyForGroups("p2", []string{"grp-sales"}, []string{"g-opus"}, "prov-x"), + } + got := buildModelPolicies(policies, byID) + assert.Equal(t, []routerModelPolicy{ + {GroupIDs: []string{"grp-eng"}, Models: []string{"gpt-4o"}}, + {GroupIDs: []string{"grp-sales"}, Models: []string{"claude-opus-4"}}, + }, got["prov-x"], + "the two policies must stay separable so neither group is offered the other's models") + }) + + t.Run("an unrestricted policy carries nil models", func(t *testing.T) { + policies := []*types.Policy{ + policyForGroups("p1", []string{"grp-eng"}, []string{"g-4o"}, "prov-x"), + policyForGroups("p2", []string{"grp-admin"}, nil, "prov-x"), + } + got := buildModelPolicies(policies, byID) + 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) + assert.Nil(t, got["prov-x"][0].Models, + "a guardrail with the allowlist check off restricts nothing") + }) + + t.Run("an enabled allowlist with no models permits nothing", func(t *testing.T) { + byIDEmpty := map[string]*types.Guardrail{ + "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) + 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) + }) + + 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), + "a policy with no source groups authorises nobody, so it bounds nobody's listing") + }) +} diff --git a/proxy/internal/middleware/builtin/llm_router/factory.go b/proxy/internal/middleware/builtin/llm_router/factory.go index 938a23ebe..ae3d44a40 100644 --- a/proxy/internal/middleware/builtin/llm_router/factory.go +++ b/proxy/internal/middleware/builtin/llm_router/factory.go @@ -44,6 +44,12 @@ type ProviderRoute struct { AuthHeaderName string `json:"auth_header_name"` AuthHeaderValue string `json:"auth_header_value"` AllowedGroupIDs []string `json:"allowed_group_ids"` + // ModelPolicies carries, per authorising policy, the source groups it + // binds and the models it permits. The router uses it to bound a model + // listing to what THIS caller may use: a provider reachable by two groups + // under different allowlists must not offer either group the other's + // models. Empty means no policy restricts models on this route. + ModelPolicies []ModelPolicyRule `json:"model_policies,omitempty"` // Vertex marks a Google Vertex AI provider. Vertex requests carry the // model in the URL path, so the router selects this route by path // (isVertexPath) and bypasses the model/vendor table entirely. @@ -65,6 +71,18 @@ type ProviderRoute struct { SkipTLSVerify bool `json:"skip_tls_verify,omitempty"` } +// ModelPolicyRule is one authorising policy's contribution to what a caller +// may use on a route: the source groups it binds, and the models it permits. +// +// Models is nil when the policy sets no model allowlist — an unrestricted +// policy, which lifts the restriction for the groups it binds. That is why +// nil and empty must stay distinct: an empty list is a guardrail that permits +// nothing, and collapsing the two would let a listing fail open. +type ModelPolicyRule struct { + GroupIDs []string `json:"group_ids"` + Models []string `json:"models"` +} + // Config is the on-wire configuration accepted by the factory. An // empty Providers slice yields a router that denies every request as // not-routable; the synthesiser is responsible for stamping the diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index e6ad332fc..01981666c 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -242,12 +242,13 @@ func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix { stripBedrockNamespace(out) } - // A route that enumerates its models bounds what the caller may use, - // so the picker must not offer the rest: every entry outside the list - // is a request the chain will deny. - if reqPath == modelListingPath && len(route.Models) > 0 && - out.Mutations != nil && out.Mutations.RewriteUpstream != nil { - out.Mutations.RewriteUpstream.DiscoveryModels = append([]string(nil), route.Models...) + // What the caller may actually use bounds what the picker may offer: + // every entry outside it is a request the chain will deny a moment + // later. + if reqPath == modelListingPath && out.Mutations != nil && out.Mutations.RewriteUpstream != nil { + if models, bounded := discoverableModels(route, userGroups); bounded { + out.Mutations.RewriteUpstream.DiscoveryModels = models + } } return out case matchOutcomeUnauthorised: @@ -271,6 +272,96 @@ func isNonInferenceMethod(method string) bool { return method == http.MethodGet || method == http.MethodHead } +// discoverableModels returns the model ids a caller in userGroups may actually +// use on this route, and whether the listing should be bounded to them at all. +// +// Two things narrow a listing, and both must apply or the picker offers models +// the very next request refuses: +// +// - the provider's own enumerated models, when it lists any (a gateway record +// enumerates nothing and claims everything); +// - the model allowlists of the policies that authorise THIS caller. A +// provider reachable by two groups under different allowlists must not +// offer either group the other's models, which is why the rules carry their +// source groups rather than arriving pre-flattened. +// +// A policy that sets no allowlist lifts the restriction for the groups it +// binds, so a caller holding one unrestricted policy sees the provider's full +// list. bounded is false when nothing narrows the listing — an unrestricted +// caller on a route that enumerates nothing — in which case the upstream's own +// answer passes through untouched. +func discoverableModels(route ProviderRoute, userGroups []string) ([]string, bool) { + permitted, restricted := policyPermittedModels(route, userGroups) + + switch { + case !restricted && len(route.Models) == 0: + return nil, false + case !restricted: + return append([]string(nil), route.Models...), true + case len(route.Models) == 0: + // A gateway record enumerates nothing, so the allowlist is the whole + // bound — previously such a record offered the upstream's entire + // catalogue however narrow the policy was. + return sortedModels(permitted), true + } + + // Both bound: only what the provider serves and the policy permits. + intersection := make(map[string]struct{}, len(route.Models)) + for _, m := range route.Models { + if _, ok := permitted[m]; ok { + intersection[m] = struct{}{} + } + } + return sortedModels(intersection), true +} + +// policyPermittedModels folds the rules whose groups intersect the caller's +// into the set of models they permit. restricted is false when the caller +// holds at least one authorising policy that sets no allowlist, or when no +// rule binds them at all. +func policyPermittedModels(route ProviderRoute, userGroups []string) (map[string]struct{}, bool) { + permitted := make(map[string]struct{}) + restricted := false + for _, rule := range route.ModelPolicies { + if !groupsIntersect(rule.GroupIDs, userGroups) { + continue + } + if rule.Models == nil { + // An unrestricted policy the caller holds lifts the restriction + // entirely, whatever the others say. + return nil, false + } + restricted = true + for _, m := range rule.Models { + permitted[m] = struct{}{} + } + } + return permitted, restricted +} + +// groupsIntersect reports whether the two group-id sets share a member. +func groupsIntersect(a, b []string) bool { + for _, x := range a { + for _, y := range b { + if x == y { + return true + } + } + } + return false +} + +// sortedModels flattens a model set into a stable slice so the bound the proxy +// applies — and any test asserting on it — does not depend on map order. +func sortedModels(set map[string]struct{}) []string { + out := make([]string, 0, len(set)) + for m := range set { + out = append(out, m) + } + sort.Strings(out) + return out +} + // markNonInference tags an allow as a request that spends no tokens, so the // limit check skips the management pre-flight it would charge nothing against. func markNonInference(out *middleware.Output) { diff --git a/proxy/internal/middleware/builtin/llm_router/middleware_test.go b/proxy/internal/middleware/builtin/llm_router/middleware_test.go index 336cdb9fe..5a1d32480 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware_test.go @@ -1145,3 +1145,144 @@ func TestRouter_PinnedDatedModelStaysDistinct(t *testing.T) { "declaration order must not decide between two deliberately pinned builds") }) } + +// TestRouter_DiscoveryBoundToCallersPolicies pins that a model listing is +// bounded by the policies that authorise the caller, not by the union across +// everyone who can reach the provider. Two teams sharing one provider record +// under different allowlists is the case that makes the difference visible: a +// flattened per-provider list would offer each team the other's models, and +// every one of those entries is a request the guardrail then refuses. +func TestRouter_DiscoveryBoundToCallersPolicies(t *testing.T) { + const ( + eng = "grp-eng" + sales = "grp-sales" + ) + route := ProviderRoute{ + ID: "shared-gateway", + Models: []string{"claude-sonnet-5", "claude-haiku-4-5", "gpt-4o"}, + AllowedGroupIDs: []string{eng, sales}, + UpstreamScheme: "https", + UpstreamHost: "gateway.example.com", + ModelPolicies: []ModelPolicyRule{ + {GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}}, + {GroupIDs: []string{sales}, Models: []string{"gpt-4o"}}, + }, + } + + listingFor := func(t *testing.T, group string) []string { + t.Helper() + mw := New(Config{Providers: []ProviderRoute{route}}) + in := newModellessInput(modelListingPath) + in.UserGroups = []string{group} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + return out.Mutations.RewriteUpstream.DiscoveryModels + } + + t.Run("each group sees only its own policy's models", func(t *testing.T) { + assert.Equal(t, []string{"claude-sonnet-5"}, listingFor(t, eng), + "engineering must not be offered the model only sales may use") + assert.Equal(t, []string{"gpt-4o"}, listingFor(t, sales), + "sales must not be offered the model only engineering may use") + }) + + t.Run("a model no policy allows is offered to nobody", func(t *testing.T) { + for _, group := range []string{eng, sales} { + assert.NotContains(t, listingFor(t, group), "claude-haiku-4-5", + "the provider serves it, but no policy permits it") + } + }) +} + +// TestRouter_DiscoveryUnrestrictedPolicy covers the lifting rule: a caller +// holding one policy without a model allowlist sees everything the provider +// enumerates, whatever the other policies say. +func TestRouter_DiscoveryUnrestrictedPolicy(t *testing.T) { + const ( + eng = "grp-eng" + admin = "grp-admin" + ) + route := ProviderRoute{ + ID: "shared-gateway", + Models: []string{"claude-sonnet-5", "gpt-4o"}, + AllowedGroupIDs: []string{eng, admin}, + UpstreamScheme: "https", + UpstreamHost: "gateway.example.com", + ModelPolicies: []ModelPolicyRule{ + {GroupIDs: []string{eng}, Models: []string{"claude-sonnet-5"}}, + // nil Models: a policy that sets no allowlist at all. + {GroupIDs: []string{admin}}, + }, + } + mw := New(Config{Providers: []ProviderRoute{route}}) + + in := newModellessInput(modelListingPath) + in.UserGroups = []string{eng, admin} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.ElementsMatch(t, []string{"claude-sonnet-5", "gpt-4o"}, + out.Mutations.RewriteUpstream.DiscoveryModels, + "an unrestricted policy the caller holds lifts the restriction") +} + +// TestRouter_DiscoveryOnGatewayRecord covers a record that enumerates no +// models. It previously offered the upstream's whole catalogue however narrow +// the policy was, because there was nothing to intersect against; the policy +// allowlist is now the bound on its own. +func TestRouter_DiscoveryOnGatewayRecord(t *testing.T) { + const eng = "grp-eng" + base := ProviderRoute{ + ID: "litellm", + AllowedGroupIDs: []string{eng}, + UpstreamScheme: "https", + UpstreamHost: "litellm.internal", + } + + t.Run("a policy allowlist bounds it", func(t *testing.T) { + route := base + route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{"gpt-4o"}}} + mw := New(Config{Providers: []ProviderRoute{route}}) + + in := newModellessInput(modelListingPath) + in.UserGroups = []string{eng} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, []string{"gpt-4o"}, out.Mutations.RewriteUpstream.DiscoveryModels, + "a catch-all record must still be bounded by what policy permits") + }) + + t.Run("an allowlist permitting nothing offers nothing", func(t *testing.T) { + route := base + route.ModelPolicies = []ModelPolicyRule{{GroupIDs: []string{eng}, Models: []string{}}} + mw := New(Config{Providers: []ProviderRoute{route}}) + + in := newModellessInput(modelListingPath) + in.UserGroups = []string{eng} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels, + "an empty allowlist permits nothing, and must not be read as unrestricted") + }) + + t.Run("no policy restriction leaves the listing alone", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{base}}) + + in := newModellessInput(modelListingPath) + in.UserGroups = []string{eng} + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Nil(t, out.Mutations.RewriteUpstream.DiscoveryModels, + "nothing narrows the listing, so the upstream's own answer passes through") + }) +}