From 7be45c2dd8b93434409037a05cc8bbdff8cc6a0a Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Sun, 23 Aug 2026 20:13:35 +0200 Subject: [PATCH] [proxy,management] Bound model discovery to the caller's own policies (#7239) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [proxy,management] Bound model discovery to the caller's own policies GET /v1/models was bounded by the provider record's enumerated models, which is the right bound only while a single policy reaches a provider. Where two teams share one provider under different allowlists, every caller was offered the union — each model outside their own policy being a request the guardrail refuses a moment later. A gateway record enumerating nothing was worse: it offered the upstream's entire catalogue however narrow the policy was. Each route now carries one rule per authorising policy — its source groups and the models it permits — instead of a single flattened list. At request time the router keeps the rules whose groups intersect the caller's, unions their models, and intersects that with what the provider serves. nil and [] stay distinct end to end: a policy setting no allowlist reaches the router as nil and lifts the restriction for the groups it binds, while an enabled allowlist with no models arrives as [] and permits nothing. Collapsing them would let a listing that should offer nothing fall open to everything. The guardrail's own per-provider allowlist is untouched. It is a fail-closed backstop that cannot tell who is asking, so discovery is now narrower than the backstop rather than wider. --- e2e/agentnetwork/custom_pricing_test.go | 7 +- e2e/agentnetwork/discovery_live_test.go | 400 ++++++++++++++++++ .../discovery_multipolicy_test.go | 170 ++++++++ e2e/harness/client.go | 36 +- .../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 ++++++ 9 files changed, 993 insertions(+), 25 deletions(-) create mode 100644 e2e/agentnetwork/discovery_live_test.go create mode 100644 e2e/agentnetwork/discovery_multipolicy_test.go diff --git a/e2e/agentnetwork/custom_pricing_test.go b/e2e/agentnetwork/custom_pricing_test.go index b3ca5028f..90e198d3d 100644 --- a/e2e/agentnetwork/custom_pricing_test.go +++ b/e2e/agentnetwork/custom_pricing_test.go @@ -174,7 +174,12 @@ func chatOnce(t *testing.T, ctx context.Context, env pricedEnv, model, sessionID // accessLogIngestWindow is how long a single request's access-log row is given // to appear before the caller gives up on it. -const accessLogIngestWindow = 30 * time.Second +// accessLogIngestWindow bounds how long a row may take to appear after its +// request returned. The proxy streams each entry to management with a 10s send +// timeout of its own, so a request whose send hits one full timeout and is +// retried has not yet missed anything real — 30s left barely three send +// attempts of headroom and lost the race on a loaded runner. +const accessLogIngestWindow = 60 * time.Second // accessLogPollInterval is how long the lookup waits between pages. Ingest is // asynchronous, so the row lands somewhere inside the window rather than on diff --git a/e2e/agentnetwork/discovery_live_test.go b/e2e/agentnetwork/discovery_live_test.go new file mode 100644 index 000000000..321e751bf --- /dev/null +++ b/e2e/agentnetwork/discovery_live_test.go @@ -0,0 +1,400 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "encoding/json" + "os" + "sort" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/e2e/harness" + sharedllm "github.com/netbirdio/netbird/shared/llm" + "github.com/netbirdio/netbird/shared/management/http/api" +) + +// TestLiveModelDiscovery drives model discovery against the REAL vendor +// endpoints — OpenAI, Anthropic, Bedrock and Vertex — rather than the mock. +// +// The mock upstream proves the filter's mechanics: it advertises ids we chose, +// so a listing narrowing to the ones we authorised is arithmetic we already +// controlled both sides of. What it cannot prove is that the filter survives +// contact with a real catalogue — ids we never enumerated, dated builds whose +// suffix the vendor picks, surfaces that answer a listing request with +// something other than a listing. That is what this covers, and it is the part +// a QA engineer would otherwise have to walk through by hand. +// +// One proxy serves every case. Each provider gets its own group, policy and +// client, because a model-less request matches exactly ONE route +// (matchModelless): with two providers authorised for the same caller, the +// listing would go to whichever won the tiebreak and the other would go +// untested. Group-scoping the caller makes each provider the only candidate +// for its own client. +func TestLiveModelDiscovery(t *testing.T) { + cases := liveDiscoveryCases() + if len(cases) == 0 { + t.Skip("no provider keys set; source ~/.llm-keys to run live model discovery") + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute) + defer cancel() + + t.Logf("[discovery] live matrix: %s", strings.Join(caseNames(cases), ", ")) + + // Provision every provider, group and policy before the proxy starts: the + // proxy takes a configuration snapshot at connect time and does not + // reconcile provider changes made afterwards. + keys := make(map[string]string, len(cases)) + for i := range cases { + keys[cases[i].name] = provisionLiveDiscovery(t, ctx, &cases[i]) + } + + endpoint, firstIP, firstClient, px := connectClient(t, ctx, "disc-live", keys[cases[0].name]) + clients := map[string]*harness.Client{cases[0].name: firstClient} + ips := map[string]string{cases[0].name: firstIP} + for _, tc := range cases[1:] { + cl := joinClient(t, ctx, px, endpoint, keys[tc.name]) + ip, err := cl.ResolveProxyIP(ctx, endpoint) + require.NoError(t, err, "resolve endpoint from the %s client", tc.name) + clients[tc.name] = cl + ips[tc.name] = ip + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + runLiveDiscoveryCase(t, ctx, tc, clients[tc.name], endpoint, ips[tc.name]) + }) + } +} + +// discoveryOutcome is what a discovery request must produce end to end. The +// three are genuinely different contracts, not degrees of success: only the +// first puts a bounded listing in front of the caller. +type discoveryOutcome int + +const ( + // outcomeFiltered: the proxy routes the request and bounds the response to + // what the caller may use. + outcomeFiltered discoveryOutcome = iota + // outcomeDenied: no provider of this shape can serve the surface, so the + // proxy refuses rather than rewriting the request onto an upstream that + // would 404 it. The caller gets a NetBird error, not a vendor one. + outcomeDenied + // outcomeUpstreamNoListing: the proxy routes the request to the configured + // upstream, and the vendor does not implement the endpoint there. Proxy + // side correct, product side a dead end — see the Bedrock case. + outcomeUpstreamNoListing +) + +// liveDiscoveryCase is one provider's discovery surface and what the proxy +// must make of it. +type liveDiscoveryCase struct { + name string + catalogID string + upstream string + apiKey string + + // path is the discovery endpoint the client calls. Not every surface uses + // /v1/models: Bedrock lists inference profiles instead. + path string + // headers the vendor requires on a bare GET (Anthropic versions its API + // through a header, and rejects a request without one). + headers []string + + // models the provider record enumerates. Empty models a gateway record, + // which enumerates nothing and claims everything. + models []string + // allowlist, when non-empty, is a guardrail narrowing the policy below the + // provider's own enumeration — the second of the two bounds discovery + // applies, and the only one a provider record alone cannot demonstrate. + allowlist []string + + // outcome is what this surface must produce end to end. + outcome discoveryOutcome + + // permitted is every id allowed to survive filtering, in the form the + // provider record registers it. A surviving id counts as permitted when it + // matches one of these outright or after Anthropic date-normalisation. + permitted []string + // wantHidden are ids the upstream is known to advertise and the bound must + // remove. Only set where we enumerate the model ourselves, so the + // expectation cannot rot when a vendor changes its catalogue. + wantHidden []string +} + +// liveDiscoveryCases builds the matrix from whichever provider credentials are +// present, mirroring availableProviders' env-var gating so a partial key set +// still yields partial coverage. +func liveDiscoveryCases() []liveDiscoveryCase { + var cases []liveDiscoveryCase + + // OpenAI enumerates TWO real models and the policy permits one. That is + // the only case here where both bounds are observable at once: the + // upstream advertises dozens of ids, the provider record cuts them to two, + // and the guardrail cuts those to one. + if k := os.Getenv("OPENAI_TOKEN"); k != "" { + cases = append(cases, liveDiscoveryCase{ + name: "openai", catalogID: "openai_api", upstream: "https://api.openai.com", apiKey: k, + path: "/v1/models", + models: []string{"gpt-4o-mini", "gpt-4o"}, + allowlist: []string{"gpt-4o-mini"}, + outcome: outcomeFiltered, + permitted: []string{"gpt-4o-mini"}, + wantHidden: []string{"gpt-4o"}, + }) + } + + // Anthropic is the surface Claude Code actually calls. Its listing returns + // DATED build ids (claude-haiku-4-5-20251001) while the provider record + // registers the undated id, so this is the case that proves the filter's + // date-normalisation against ids the vendor chose rather than ids we wrote. + if k := os.Getenv("ANTHROPIC_TOKEN"); k != "" { + cases = append(cases, liveDiscoveryCase{ + name: "anthropic", catalogID: "anthropic_api", upstream: "https://api.anthropic.com", apiKey: k, + path: "/v1/models", + headers: []string{"anthropic-version: 2023-06-01"}, + models: []string{"claude-haiku-4-5"}, + outcome: outcomeFiltered, + permitted: []string{"claude-haiku-4-5"}, + }) + } + + // Bedrock lists inference profiles, not models: matchModelless routes + // /inference-profiles to a Bedrock route and refuses /v1/models for one. + // + // The request reaches AWS and AWS refuses it — bedrock-runtime answers + // , because ListInferenceProfiles is a CONTROL + // PLANE operation served by bedrock..amazonaws.com, not the runtime + // host. A provider record carries one upstream and it has to be the runtime + // host for InvokeModel to work, so no Bedrock record can serve a listing as + // the model stands today. + // + // The mock upstream hides this entirely: it answers /inference-profiles on + // the same listener as everything else, so the routing test passes there + // while the real endpoint 404s. That is the whole reason this file exists, + // so the case is kept, asserting what actually happens. + if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" { + region := os.Getenv("AWS_REGION") + if region == "" { + region = "eu-central-1" + } + model := os.Getenv("AWS_BEDROCK_MODEL") + if model == "" { + model = "global.anthropic.claude-sonnet-4-6" + } + cases = append(cases, liveDiscoveryCase{ + name: "bedrock", catalogID: "bedrock_api", + upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, + path: "/inference-profiles", + models: []string{sharedllm.NormalizeAnthropicModel(strings.TrimPrefix(model, "global."))}, + outcome: outcomeUpstreamNoListing, + }) + } + + // Vertex carries the model in the rawPredict path and serves no listing + // endpoint at all, so the proxy must refuse discovery rather than rewrite + // it onto an upstream that would 404. + if sa := os.Getenv("GOOGLE_VERTEX_SA_BASE64"); sa != "" { + if project := os.Getenv("GOOGLE_VERTEX_PROJECT"); project != "" { + region := os.Getenv("GOOGLE_VERTEX_REGION") + if region == "" { + region = "global" + } + host := "aiplatform.googleapis.com" + if region != "global" { + host = region + "-aiplatform.googleapis.com" + } + cases = append(cases, liveDiscoveryCase{ + name: "vertex", catalogID: "vertex_ai_api", upstream: "https://" + host, + apiKey: "keyfile::" + sa, + path: "/v1/models", + outcome: outcomeDenied, + }) + } + } + + return cases +} + +// provisionLiveDiscovery creates the group, provider, optional guardrail and +// policy for one case, and returns the setup key a client joins that group +// with. Scoping each provider to its own group is what keeps it the only +// candidate for its own client's model-less request. +func provisionLiveDiscovery(t *testing.T, ctx context.Context, tc *liveDiscoveryCase) string { + t.Helper() + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-disc-live-" + tc.name}) + require.NoError(t, err, "create group for %s", tc.name) + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-disc-live-" + tc.name, + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key for %s", tc.name) + require.NotEmpty(t, sk.Key, "setup key plaintext for %s", tc.name) + + req := api.AgentNetworkProviderRequest{ + Name: "e2e-disc-live-" + tc.name, + ProviderId: tc.catalogID, + UpstreamUrl: tc.upstream, + ApiKey: &tc.apiKey, + Enabled: ptr(true), + } + if len(tc.models) > 0 { + models := make([]api.AgentNetworkProviderModel, 0, len(tc.models)) + for _, id := range tc.models { + models = append(models, api.AgentNetworkProviderModel{Id: id, InputPer1k: 0.001, OutputPer1k: 0.002}) + } + req.Models = &models + } + prov, err := srv.CreateProvider(ctx, req) + require.NoError(t, err, "create provider %s", tc.name) + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + polReq := api.AgentNetworkPolicyRequest{ + Name: "e2e-disc-live-" + tc.name, + Enabled: ptr(true), + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + } + if len(tc.allowlist) > 0 { + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-disc-live-" + tc.name + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = tc.allowlist + g, gerr := srv.CreateGuardrail(ctx, gr) + require.NoError(t, gerr, "create guardrail for %s", tc.name) + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), g.Id) }) + polReq.GuardrailIds = &[]string{g.Id} + } + pol, err := srv.CreatePolicy(ctx, polReq) + require.NoError(t, err, "create policy for %s", tc.name) + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + return sk.Key +} + +// runLiveDiscoveryCase issues the discovery request and reports everything the +// vendor said before asserting on any of it. The log is the point on the first +// run: a live catalogue is the one input we do not control, so a failure has to +// arrive with the response that caused it rather than just a count. +func runLiveDiscoveryCase(t *testing.T, ctx context.Context, tc liveDiscoveryCase, cl *harness.Client, endpoint, proxyIP string) { + t.Helper() + + // A single request is enough for the two non-listing outcomes, and retrying + // them would burn the retry window waiting for a status that is never + // coming. + if tc.outcome != outcomeFiltered { + code, body, err := cl.Get(ctx, endpoint, proxyIP, tc.path, tc.headers) + require.NoError(t, err, "request must reach the proxy") + t.Logf("[discovery] %s GET %s -> %d; body: %s", tc.name, tc.path, code, truncate(body, 2000)) + assert.NotEqual(t, 200, code, + "%s serves no bounded listing, so a 200 here would mean the caller was handed a picker nothing narrows; body: %s", + tc.name, truncate(body, 2000)) + + // Which side refused is the whole distinction between these two + // outcomes, and a NetBird error is the thing that tells them apart: the + // middleware chain stamps its own name on anything it generates. + if tc.outcome == outcomeDenied { + assert.True(t, isProxyError(body), + "%s serves no listing endpoint at all, so the proxy must refuse the request itself rather than forward it to an upstream that would answer for us; body: %s", + tc.name, truncate(body, 2000)) + return + } + assert.False(t, isProxyError(body), + "%s discovery must be routed to the configured upstream and refused by the vendor, not blocked by the proxy; body: %s", + tc.name, truncate(body, 2000)) + return + } + + code, body := callUntil(t, func() (int, string, error) { + return cl.Get(ctx, endpoint, proxyIP, tc.path, tc.headers) + }, 200) + t.Logf("[discovery] %s GET %s -> %d; body: %s", tc.name, tc.path, code, truncate(body, 4000)) + require.Equal(t, 200, code, "%s discovery must be served; body: %s", tc.name, truncate(body, 2000)) + + ids, ok := listingIDs(body) + require.Truef(t, ok, + "%s answered discovery with something other than a {\"data\":[{\"id\":…}]} listing, which the filter forwards untouched — the caller would get an unbounded picker; body: %s", + tc.name, truncate(body, 2000)) + sort.Strings(ids) + t.Logf("[discovery] %s: %d ids after filtering: %s", tc.name, len(ids), strings.Join(ids, ", ")) + + require.NotEmpty(t, ids, "%s filtered the listing down to nothing; the caller would see an empty picker", tc.name) + + permitted := make(map[string]struct{}, len(tc.permitted)*2) + for _, id := range tc.permitted { + permitted[id] = struct{}{} + permitted[sharedllm.NormalizeAnthropicModel(id)] = struct{}{} + } + for _, id := range ids { + _, direct := permitted[id] + _, normalised := permitted[sharedllm.NormalizeAnthropicModel(id)] + assert.Truef(t, direct || normalised, + "%s offered %q, which no policy on this route permits — every entry the picker shows must be a request the guardrail would allow", tc.name, id) + } + for _, hidden := range tc.wantHidden { + assert.NotContainsf(t, ids, hidden, + "%s offered %q, which the provider enumerates but the policy does not permit", tc.name, hidden) + } +} + +// isProxyError reports whether a response body was generated by the middleware +// chain rather than forwarded from a vendor. Every chain-generated error names +// the middleware that raised it, which no upstream's error body does — so this +// separates "the proxy refused" from "the proxy routed it and the vendor +// refused", the two failures that otherwise look alike from the client side. +func isProxyError(body string) bool { + return strings.Contains(body, `"middleware":`) +} + +// listingIDs pulls the model ids out of a listing response. ok is false when +// the body is not the {"data":[{"id":…}]} shape the filter recognises. +func listingIDs(body string) ([]string, bool) { + var doc struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + if err := json.Unmarshal([]byte(body), &doc); err != nil { + return nil, false + } + if doc.Data == nil { + return nil, false + } + ids := make([]string, 0, len(doc.Data)) + for _, entry := range doc.Data { + ids = append(ids, entry.ID) + } + return ids, true +} + +func caseNames(cases []liveDiscoveryCase) []string { + names := make([]string, 0, len(cases)) + for _, c := range cases { + names = append(names, c.name) + } + return names +} + +// truncate bounds a logged response body. A live catalogue can run to tens of +// kilobytes, and the useful part is the front. +func truncate(s string, limit int) string { + if len(s) <= limit { + return s + } + return s[:limit] + "… (" + strconv.Itoa(len(s)-limit) + " more bytes)" +} diff --git a/e2e/agentnetwork/discovery_multipolicy_test.go b/e2e/agentnetwork/discovery_multipolicy_test.go new file mode 100644 index 000000000..447c1314c --- /dev/null +++ b/e2e/agentnetwork/discovery_multipolicy_test.go @@ -0,0 +1,170 @@ +//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 + mkKey := func(name, groupID string) string { + sk, kerr := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: name, + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{groupID}, + Ephemeral: &ephemeral, + }) + require.NoError(t, kerr, "mint setup key %s", name) + require.NotEmpty(t, sk.Key, "setup key plaintext") + return sk.Key + } + // One client per group. The second is what makes the first assertion mean + // something: without a client that DOES see the other team's model, its + // absence from the main client's listing could equally be a policy that + // never propagated. + keyMain := mkKey("e2e-disc-mp-main-client", grpMain.Id) + keyOther := mkKey("e2e-disc-mp-other-client", grpOther.Id) + + // 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, clMain, px := connectClient(t, ctx, "disc-mp", keyMain) + clOther := joinClient(t, ctx, px, endpoint, keyOther) + + listing := func(t *testing.T, cl *harness.Client, ip string) string { + t.Helper() + code, body := callUntil(t, func() (int, string, error) { + return cl.Get(ctx, endpoint, ip, "/v1/models?limit=1000", nil) + }, 200) + require.Equal(t, 200, code, "discovery must be served; body: %s", body) + return body + } + + otherIP, err := clOther.ResolveProxyIP(ctx, endpoint) + require.NoError(t, err, "resolve endpoint from the other client") + + // The other team's client first: seeing its own model proves polOther is + // live, so the main client's listing is narrowed by policy scoping rather + // than by the other policy having failed to apply at all. + otherBody := listing(t, clOther, otherIP) + assert.Contains(t, otherBody, harness.VLLMUnlistedModel, + "the other group's policy must be in force, or this test proves nothing") + assert.NotContains(t, otherBody, harness.VLLMModel, + "and it must not be offered the main group's model either — isolation runs both ways") + + mainBody := listing(t, clMain, proxyIP) + assert.Contains(t, mainBody, harness.VLLMModel, + "the model the caller's own policy permits must reach the picker") + assert.NotContains(t, mainBody, harness.VLLMUnlistedModel, + "a model only another group's policy permits must not be offered to this caller") +} + +// joinClient starts a second tunnel client against an already-running proxy, so +// a test can drive the same endpoint as two different group memberships without +// paying for a second proxy. +func joinClient(t *testing.T, ctx context.Context, px *harness.Proxy, endpoint, setupKey string) *harness.Client { + t.Helper() + + cl, err := harness.StartClient(ctx, srv, setupKey) + require.NoError(t, err, "start second client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "second client must connect to management") + _, err = cl.ResolveProxyIP(ctx, endpoint) + require.NoError(t, err, "second client could not resolve the endpoint") + // Guarded rather than passed straight to require: px.Logs pulls the whole + // proxy container log, which is only worth fetching when the wait failed. + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + require.NoError(t, err, "second client did not see the proxy peer\n=== proxy logs ===\n%s", + px.Logs(context.Background())) + } + return cl +} diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 9e9e7b34a..73931027d 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -200,12 +200,18 @@ func (cl *Client) pollStatus(ctx context.Context, timeout time.Duration, want st const ( // curlExitCouldNotResolve is curl's exit code for a DNS resolution failure, distinct from connection-level failures. curlExitCouldNotResolve = 6 - // dnsProbeRetryWindow bounds DNS-failure retries: the synthesized zone lands a beat after management connects, so early NXDOMAIN is propagation; a zone still absent after this window is a real failure. - dnsProbeRetryWindow = 30 * time.Second - dnsProbeRetryInterval = 2 * time.Second + // curlExitCouldNotConnect is curl's exit code for a connection that never + // established. The probe exists to WAKE the lazy proxy peer, so the first + // attempt legitimately arrives before WireGuard has brought the tunnel up + // and fails here — which is propagation, exactly like an early NXDOMAIN, + // and belongs inside the retry window rather than failing the test outright. + curlExitCouldNotConnect = 7 + // endpointProbeRetryWindow bounds retries of the transient failures above: the synthesized zone and the tunnel both land a beat after management connects. Still failing after this window is a real failure. + endpointProbeRetryWindow = 30 * time.Second + endpointProbeRetryInterval = 2 * time.Second ) -// ResolveProxyIP GETs https:/// from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; only DNS failures retry, within dnsProbeRetryWindow. Returns the connected IP for --resolve pinning. +// ResolveProxyIP GETs https:/// from the client's netns: any HTTP status proves DNS + tunnel and wakes the lazy proxy peer; DNS and connect failures retry, within endpointProbeRetryWindow. Returns the connected IP for --resolve pinning. func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, error) { args := []string{ "run", "--rm", @@ -216,7 +222,7 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, "-w", "%{remote_ip}", "https://" + endpoint + "/", } - deadline := time.Now().Add(dnsProbeRetryWindow) + deadline := time.Now().Add(endpointProbeRetryWindow) for { cmd := exec.CommandContext(ctx, "docker", args...) var stdout, stderr strings.Builder @@ -232,21 +238,29 @@ func (cl *Client) ResolveProxyIP(ctx context.Context, endpoint string) (string, } var exitErr *exec.ExitError - if !errors.As(err, &exitErr) || exitErr.ExitCode() != curlExitCouldNotResolve { + if !errors.As(err, &exitErr) || !isTransientProbeExit(exitErr.ExitCode()) { return "", fmt.Errorf("no HTTP response from %s: %w (%s)", endpoint, err, strings.TrimSpace(stderr.String())) } - dnsErr := fmt.Errorf("DNS resolution failed for %s: %s", endpoint, strings.TrimSpace(stderr.String())) - if time.Until(deadline) < dnsProbeRetryInterval { - return "", dnsErr + probeErr := fmt.Errorf("endpoint %s not reachable yet: %s", endpoint, strings.TrimSpace(stderr.String())) + if time.Until(deadline) < endpointProbeRetryInterval { + return "", probeErr } select { case <-ctx.Done(): - return "", fmt.Errorf("%w (%w)", dnsErr, ctx.Err()) - case <-time.After(dnsProbeRetryInterval): + return "", fmt.Errorf("%w (%w)", probeErr, ctx.Err()) + case <-time.After(endpointProbeRetryInterval): } } } +// isTransientProbeExit reports whether a curl exit code describes a state the +// endpoint is expected to pass THROUGH on its way up, rather than a settled +// failure. Anything else — TLS refusal, a protocol error, a bad argument — +// would still be failing after the retry window, so it fails immediately. +func isTransientProbeExit(code int) bool { + return code == curlExitCouldNotResolve || code == curlExitCouldNotConnect +} + // Wire shapes for Chat. const ( // WireChat is the OpenAI-compatible /v1/chat/completions shape. 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") + }) +}