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