From acfad1a384d20dbb47d91c3195055899efa70220 Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sun, 23 Aug 2026 06:35:27 +0000 Subject: [PATCH] [proxy] Strip only a gateway's own namespace before matching a model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three open review findings. The discovery filter treated every slash in a listed id as a gateway prefix and matched the tail against the policy. A self-hosted id carries slashes of its own, and an upstream may scope ids per tenant, so "tenant-b/claude-sonnet-5" matched a permitted "claude-sonnet-5" and reached the picker — a model the policy never named, and one the guardrail denies on sight, since enforcement compares the id as written. Strip only the namespaces a gateway is known to prepend, taken from the first slash rather than the last. The e2e retry loops slept between attempts without watching the context, so a cancelled run kept retrying calls that fail instantly and spent its remaining window sleeping between them. They now stop when the context is done. The streamed provider's setup key outlived its test: deleting the group does not delete the key that auto-joins it. --- e2e/agentnetwork/custom_pricing_test.go | 12 ++++-- e2e/agentnetwork/main_test.go | 16 ++++++++ e2e/agentnetwork/streaming_test.go | 7 +++- proxy/internal/proxy/discovery_filter.go | 38 +++++++++++++++---- proxy/internal/proxy/discovery_filter_test.go | 18 +++++++++ 5 files changed, 80 insertions(+), 11 deletions(-) diff --git a/e2e/agentnetwork/custom_pricing_test.go b/e2e/agentnetwork/custom_pricing_test.go index 33bb13db7..1083aee04 100644 --- a/e2e/agentnetwork/custom_pricing_test.go +++ b/e2e/agentnetwork/custom_pricing_test.go @@ -163,7 +163,9 @@ func chatOnce(t *testing.T, ctx context.Context, env pricedEnv, model, sessionID break } } - time.Sleep(5 * time.Second) + if !waitBeforeRetry(ctx, 5*time.Second) { + break + } } require.Equal(t, 200, code, "chat for %s must return 200; body: %s\n=== proxy logs ===\n%s", model, body, env.proxy.Logs(context.Background())) @@ -430,7 +432,9 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) { lastSession = fmt.Sprintf("e2e-session-reprice-b-%d", time.Now().UnixNano()) code, _, cerr := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, customModel, "Reply with exactly: pong", lastSession) if cerr != nil || code != 200 { - time.Sleep(5 * time.Second) + if !waitBeforeRetry(ctx, 5*time.Second) { + break + } continue } row, ok := lookupAccessLogBySession(ctx, lastSession, repriceIngestWindow) @@ -448,7 +452,9 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) { } // Still priced at the old rate — the push hasn't landed yet; retry. lastCost, sawRow = row.InputCostUsd, true - time.Sleep(5 * time.Second) + if !waitBeforeRetry(ctx, 5*time.Second) { + break + } } lastSeen := "no row was ever read" if sawRow { diff --git a/e2e/agentnetwork/main_test.go b/e2e/agentnetwork/main_test.go index cc366b3fb..687af1d4d 100644 --- a/e2e/agentnetwork/main_test.go +++ b/e2e/agentnetwork/main_test.go @@ -54,3 +54,19 @@ func run(m *testing.M) int { return m.Run() } + +// waitBeforeRetry pauses between attempts of a polling loop and reports +// whether the caller should keep going. A cancelled context ends the loop +// where a plain sleep would keep retrying against it: every call fails +// instantly once ctx is done, so the loop would spend its whole remaining +// window sleeping between failures nobody is waiting for any more. +func waitBeforeRetry(ctx context.Context, d time.Duration) bool { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/e2e/agentnetwork/streaming_test.go b/e2e/agentnetwork/streaming_test.go index 4637da10c..a5fa8df3f 100644 --- a/e2e/agentnetwork/streaming_test.go +++ b/e2e/agentnetwork/streaming_test.go @@ -129,6 +129,9 @@ func provisionStreamingProvider(t *testing.T, ctx context.Context, catalogID str Ephemeral: &ephemeral, }) require.NoError(t, err, "mint setup key") + // Deleting the group does not delete the key it auto-joins, so the key + // needs a cleanup of its own. + t.Cleanup(func() { _ = srv.API().SetupKeys.Delete(context.Background(), sk.Id) }) require.NotEmpty(t, sk.Key, "setup key plaintext") dummyKey := "sk-stream-e2e" @@ -195,7 +198,9 @@ func chatStreamUntil(t *testing.T, ctx context.Context, env pricedEnv, kind, mod break } } - time.Sleep(5 * time.Second) + if !waitBeforeRetry(ctx, 5*time.Second) { + break + } } if code != 200 { t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background())) diff --git a/proxy/internal/proxy/discovery_filter.go b/proxy/internal/proxy/discovery_filter.go index 7b5f8d415..c9d606970 100644 --- a/proxy/internal/proxy/discovery_filter.go +++ b/proxy/internal/proxy/discovery_filter.go @@ -152,19 +152,43 @@ func entryPermitted(entry map[string]json.RawMessage, permitted map[string]struc return false } +// gatewayNamespaces are the provider prefixes a gateway prepends to a model +// it re-exports: LiteLLM lists a Bedrock model the operator registered as +// "anthropic.claude-opus-5" under "bedrock/anthropic.claude-opus-5". Only +// these are stripped before matching. +// +// A slash is not by itself a namespace separator. Self-hosted backends ship +// ids that carry one ("Qwen/Qwen2.5-0.5B-Instruct"), and an upstream is free +// to scope ids per tenant ("tenant-b/claude-sonnet-5"). Treating every slash +// as a prefix let any such id match an allowed model by its tail, so the +// picker offered models the policy never named. +var gatewayNamespaces = map[string]struct{}{ + "anthropic": {}, + "azure": {}, + "bedrock": {}, + "mistral": {}, + "openai": {}, + "vertex_ai": {}, +} + // modelIDForms returns the forms a single model id may be written in: the id -// itself, its undated form, and the same two with a gateway's provider -// prefix removed ("vertex_ai/claude-sonnet-5"). The bare id is tried first, -// because a self-hosted id can legitimately contain a slash of its own -// ("Qwen/Qwen2.5-0.5B-Instruct") and must not be cut down to its tail. +// itself, its undated form, and — when the id is namespaced by a gateway we +// recognise — the same two with that namespace removed +// ("vertex_ai/claude-sonnet-5"). The bare id is always tried first. +// +// The namespace is what precedes the FIRST slash: it is a prefix the gateway +// put in front of the whole id, and everything after it is the id the +// operator would have registered, separators included. func modelIDForms(id string) []string { if id == "" { return nil } forms := []string{id, sharedllm.NormalizeAnthropicModel(id)} - if slash := strings.LastIndex(id, "/"); slash >= 0 { - tail := id[slash+1:] - forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail)) + if slash := strings.Index(id, "/"); slash > 0 { + if _, ok := gatewayNamespaces[id[:slash]]; ok { + tail := id[slash+1:] + forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail)) + } } return forms } diff --git a/proxy/internal/proxy/discovery_filter_test.go b/proxy/internal/proxy/discovery_filter_test.go index 85dddf230..103eac594 100644 --- a/proxy/internal/proxy/discovery_filter_test.go +++ b/proxy/internal/proxy/discovery_filter_test.go @@ -171,6 +171,24 @@ func TestModelDiscoveryFilter_KeepsSlashBearingIDs(t *testing.T) { "a slash inside the model id is part of the id, not a provider prefix") } +// TestModelDiscoveryFilter_RejectsTailMatchOnUnknownNamespace covers the id +// an upstream scopes with a prefix of its own. "tenant-b/claude-sonnet-5" +// ends in a model the policy permits, but it is a different model on a +// different tenant, and the guardrail denies that string outright — so +// offering it hands the picker an entry the next request refuses. +func TestModelDiscoveryFilter_RejectsTailMatchOnUnknownNamespace(t *testing.T) { + ids := listedIDs(t, []string{"claude-sonnet-5"}, `{ + "data": [ + {"id": "claude-sonnet-5"}, + {"id": "tenant-b/claude-sonnet-5"}, + {"id": "Qwen/claude-sonnet-5"} + ] + }`) + + assert.Equal(t, []string{"claude-sonnet-5"}, ids, + "only a namespace a gateway is known to prepend may be stripped before matching") +} + // TestModelDiscoveryFilter_ForwardsOversizedBodyIntact covers a listing past // the buffering cap. The filter reads one byte beyond the cap to detect the // size; forwarding only what it read would hand the client a body truncated