From 766fcae3f8a9d6ba445fe8b2f2d872506f8d72e3 Mon Sep 17 00:00:00 2001 From: Maycon Santos Date: Sun, 23 Aug 2026 20:02:33 +0200 Subject: [PATCH] [proxy,management] Conform the Agent Network endpoint to the LLM gateway protocol (#7154) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [proxy,management] Conform the Agent Network endpoint to the LLM gateway protocol Reviewed the proxy against Claude Code's published gateway contract. The transport layer already held up; fourteen gaps sat one layer up, in the model catalog and in the non-inference endpoints clients call. Two of them cost money. The catalog carried no claude-opus-5 or claude-sonnet-5, so an operator could not authorise the models coding agents default to — those requests denied as not-routable, or priced at zero where a catch-all carried them. And gateway records pin ParserID "openai" while the same record serves /v1/messages, so Anthropic responses were read with the OpenAI parser, which never looks at message_start where input tokens live: input metered as roughly zero on every stream and cost was skipped entirely. The rest fix requests refused for structural rather than policy reasons: model discovery denied for every account with a model allowlist, token counting denied on Bedrock and mis-parsed on Vertex, startup probes refused and written into the access log at every session start, and denials rendered in a shape no LLM client parses. Two changes are additive by design — the deny body keeps every field it had and adds the vendor's error object alongside, and body-level identity injection is now gated on the request's dialect so it stops sending OpenAI-shape fields into Anthropic bodies that reject them. The end-to-end work turned up one more: the discovery filter treated any slash in a model id as a gateway prefix, which would have dropped every self-hosted "Qwen/..." model from the picker. --- agent-network/README.md | 29 ++ e2e/agentnetwork/custom_pricing_test.go | 189 +++++++- e2e/agentnetwork/gateway_protocol_test.go | 455 ++++++++++++++++++ e2e/agentnetwork/gateway_review_test.go | 242 ++++++++++ e2e/agentnetwork/main_test.go | 16 + e2e/agentnetwork/streaming_test.go | 209 ++++++++ e2e/harness/client.go | 53 +- e2e/harness/vllm.go | 135 +++++- .../modules/agentnetwork/catalog/catalog.go | 6 + .../agentnetwork/catalog/catalog_test.go | 36 ++ .../modules/agentnetwork/pricing/defaults.go | 6 - .../pricing/defaults_llm_pricing.example.yaml | 10 + .../agentnetwork/pricing/defaults_test.go | 8 +- proxy/internal/llm/model.go | 8 + proxy/internal/llm/pricing/pricing.go | 16 +- proxy/internal/llm/pricing/pricing_test.go | 19 + .../builtin/cost_meter/middleware.go | 18 +- .../builtin/llm_guardrail/middleware.go | 26 +- .../builtin/llm_guardrail/middleware_test.go | 49 ++ .../builtin/llm_identity_inject/middleware.go | 31 ++ .../llm_identity_inject/middleware_test.go | 54 +++ .../builtin/llm_limit_check/middleware.go | 14 +- .../llm_limit_check/middleware_test.go | 32 ++ .../llm_request_parser/bedrock_test.go | 26 + .../builtin/llm_request_parser/middleware.go | 76 ++- .../llm_request_parser/middleware_test.go | 105 ++++ .../builtin/llm_router/bedrock_route_test.go | 87 ++++ .../builtin/llm_router/middleware.go | 302 +++++++++--- .../builtin/llm_router/middleware_test.go | 276 ++++++++++- proxy/internal/middleware/decision.go | 68 +++ proxy/internal/middleware/decision_test.go | 92 ++++ proxy/internal/middleware/keys.go | 17 + proxy/internal/middleware/types.go | 12 + proxy/internal/proxy/discovery_filter.go | 215 +++++++++ proxy/internal/proxy/discovery_filter_test.go | 235 +++++++++ proxy/internal/proxy/reverseproxy.go | 3 + proxy/server.go | 26 +- shared/llm/model.go | 21 + shared/llm/model_test.go | 26 + 39 files changed, 3094 insertions(+), 154 deletions(-) create mode 100644 e2e/agentnetwork/gateway_protocol_test.go create mode 100644 e2e/agentnetwork/gateway_review_test.go create mode 100644 e2e/agentnetwork/streaming_test.go create mode 100644 management/internals/modules/agentnetwork/catalog/catalog_test.go create mode 100644 proxy/internal/middleware/decision_test.go create mode 100644 proxy/internal/proxy/discovery_filter.go create mode 100644 proxy/internal/proxy/discovery_filter_test.go diff --git a/agent-network/README.md b/agent-network/README.md index 1997ea299..5211fe8f9 100644 --- a/agent-network/README.md +++ b/agent-network/README.md @@ -40,6 +40,35 @@ You can then use this private endpoint to configure your AI agents, whether that Full step-by-step setup: **https://docs.netbird.io/agent-network/quickstart** +## Client settings that don't follow the endpoint + +Most of an agent's traffic follows the base URL you hand it, but a few +client-side checks call their vendor directly and never reach the proxy. On a +network that blocks direct egress they fail even though inference works, so +they are worth setting once when you roll the endpoint out. + +For Claude Code: + +- **Fast mode** checks availability against `api.anthropic.com` rather than the + configured base URL. Set `CLAUDE_CODE_SKIP_FAST_MODE_ORG_CHECK=1` when the + agent authenticates with `ANTHROPIC_AUTH_TOKEN` alone (the usual shape when + the proxy injects the real provider key) or when a TLS-inspecting proxy + answers the check itself. Set + `CLAUDE_CODE_SKIP_FAST_MODE_NETWORK_ERRORS=1` when the network refuses the + connection outright. Fast mode is an Anthropic-API feature, so it is + unavailable on a Bedrock- or Vertex-backed endpoint whatever you set. +- **Model discovery** is off by default. Set + `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` for the picker to list the + models your policies authorise; the proxy filters the response to that set. + The client gives discovery a three-second budget and treats any redirect as + a failure, so the endpoint must serve `/v1/models` directly. +- **The WebFetch domain safety check** also calls `api.anthropic.com` directly + and is unaffected by the variables above. + +Allowing direct egress to `api.anthropic.com` covers the network cases but not +the credential one, where the check reaches Anthropic and is rejected because +the agent presents a proxy-issued key. + ## Architecture Agent Network is built on two existing NetBird capabilities: diff --git a/e2e/agentnetwork/custom_pricing_test.go b/e2e/agentnetwork/custom_pricing_test.go index e3750258f..b3ca5028f 100644 --- a/e2e/agentnetwork/custom_pricing_test.go +++ b/e2e/agentnetwork/custom_pricing_test.go @@ -23,9 +23,10 @@ import ( // model the client asks for. The proxy prices off the REQUEST model, not the // upstream response model, so a made-up model id billed at operator rates lets // these tests assert exact costs without a real vendor key. +// Sourced from the harness so the counts can't drift from the mock's config. const ( - vllmPromptTokens = 11 - vllmCompletionTokens = 2 + vllmPromptTokens = harness.VLLMChatInputTokens + vllmCompletionTokens = harness.VLLMChatOutputTokens ) // pricedEnv is a connected single-provider agent-network deployment pointed at @@ -162,30 +163,85 @@ 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())) return body } -// findAccessLogBySession polls the access-log page for the row carrying sessionID. -func findAccessLogBySession(t *testing.T, ctx context.Context, sessionID string) api.AgentNetworkAccessLog { - t.Helper() - var row api.AgentNetworkAccessLog - require.Eventually(t, func() bool { - logs, lerr := srv.ListAccessLogs(ctx) - if lerr != nil { - return false - } - for _, r := range logs.Data { - if r.SessionId != nil && *r.SessionId == sessionID { - row = r - return true +// 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 + +// accessLogPollInterval is how long the lookup waits between pages. Ingest is +// asynchronous, so the row lands somewhere inside the window rather than on +// any particular poll. +const accessLogPollInterval = 2 * time.Second + +// lookupAccessLogBySession polls the access-log page for the row carrying +// sessionID and reports whether it arrived within the window. It never fails +// the test: callers that can recover — by firing a fresh request under a new +// session — need to see the miss rather than die on it. +func lookupAccessLogBySession(ctx context.Context, sessionID string, within time.Duration) (api.AgentNetworkAccessLog, bool) { + deadline := time.Now().Add(within) + for { + // Each poll is bounded by what is left of the window rather than by the + // caller's context: a single stalled request would otherwise hold the + // loop open long past the ingest window it is meant to enforce, and the + // caller would read the delay as a missing row. + if logs, lerr := listAccessLogsBy(ctx, deadline); lerr == nil { + for _, r := range logs.Data { + if r.SessionId != nil && *r.SessionId == sessionID { + return r, true + } } } - return false - }, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row", sessionID) + // The wait is bounded by the window as well, so the answer arrives when + // the caller's budget runs out rather than a poll interval later: a + // full interval slept past the deadline reports "no row" up to two + // seconds late, which reads as a slower lookup than the one asked for. + wait := time.Until(deadline) + if wait > accessLogPollInterval { + wait = accessLogPollInterval + } + if wait <= 0 { + return api.AgentNetworkAccessLog{}, false + } + timer := time.NewTimer(wait) + select { + case <-ctx.Done(): + timer.Stop() + return api.AgentNetworkAccessLog{}, false + case <-timer.C: + } + // Checked after the wait rather than before the request: a poll issued + // past the deadline carries no budget and would fail on arrival. + if !time.Now().Before(deadline) { + return api.AgentNetworkAccessLog{}, false + } + } +} + +// listAccessLogsBy fetches one access-log page under a context that expires at +// deadline, so no single call can outlive the window its caller is polling +// within. The parent's cancellation still applies: the child inherits it. +func listAccessLogsBy(ctx context.Context, deadline time.Time) (api.AgentNetworkAccessLogsResponse, error) { + reqCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + return srv.ListAccessLogs(reqCtx) +} + +// findAccessLogBySession polls the access-log page for the row carrying +// sessionID, failing the test if it never lands. Use it for a request whose row +// must exist; where a missing row is a recoverable race, use +// lookupAccessLogBySession and retry. +func findAccessLogBySession(t *testing.T, ctx context.Context, sessionID string) api.AgentNetworkAccessLog { + t.Helper() + row, ok := lookupAccessLogBySession(ctx, sessionID, accessLogIngestWindow) + require.True(t, ok, "session id %q must be recorded in an access-log row", sessionID) return row } @@ -319,6 +375,11 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) { outRateA = 0.020 inRateB = 0.050 // 5x / 4x the original, so a repriced row is unmistakable outRateB = 0.080 + // Per-attempt ingest wait, shorter than the default so a request that + // produces no row costs one retry rather than most of the budget, and an + // overall deadline long enough to hold several attempts. + repriceIngestWindow = 20 * time.Second + repriceDeadline = 180 * time.Second ) env := provisionPricedProvider(t, ctx, "reprice", []api.AgentNetworkProviderModel{ @@ -353,27 +414,61 @@ func TestPriceChangeUpdatesRecordedCost(t *testing.T) { // reading its cost, so an un-ingested row is never mistaken for "still rate A". // The expected new input cost is unmistakably higher than rate A, so a // lingering old-rate row can't satisfy the check. + // + // Every way an iteration can come up short — the request failing, its row not + // landing, or the row still carrying rate A — is a symptom of the same + // in-flight rebuild, so each one retries under a fresh session rather than + // ending the test. Only the outer deadline is fatal. wantInputB := float64(vllmPromptTokens) / 1000 * inRateB var repriced api.AgentNetworkAccessLog var lastSession string - deadline := time.Now().Add(90 * time.Second) + // The cost last read, kept separately: repriced is the zero value on every + // path that gives up, so reporting its cost would say "$0.000000" whether + // the rows were still at rate A or no row was ever read. + var lastCost float64 + var sawRow bool + deadline := time.Now().Add(repriceDeadline) + // Everything inside the loop runs under the deadline rather than the + // test's own context. An attempt started just before it would otherwise + // run well past it: the chat container is capped at 90s of its own and the + // row lookup at another 20s, so the loop could report a repricing failure + // nearly two minutes after the window it was given had closed. + repriceCtx, cancelReprice := context.WithDeadline(ctx, deadline) + defer cancelReprice() for time.Now().Before(deadline) { 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) + code, _, cerr := env.client.Chat(repriceCtx, env.endpoint, env.proxyIP, harness.WireChat, customModel, "Reply with exactly: pong", lastSession) if cerr != nil || code != 200 { - time.Sleep(5 * time.Second) + if !waitBeforeRetry(repriceCtx, 5*time.Second) { + break + } + continue + } + row, ok := lookupAccessLogBySession(repriceCtx, lastSession, repriceIngestWindow) + if !ok { + // No row for this request. The proxy now publishes a rebuilt chain + // before the route that reaches it, so a request can no longer be + // served unattributed mid-update; this retry covers the ingest + // window alone. Fire another one under a fresh session. + t.Logf("no access-log row for session %q within %s; retrying under a fresh session", lastSession, repriceIngestWindow) continue } - row := findAccessLogBySession(t, ctx, lastSession) if inDelta(row.InputCostUsd, wantInputB, 1e-6) { repriced = row break } // Still priced at the old rate — the push hasn't landed yet; retry. - time.Sleep(5 * time.Second) + lastCost, sawRow = row.InputCostUsd, true + if !waitBeforeRetry(repriceCtx, 5*time.Second) { + break + } } - require.NotEmpty(t, repriced.Id, "a request after the price change must be priced at the new rate B; last input_cost_usd=$%.6f, wanted $%.6f\n=== proxy logs ===\n%s", - repriced.InputCostUsd, wantInputB, env.proxy.Logs(context.Background())) + lastSeen := "no row was ever read" + if sawRow { + lastSeen = fmt.Sprintf("last input_cost_usd=$%.6f", lastCost) + } + require.NotEmpty(t, repriced.Id, "a request after the price change must be priced at the new rate B; %s, wanted $%.6f\n=== proxy logs ===\n%s", + lastSeen, wantInputB, env.proxy.Logs(context.Background())) assertOpenAICostAtRates(t, repriced, inRateB, outRateB) verifyUsageRowForSession(t, lastSession, inRateB, outRateB) @@ -630,3 +725,47 @@ func inDelta(a, b, tol float64) bool { } return d <= tol } + +// TestCustomDatedModelKeepsItsOwnPrice covers the review fix that anchored the +// release-date fallback to Claude ids. Pricing looks every model up through +// that helper, so while it matched a bare trailing date any operator id ending +// in eight digits inherited the rate of its undated sibling — a silent +// mis-bill on models NetBird knows nothing about. +func TestCustomDatedModelKeepsItsOwnPrice(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + const ( + baseModel = "internal-llm" + datedModel = "internal-llm-20250101" + baseIn = 0.010 + baseOut = 0.020 + // An order of magnitude apart, so a row billed at the wrong entry is + // unmistakable rather than a rounding argument. + datedIn = 0.100 + datedOut = 0.200 + ) + + env := provisionPricedProvider(t, ctx, "customdated", []api.AgentNetworkProviderModel{ + {Id: baseModel, InputPer1k: baseIn, OutputPer1k: baseOut}, + {Id: datedModel, InputPer1k: datedIn, OutputPer1k: datedOut}, + }) + + t.Run("the undated id bills at its own rate", func(t *testing.T) { + session := fmt.Sprintf("e2e-session-customdated-base-%d", time.Now().UnixNano()) + chatOnce(t, ctx, env, baseModel, session) + assertOpenAICostAtRates(t, findAccessLogBySession(t, ctx, session), baseIn, baseOut) + }) + + t.Run("the dated id keeps its own rate", func(t *testing.T) { + session := fmt.Sprintf("e2e-session-customdated-dated-%d", time.Now().UnixNano()) + chatOnce(t, ctx, env, datedModel, session) + row := findAccessLogBySession(t, ctx, session) + assertOpenAICostAtRates(t, row, datedIn, datedOut) + + // Spelled out because it is the regression: inheriting the sibling's + // rate would bill this request at a tenth of its price. + assert.Greater(t, row.InputCostUsd, float64(vllmPromptTokens)/1000*baseIn*2, + "a custom dated id must not inherit the undated entry's rate") + }) +} diff --git a/e2e/agentnetwork/gateway_protocol_test.go b/e2e/agentnetwork/gateway_protocol_test.go new file mode 100644 index 000000000..c21a4fc53 --- /dev/null +++ b/e2e/agentnetwork/gateway_protocol_test.go @@ -0,0 +1,455 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "fmt" + "strings" + "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" +) + +// Models each catalog surface is registered with in the matrix below. They +// differ per provider so the router's choice is unambiguous: a request that +// lands on the wrong provider record fails the surface assertion instead of +// passing by coincidence. +const ( + matrixAnthropicModel = "claude-sonnet-5" + matrixBedrockModel = "anthropic.claude-sonnet-5" + // matrixBedrockPathModel is what a Bedrock SDK client puts in the URL: a + // cross-region inference profile with a release date and version suffix. + // The proxy must normalise it back to matrixBedrockModel to route and price. + matrixBedrockPathModel = "us.anthropic.claude-sonnet-5-20250101-v1:0" + // matrixVertexModel differs from the Anthropic record's model on purpose: + // a shared id would leave two routes claiming it and make which one serves + // /v1/messages depend on declaration order. + matrixVertexModel = "claude-haiku-4-5" + matrixVertexProject = "e2e-project" + matrixVertexRegion = "us-east5" +) + +// gatewayEnv is a connected client plus a set of provider records, all pointed +// at one mock upstream, so several wire shapes can be driven over a single +// tunnel. +type gatewayEnv struct { + endpoint string + proxyIP string + client *harness.Client + proxy *harness.Proxy + vllm *harness.VLLM + // providerIDs maps the catalog id to the created provider record id. + providerIDs map[string]string +} + +// provisionGatewayMatrix brings up one mock upstream and one provider record +// per catalog surface, all authorised for the same group by a single policy. +// Sharing one proxy and client keeps the wire-shape cases to one tunnel setup; +// each case still creates its own session id so its access-log row is findable. +func provisionGatewayMatrix(t *testing.T, ctx context.Context) gatewayEnv { + t.Helper() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-matrix"}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-gw-matrix-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + // The mock ignores auth, so a dummy credential satisfies each catalog + // entry's auth template. Vertex is the exception: its api_key is a GCP + // service-account keyfile the proxy mints an OAuth token from, and a dummy + // one cannot mint. That is deliberate — the Vertex case below asserts on + // routing, which happens before the token mint. + dummyKey := "sk-gw-e2e" + dummyKeyfile := "keyfile::" + "e2e-not-a-real-service-account-key" + + specs := []struct { + name string + catalogID string + apiKey string + models []api.AgentNetworkProviderModel + }{ + { + name: "openai", catalogID: "openai_api", apiKey: dummyKey, + models: []api.AgentNetworkProviderModel{{Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002}}, + }, + { + name: "anthropic", catalogID: "anthropic_api", apiKey: dummyKey, + models: []api.AgentNetworkProviderModel{{Id: matrixAnthropicModel, InputPer1k: 0.003, OutputPer1k: 0.015}}, + }, + { + name: "bedrock", catalogID: "bedrock_api", apiKey: dummyKey, + models: []api.AgentNetworkProviderModel{{Id: matrixBedrockModel, InputPer1k: 0.003, OutputPer1k: 0.015}}, + }, + { + name: "vertex", catalogID: "vertex_ai_api", apiKey: dummyKeyfile, + models: []api.AgentNetworkProviderModel{{Id: matrixVertexModel, InputPer1k: 0.001, OutputPer1k: 0.005}}, + }, + } + + providerIDs := make(map[string]string, len(specs)) + ids := make([]string, 0, len(specs)) + for _, spec := range specs { + key := spec.apiKey + models := spec.models + prov, perr := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-gw-" + spec.name, + ProviderId: spec.catalogID, + UpstreamUrl: vllm.URL, + ApiKey: &key, + Enabled: ptr(true), + Models: &models, + }) + require.NoError(t, perr, "create %s provider", spec.name) + id := prov.Id + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), id) }) + providerIDs[spec.catalogID] = id + ids = append(ids, id) + } + + // Uncapped token limit: never blocks the handful of tokens driven here, but + // switches on usage metering so consumption and cost land in the row. + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-gw-matrix", + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: ids, + Limits: &api.AgentNetworkPolicyLimits{ + TokenLimit: api.AgentNetworkPolicyTokenLimit{ + Enabled: true, + GroupCap: 10_000_000, + UserCap: 10_000_000, + WindowSeconds: 60, + }, + }, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-matrix", sk.Key) + return gatewayEnv{ + endpoint: endpoint, + proxyIP: proxyIP, + client: cl, + proxy: px, + vllm: vllm, + providerIDs: providerIDs, + } +} + +// connectClient starts a proxy and a tunnel client for the shared account and +// waits until the client can reach the proxy peer, returning the endpoint and +// the proxy's tunnel IP to pin requests to. +func connectClient(t *testing.T, ctx context.Context, name, setupKey string) (string, string, *harness.Client, *harness.Proxy) { + t.Helper() + + settings, err := srv.GetSettings(ctx) + require.NoError(t, err, "read settings") + require.NotEmpty(t, settings.Endpoint, "endpoint must be assigned") + + proxyToken, err := srv.CreateProxyTokenCLI(ctx, "e2e-"+name+"-proxy") + require.NoError(t, err, "mint proxy token") + px, err := harness.StartProxy(ctx, srv, proxyToken) + require.NoError(t, err, "start proxy") + t.Cleanup(func() { _ = px.Terminate(context.Background()) }) + + cl, err := harness.StartClient(ctx, srv, setupKey) + require.NoError(t, err, "start client") + t.Cleanup(func() { _ = cl.Terminate(context.Background()) }) + + require.NoError(t, cl.WaitConnected(ctx, 90*time.Second), "client must connect to management") + // The probe resolves the endpoint and its first packet wakes the lazy proxy + // peer, so WaitProxyPeer then observes it connected. + proxyIP, err := cl.ResolveProxyIP(ctx, settings.Endpoint) + require.NoError(t, err, "resolve endpoint to proxy IP") + if err := cl.WaitProxyPeer(ctx, 180*time.Second); err != nil { + t.Fatalf("client did not see the proxy peer: %v\n=== proxy logs ===\n%s", err, px.Logs(context.Background())) + } + return settings.Endpoint, proxyIP, cl, px +} + +// callUntil retries an HTTP call through the tunnel until it returns one of the +// wanted statuses or the deadline passes, absorbing the DNS and tunnel jitter +// the first call through a fresh tunnel can hit. The last status and body are +// returned either way so the caller can assert with real detail. +func callUntil(t *testing.T, call func() (int, string, error), want ...int) (int, string) { + t.Helper() + wanted := make(map[int]struct{}, len(want)) + for _, w := range want { + wanted[w] = struct{}{} + } + + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + c, b, err := call() + if err == nil { + code, body = c, b + if _, ok := wanted[code]; ok { + return code, body + } + } + time.Sleep(5 * time.Second) + } + return code, body +} + +// TestGatewayProtocolProviderMatrix drives one request per wire shape over a +// single tunnel, with a provider record per catalog surface behind it. It is +// the regression net for the routing and parser-selection changes: each case +// asserts the surface the request was metered under and the token counts that +// surface's own usage block carries, so a request parsed by the wrong provider's +// parser meters zero and fails rather than passing on a coincidence. +func TestGatewayProtocolProviderMatrix(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + env := provisionGatewayMatrix(t, ctx) + diag := func() string { + return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s", + env.vllm.Logs(context.Background()), env.proxy.Logs(context.Background())) + } + + t.Run("openai chat completions", func(t *testing.T) { + session := "e2e-gw-openai" + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireChat, harness.VLLMModel, "ping", session) + }, 200) + require.Equal(t, 200, code, "openai chat must succeed; body: %s%s", body, diag()) + require.Contains(t, body, "chat.completion", "body must be an OpenAI completion; got: %s", body) + + row := findAccessLogBySession(t, ctx, session) + require.NotNil(t, row.Provider) + assert.Equal(t, "openai", *row.Provider, "the OpenAI chat path must meter under the openai surface") + assert.Equal(t, int64(harness.VLLMChatInputTokens), row.InputTokens, "OpenAI usage block must be read") + assert.Equal(t, int64(harness.VLLMChatOutputTokens), row.OutputTokens) + }) + + t.Run("anthropic messages", func(t *testing.T) { + session := "e2e-gw-anthropic" + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, matrixAnthropicModel, "ping", session) + }, 200) + require.Equal(t, 200, code, "anthropic messages must succeed; body: %s%s", body, diag()) + + row := findAccessLogBySession(t, ctx, session) + require.NotNil(t, row.Provider) + assert.Equal(t, "anthropic", *row.Provider, "the /v1/messages path must meter under the anthropic surface") + // These counts only appear if the Anthropic parser read the response: + // its usage fields are named differently from the OpenAI block. + assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens, + "Anthropic input_tokens must be read; zero here means the wrong parser ran") + assert.Equal(t, int64(harness.VLLMMessagesOutputTokens), row.OutputTokens) + assert.Positive(t, row.CachedInputTokens, "the Anthropic cache-read bucket must be recorded") + assert.Positive(t, row.CostUsd, "a metered request must carry a cost") + require.NotNil(t, row.ResolvedProviderId) + assert.Equal(t, env.providerIDs["anthropic_api"], *row.ResolvedProviderId, + "a vendor-tagged request must not cross to another provider's record") + }) + + t.Run("bedrock invoke normalises the path model", func(t *testing.T) { + session := "e2e-gw-bedrock" + code, body := callUntil(t, func() (int, string, error) { + return env.client.Bedrock(ctx, env.endpoint, env.proxyIP, matrixBedrockPathModel, "ping", session) + }, 200) + require.Equal(t, 200, code, "bedrock invoke must succeed; body: %s%s", body, diag()) + + row := findAccessLogBySession(t, ctx, session) + require.NotNil(t, row.Provider) + assert.Equal(t, "bedrock", *row.Provider, "a native Bedrock path must meter under the bedrock surface") + require.NotNil(t, row.Model) + assert.Equal(t, matrixBedrockModel, *row.Model, + "the inference-profile prefix, release date and version suffix must be normalised away") + assert.Equal(t, int64(harness.VLLMMessagesInputTokens), row.InputTokens) + }) + + t.Run("anthropic token counting", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/messages/count_tokens", + fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"ping"}]}`, matrixAnthropicModel), + []string{"anthropic-version: 2023-06-01"}) + }, 200) + assert.Equal(t, 200, code, "token counting must route rather than deny; body: %s%s", body, diag()) + }) + + t.Run("bedrock token counting", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, + "/model/"+matrixBedrockPathModel+"/count-tokens", + `{"input":{"converse":{"messages":[{"role":"user","content":[{"text":"ping"}]}]}}}`, nil) + }, 200) + assert.Equal(t, 200, code, + "the Bedrock count-tokens action must route; denying it pushes counting onto the billable inference path; body: %s%s", + body, diag()) + }) + + t.Run("vertex token counting reaches its provider", func(t *testing.T) { + // The dummy service-account key cannot mint an OAuth token, so the + // request stops at the upstream credential. Both outcomes render as + // 403, so the deny code is what distinguishes them: upstream_auth_failed + // means the path resolved to the Vertex route and only the credential + // failed, while model_not_routable would mean the method segment was + // swallowed into the model id and no route ever claimed it. + path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s/count-tokens:rawPredict", + matrixVertexProject, matrixVertexRegion, matrixVertexModel) + _, body := callUntil(t, func() (int, string, error) { + return env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path, + `{"anthropic_version":"vertex-2023-10-16","messages":[{"role":"user","content":"ping"}]}`, nil) + }, 403) + assert.NotContains(t, body, "model_not_routable", + "the count-tokens method segment must not be parsed as part of the model id; body: %s%s", body, diag()) + assert.Contains(t, body, "llm_policy.upstream_auth_failed", + "the request must reach the Vertex route and fail only at the credential; body: %s%s", body, diag()) + }) + + t.Run("connection warming probe", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.Get(ctx, env.endpoint, env.proxyIP, "/api/hello", nil) + }, 200) + assert.NotEqual(t, 403, code, + "the warm-up probe carries no model and must not be refused as unroutable; body: %s%s", body, diag()) + }) + + t.Run("unknown model denies in the caller's error shape", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, + "claude-not-a-real-model-9", "ping", "e2e-gw-unknown") + }, 403) + require.Equal(t, 403, code, "a model no provider claims must still be refused; body: %s%s", body, diag()) + + // The NetBird fields stay where they were for existing consumers. + assert.Contains(t, body, "llm_policy.model_not_routable", "the deny code must be preserved") + // And the vendor's own envelope rides alongside, so the client can show + // the reason instead of an unexplained API error. + assert.Contains(t, body, `"type":"error"`, "an Anthropic caller must get the Anthropic error envelope") + assert.Contains(t, body, "permission_error", "403 must map to the vendor's permission error type") + }) +} + +// TestModelDiscoveryWithModelAllowlist covers gateway model discovery on an +// account that restricts models, which is the configuration that broke: the +// listing carries no model, and the per-model allowlist fails closed on an +// undetermined one, so discovery denied for exactly the accounts using the +// feature. It also asserts the allowlist still refuses a model outside it, so +// the exemption cannot be read as a way around the gate. +func TestModelDiscoveryWithModelAllowlist(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gw-discovery"}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-gw-discovery-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + // One provider enumerating a single model, while the upstream's own listing + // advertises two. The proxy must serve the shorter list. + dummyKey := "sk-discovery-e2e" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-gw-discovery", + ProviderId: "openai_api", + UpstreamUrl: vllm.URL, + ApiKey: &dummyKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{ + {Id: harness.VLLMModel, InputPer1k: 0.001, OutputPer1k: 0.002}, + }, + }) + require.NoError(t, err, "create provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + // The model allowlist is what makes this a regression test: without a + // guardrail enabled, discovery was never gated in the first place. + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-gw-discovery-allowlist" + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel} + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) }) + + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-gw-discovery", + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + GuardrailIds: &[]string{guard.Id}, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + endpoint, proxyIP, cl, px := connectClient(t, ctx, "gw-discovery", sk.Key) + diag := func() string { + return fmt.Sprintf("\n=== upstream logs ===\n%s\n=== proxy logs ===\n%s", + vllm.Logs(context.Background()), px.Logs(context.Background())) + } + + t.Run("listing is served and bounded by policy", func(t *testing.T) { + 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 not be refused because the request carries no model; body: %s%s", body, diag()) + + assert.Contains(t, body, harness.VLLMModel, "the authorised model must reach the picker") + assert.NotContains(t, body, harness.VLLMUnlistedModel, + "a model the policy does not authorise must not be offered; body: %s", body) + }) + + t.Run("allowlist still refuses a model outside it", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat, + harness.VLLMUnlistedModel, "ping", "e2e-gw-discovery-blocked") + }, 403) + require.Equal(t, 403, code, + "exempting model-less endpoints must not exempt inference; body: %s%s", body, diag()) + assert.True(t, + strings.Contains(body, "llm_policy.model_blocked") || strings.Contains(body, "llm_policy.model_not_routable"), + "the refusal must name a model policy code; body: %s", body) + }) + + t.Run("allowlisted model still routes", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return cl.Chat(ctx, endpoint, proxyIP, harness.WireChat, + harness.VLLMModel, "ping", "e2e-gw-discovery-allowed") + }, 200) + require.Equal(t, 200, code, "the allowlisted model must still be served; body: %s%s", body, diag()) + }) +} diff --git a/e2e/agentnetwork/gateway_review_test.go b/e2e/agentnetwork/gateway_review_test.go new file mode 100644 index 000000000..556bc4a53 --- /dev/null +++ b/e2e/agentnetwork/gateway_review_test.go @@ -0,0 +1,242 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "fmt" + "strings" + "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" +) + +// The cases in this file cover behaviour that arrived from code review, after +// the gateway-protocol end-to-end tests were written. Each had unit coverage +// only; none needed a new harness capability, which is why they belong here +// rather than on a manual checklist. + +// TestNonInferenceEndpointsAreAuthorised covers the two review findings on the +// endpoints that carry no body: the per-model lookup must be authorised +// against the same allowlist that bounds the listing beside it, and only a read +// method may claim the non-inference exemption that skips the token pre-flight. +func TestNonInferenceEndpointsAreAuthorised(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + env := provisionDiscoveryProvider(t, ctx) + + t.Run("lookup of an authorised model succeeds", func(t *testing.T) { + code, body := callUntil(t, func() (int, string, error) { + return env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMModel, nil) + }, 200) + assert.Equal(t, 200, code, "an allowlisted model must remain reachable; body: %s", body) + }) + + t.Run("lookup of an unauthorised model is refused", func(t *testing.T) { + code, body, err := env.client.Get(ctx, env.endpoint, env.proxyIP, "/v1/models/"+harness.VLLMUnlistedModel, nil) + require.NoError(t, err, "request must reach the proxy") + assert.Equal(t, 403, code, + "a model the policy does not authorise must not be confirmed by the detail lookup; body: %s", body) + }) + + // A write must not claim the exemption that lets the listing skip the token + // pre-flight. The body names no model on purpose: that is what a request + // probing for the exemption looks like, and it is the case the method gate + // exists to refuse. (A POST that does name a model is a different thing — + // it routes and meters as the inference request it is.) + for _, path := range []string{"/v1/models", "/v1/models/" + harness.VLLMModel, "/api/hello"} { + t.Run("write to "+path+" is refused", func(t *testing.T) { + code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, path, + `{"messages":[{"role":"user","content":"hi"}]}`, nil) + require.NoError(t, err, "request must reach the proxy") + assert.NotEqual(t, 200, code, + "a write to a non-inference path must not be served unmetered; body: %s", body) + }) + } + + // A request carrying the sub-agent attribution headers must still be served + // and metered normally. Asserting the ids themselves is not possible yet: + // the parser lifts them onto the request's metadata, but nothing persists + // them, so they have no queryable surface to check against. + t.Run("sub-agent headers do not disturb the request", func(t *testing.T) { + sessionID := fmt.Sprintf("e2e-session-agentid-%d", time.Now().UnixNano()) + code, body, err := env.client.PostJSON(ctx, env.endpoint, env.proxyIP, "/v1/chat/completions", + fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"Reply with exactly: pong"}]}`, harness.VLLMModel), + []string{ + "x-session-id: " + sessionID, + "x-claude-code-agent-id: agent-child-7", + "x-claude-code-parent-agent-id: agent-root-1", + }) + require.NoError(t, err, "request must reach the proxy") + require.Equal(t, 200, code, "the request must succeed; body: %s", body) + + row := findAccessLogBySession(t, ctx, sessionID) + assert.Positive(t, row.InputTokens, "the request must still be metered normally") + }) +} + +// TestDatedModelIdRouting covers both halves of the dated-id rule that review +// tightened: a dated id still reaches an undated registration, but a route +// pinned to one dated build must never serve a different one. +func TestDatedModelIdRouting(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + const ( + undated = "claude-sonnet-9" + datedA = "claude-sonnet-9-20250101" + datedB = "claude-sonnet-9-20250202" + ) + + t.Run("a dated id reaches its undated registration", func(t *testing.T) { + env := provisionModelProvider(t, ctx, "dated-undated", "anthropic_api", undated) + + sessionID := fmt.Sprintf("e2e-session-dated-%d", time.Now().UnixNano()) + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", sessionID) + }, 200) + require.Equal(t, 200, code, "a pinned release of a registered family must route; body: %s", body) + + row := findAccessLogBySession(t, ctx, sessionID) + assert.Positive(t, row.InputTokens, "the dated request must price at the registered rate, not zero") + }) + + t.Run("a route pinned to one dated build refuses another", func(t *testing.T) { + env := provisionModelProvider(t, ctx, "dated-pinned", "anthropic_api", datedA) + + code, body := callUntil(t, func() (int, string, error) { + return env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedA, "Reply with exactly: pong", "") + }, 200) + require.Equal(t, 200, code, "the exact dated id must still route; body: %s", body) + + code, body, err := env.client.Chat(ctx, env.endpoint, env.proxyIP, harness.WireMessages, datedB, "Reply with exactly: pong", "") + require.NoError(t, err, "request must reach the proxy") + assert.Equal(t, 403, code, + "a provider pinned to one dated build must not serve another; body: %s", body) + }) +} + +// TestBedrockInferenceProfilesReachTheUpstream covers the startup lookup a +// Bedrock client makes. The proxy forwards it to the configured upstream rather +// than denying it, so what comes back is the upstream's answer — never a +// NetBird policy rejection. +func TestBedrockInferenceProfilesReachTheUpstream(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Minute) + defer cancel() + + env := provisionModelProvider(t, ctx, "infprofiles", "bedrock_api", "anthropic.claude-sonnet-5") + + code, body := callUntil(t, func() (int, string, error) { + return env.client.Get(ctx, env.endpoint, env.proxyIP, "/inference-profiles", nil) + }, 200) + + assert.Equal(t, 200, code, "the lookup must reach the upstream; body: %s", body) + assert.NotContains(t, body, "llm_policy.", + "the proxy must not answer a control-plane lookup with a policy denial") + assert.Contains(t, body, "inferenceProfileSummaries", + "the upstream's own answer must come back untouched") +} + +// provisionDiscoveryProvider brings up one mock-backed provider enumerating a +// single model, with an allowlist guardrail in effect, plus a connected client. +func provisionDiscoveryProvider(t *testing.T, ctx context.Context) pricedEnv { + t.Helper() + env := provisionModelProvider(t, ctx, "noninference", "openai_api", harness.VLLMModel) + + var gr api.AgentNetworkGuardrailRequest + gr.Name = "e2e-noninference-allowlist-" + fmt.Sprint(time.Now().UnixNano()) + gr.Checks.ModelAllowlist.Enabled = true + gr.Checks.ModelAllowlist.Models = []string{harness.VLLMModel} + guard, err := srv.CreateGuardrail(ctx, gr) + require.NoError(t, err, "create guardrail") + t.Cleanup(func() { _ = srv.DeleteGuardrail(context.Background(), guard.Id) }) + + enabled := true + _, err = srv.UpdatePolicy(ctx, env.policyID, api.AgentNetworkPolicyRequest{ + Name: "e2e-noninference", + Enabled: &enabled, + SourceGroups: []string{env.groupID}, + DestinationProviderIds: []string{env.providerID}, + GuardrailIds: &[]string{guard.Id}, + }) + require.NoError(t, err, "attach guardrail to policy") + return env +} + +// provisionModelProvider brings up the mock, one provider under the given +// catalog id enumerating exactly one model, an authorising policy, and a +// connected proxy + client. +func provisionModelProvider(t *testing.T, ctx context.Context, name, catalogID, model string) pricedEnv { + t.Helper() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + suffix := strings.ToLower(name) + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-gwr-" + suffix}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-gwr-" + suffix + "-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + Ephemeral: &ephemeral, + }) + require.NoError(t, err, "mint setup key") + require.NotEmpty(t, sk.Key, "setup key plaintext") + + dummyKey := "sk-gwr-e2e" + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: "e2e-gwr-" + suffix, + ProviderId: catalogID, + UpstreamUrl: vllm.URL, + ApiKey: &dummyKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{ + {Id: model, InputPer1k: 0.001, OutputPer1k: 0.002}, + }, + }) + require.NoError(t, err, "create provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-gwr-" + suffix, + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + Limits: &api.AgentNetworkPolicyLimits{ + TokenLimit: api.AgentNetworkPolicyTokenLimit{ + Enabled: true, + GroupCap: 10_000_000, + UserCap: 10_000_000, + WindowSeconds: 60, + }, + }, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + endpoint, proxyIP, cl, px := connectClient(t, ctx, "gwr-"+suffix, sk.Key) + return pricedEnv{ + providerID: prov.Id, + groupID: grp.Id, + policyID: pol.Id, + upstream: vllm.URL, + endpoint: endpoint, + proxyIP: proxyIP, + client: cl, + proxy: px, + } +} 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 new file mode 100644 index 000000000..a5fa8df3f --- /dev/null +++ b/e2e/agentnetwork/streaming_test.go @@ -0,0 +1,209 @@ +//go:build e2e + +package agentnetwork + +import ( + "context" + "fmt" + "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" +) + +// streamedModel is priced high enough that a mis-metered request is obvious in +// the recorded cost, and named so it cannot collide with another test's route. +const streamedModel = "e2e-streamed-model" + +const ( + streamInRate = 0.010 + streamOutRate = 0.020 + // The cache-read bucket is priced separately from input, so a run that + // folded the two together fails the per-bucket assertions below. + streamCacheReadRate = 0.001 +) + +// TestStreamingResponseMetersInputTokens is the end-to-end guard for the +// metering bug this endpoint's gateway-protocol work fixed. +// +// On a streamed answer the input-token count exists only in the opening +// message_start event; every later frame reports output. A response read with +// the wrong vendor's parser — the shape a gateway record produces when it names +// one API surface and serves another — never looks at that event, so input +// metered as zero and the bulk of the bill silently vanished. Nothing in the +// suite sent stream: true before this test, so the whole branch went unrun. +// +// The provider points at the mock's streaming listener, which answers every +// request as SSE with token counts that differ from the buffered surface. That +// difference is the point: passing these assertions is only possible if the +// stream accumulator ran. +func TestStreamingResponseMetersInputTokens(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + env := provisionStreamingProvider(t, ctx, "anthropic_api") + + sessionID := fmt.Sprintf("e2e-session-stream-%d", time.Now().UnixNano()) + code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID) + require.Equal(t, 200, code, "streamed chat must succeed; body: %s", body) + assert.Contains(t, body, "message_start", + "the client must receive the event stream itself, not a buffered rewrite of it") + + row := findAccessLogBySession(t, ctx, sessionID) + + assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens), + "input tokens live in message_start; zero here is the bug this test exists for") + assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens), + "output tokens ride message_delta and supersede the message_start seed") + assert.Equal(t, harness.VLLMStreamCacheReadTokens, int(row.CachedInputTokens), + "the Anthropic cache bucket rides message_start too, and only its own parser reads it") + + // The Anthropic surface bills cache reads additively, so the input bucket + // prices the full input count rather than a remainder. + wantInput := float64(harness.VLLMStreamInputTokens) / 1000 * streamInRate + wantOutput := float64(harness.VLLMStreamOutputTokens) / 1000 * streamOutRate + wantCacheRead := float64(harness.VLLMStreamCacheReadTokens) / 1000 * streamCacheReadRate + assert.InDelta(t, wantInput, row.InputCostUsd, 1e-6, "input cost must price the streamed input tokens") + assert.InDelta(t, wantOutput, row.OutputCostUsd, 1e-6, "output cost must price the streamed output tokens") + // The total, not merely a positive number: input and output alone are + // positive, so a cache bucket parsed and then never billed would pass any + // weaker assertion. The gap is 7e-6, well outside the delta. + assert.InDelta(t, wantInput+wantOutput+wantCacheRead, row.CostUsd, 1e-6, + "the recorded cost must be every bucket the surface bills, cache reads included") +} + +// TestStreamingOnGatewayTypedProvider drives the same streamed Anthropic call +// through a provider record whose catalog id names the OpenAI surface — the +// exact misconfiguration that hid the bug, since gateway records commonly pin +// one parser while the upstream serves another shape entirely. +// +// The router must choose the parser from the request path rather than the +// record's provider id, or the Anthropic usage block goes unread and input +// meters at zero all over again. +func TestStreamingOnGatewayTypedProvider(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) + defer cancel() + + env := provisionStreamingProvider(t, ctx, "openai_api") + + sessionID := fmt.Sprintf("e2e-session-stream-gw-%d", time.Now().UnixNano()) + code, body := chatStreamUntil(t, ctx, env, harness.WireMessages, streamedModel, sessionID) + require.Equal(t, 200, code, "streamed chat through a gateway record must succeed; body: %s", body) + + row := findAccessLogBySession(t, ctx, sessionID) + + assert.Equal(t, harness.VLLMStreamInputTokens, int(row.InputTokens), + "a record typed openai_api must still read the Anthropic usage block it is actually serving") + assert.Equal(t, harness.VLLMStreamOutputTokens, int(row.OutputTokens), + "output tokens must survive the surface mismatch too") + assert.InDelta(t, float64(harness.VLLMStreamInputTokens)/1000*streamInRate, row.InputCostUsd, 1e-6, + "the request must be priced on the surface it spoke, not the one the record names") +} + +// provisionStreamingProvider brings up the mock, one provider pointed at its +// streaming listener under the given catalog id, a policy authorising it, and a +// connected proxy + client. +func provisionStreamingProvider(t *testing.T, ctx context.Context, catalogID string) pricedEnv { + t.Helper() + + vllm, err := harness.StartVLLM(ctx, srv) + require.NoError(t, err, "start mock upstream") + t.Cleanup(func() { _ = vllm.Terminate(context.Background()) }) + + name := "stream-" + catalogID + grp, err := srv.API().Groups.Create(ctx, api.PostApiGroupsJSONRequestBody{Name: "e2e-" + name}) + require.NoError(t, err, "create group") + t.Cleanup(func() { _ = srv.API().Groups.Delete(context.Background(), grp.Id) }) + + ephemeral := false + sk, err := srv.API().SetupKeys.Create(ctx, api.PostApiSetupKeysJSONRequestBody{ + Name: "e2e-" + name + "-client", + Type: "reusable", + ExpiresIn: 86400, + UsageLimit: 0, + AutoGroups: []string{grp.Id}, + 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" + cacheRead := streamCacheReadRate + prov, err := srv.CreateProvider(ctx, api.AgentNetworkProviderRequest{ + Name: name, + ProviderId: catalogID, + UpstreamUrl: vllm.StreamURL, + ApiKey: &dummyKey, + Enabled: ptr(true), + Models: &[]api.AgentNetworkProviderModel{{ + Id: streamedModel, + InputPer1k: streamInRate, + OutputPer1k: streamOutRate, + CacheReadPer1k: &cacheRead, + }}, + }) + require.NoError(t, err, "create provider") + t.Cleanup(func() { _ = srv.DeleteProvider(context.Background(), prov.Id) }) + + enabled := true + pol, err := srv.CreatePolicy(ctx, api.AgentNetworkPolicyRequest{ + Name: "e2e-" + name, + Enabled: &enabled, + SourceGroups: []string{grp.Id}, + DestinationProviderIds: []string{prov.Id}, + Limits: &api.AgentNetworkPolicyLimits{ + TokenLimit: api.AgentNetworkPolicyTokenLimit{ + Enabled: true, + GroupCap: 10_000_000, + UserCap: 10_000_000, + WindowSeconds: 60, + }, + }, + }) + require.NoError(t, err, "create policy") + t.Cleanup(func() { _ = srv.DeletePolicy(context.Background(), pol.Id) }) + + endpoint, proxyIP, cl, px := connectClient(t, ctx, name, sk.Key) + return pricedEnv{ + providerID: prov.Id, + groupID: grp.Id, + policyID: pol.Id, + upstream: vllm.StreamURL, + endpoint: endpoint, + proxyIP: proxyIP, + client: cl, + proxy: px, + } +} + +// chatStreamUntil drives one streamed chat, retrying to absorb the tunnel and +// DNS jitter a first call through a fresh peer can hit. +func chatStreamUntil(t *testing.T, ctx context.Context, env pricedEnv, kind, model, sessionID string) (int, string) { + t.Helper() + var code int + var body string + deadline := time.Now().Add(90 * time.Second) + for time.Now().Before(deadline) { + c, b, cerr := env.client.ChatStream(ctx, env.endpoint, env.proxyIP, kind, model, "Reply with exactly: pong", sessionID) + if cerr == nil { + code, body = c, b + if code == 200 { + break + } + } + if !waitBeforeRetry(ctx, 5*time.Second) { + break + } + } + if code != 200 { + t.Logf("=== proxy logs ===\n%s", env.proxy.Logs(context.Background())) + } + return code, body +} diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 0d7f016a6..9e9e7b34a 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "net/http" "os/exec" "strconv" "strings" @@ -292,6 +293,27 @@ func (cl *Client) ChatPrefixed(ctx context.Context, endpoint, proxyIP, pathPrefi return cl.post(ctx, endpoint, proxyIP, pathPrefix+path, body, withSessionID(headers, sessionID)) } +// ChatStream is Chat with "stream": true in the request body, so the proxy's +// request parser marks the call as streaming and its response parser takes the +// SSE accumulator rather than the buffered-body path. Pair it with a provider +// pointed at VLLM.StreamURL, which answers every request as an event stream. +func (cl *Client) ChatStream(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) { + var path, body string + var headers []string + switch kind { + case WireMessages: + path = "/v1/messages" + headers = []string{"anthropic-version: 2023-06-01"} + body = fmt.Sprintf(`{"model":%q,"max_tokens":2048,"stream":true,"messages":[{"role":"user","content":%q}]}`, model, prompt) + default: + path = "/v1/chat/completions" + // include_usage is what makes a real OpenAI stream emit its final usage + // frame; without it the last chunk carries no tokens at all. + body = fmt.Sprintf(`{"model":%q,"stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":%q}]}`, model, prompt) + } + return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID)) +} + // Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike // Chat, the model is carried in the request path (project/region/model), so the // proxy routes by path and mints the service-account OAuth token; the body uses @@ -322,10 +344,29 @@ func withSessionID(headers []string, sessionID string) []string { return append(headers, "x-session-id: "+sessionID) } -// post runs curl in a throwaway container sharing the client's network -// namespace so the request traverses the WireGuard tunnel, pinning the endpoint -// to the proxy IP. It returns the HTTP status and response body. +// Get issues a GET to the agent-network endpoint over the client's tunnel. +// Model discovery and the connection-warming probe are read-only endpoints +// that carry no body, so they can't go through the chat helpers. +func (cl *Client) Get(ctx context.Context, endpoint, proxyIP, path string, extraHeaders []string) (int, string, error) { + return cl.do(ctx, http.MethodGet, endpoint, proxyIP, path, "", extraHeaders) +} + +// PostJSON issues an arbitrary JSON POST over the client's tunnel, for wire +// shapes the typed helpers don't cover (token counting, say). +func (cl *Client) PostJSON(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) { + return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders) +} + +// post issues a JSON POST. Retained as the shorthand the chat helpers use. func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) { + return cl.do(ctx, http.MethodPost, endpoint, proxyIP, path, body, extraHeaders) +} + +// do runs curl in a throwaway container sharing the client's network +// namespace so the request traverses the WireGuard tunnel, pinning the endpoint +// to the proxy IP. It returns the HTTP status and response body. An empty body +// sends no payload, which is what a GET needs. +func (cl *Client) do(ctx context.Context, method, endpoint, proxyIP, path, body string, extraHeaders []string) (int, string, error) { url := "https://" + endpoint + path args := []string{ "run", "--rm", @@ -334,13 +375,15 @@ func (cl *Client) post(ctx context.Context, endpoint, proxyIP, path, body string "-sk", "--connect-timeout", "5", "--max-time", "90", "--resolve", endpoint + ":443:" + proxyIP, "-o", "/dev/stderr", "-w", "%{http_code}", - "-X", "POST", url, + "-X", method, url, "-H", "Content-Type: application/json", } for _, h := range extraHeaders { args = append(args, "-H", h) } - args = append(args, "--data", body) + if body != "" { + args = append(args, "--data", body) + } cmd := exec.CommandContext(ctx, "docker", args...) // -w writes the status code to stdout; -o /dev/stderr writes the body to // stderr so we can capture both separately. diff --git a/e2e/harness/vllm.go b/e2e/harness/vllm.go index 2f3d306cc..cf9316325 100644 --- a/e2e/harness/vllm.go +++ b/e2e/harness/vllm.go @@ -18,18 +18,63 @@ const ( vllmImage = "nginx:alpine" vllmAlias = "vllm" vllmPort = "8000/tcp" + // vllmStreamPort serves the same wire shapes as an SSE stream. See the + // nginx config for why streaming lives on its own listener. + vllmStreamPort = "8001/tcp" // VLLMModel is the served model id the mock advertises and echoes back. It // matches a real small model commonly served by vLLM so the provider's // enumerated model and the client's request line up. VLLMModel = "Qwen/Qwen2.5-0.5B-Instruct" + // VLLMUnlistedModel is a second id the mock's model listing advertises but + // no test provider enumerates, so a filtered listing is observably shorter + // than the upstream's own. + VLLMUnlistedModel = "Qwen/Qwen2.5-7B-Instruct" +) + +// Token counts the mock reports per wire shape. Tests assert on these rather +// than on "> 0" so a response parsed with the wrong provider's parser (which +// would read a different field, or none) fails loudly instead of passing on +// a coincidental non-zero. +const ( + // VLLMChatInputTokens / VLLMChatOutputTokens ride the OpenAI usage block. + VLLMChatInputTokens = 11 + VLLMChatOutputTokens = 2 + // VLLMMessagesInputTokens / VLLMMessagesOutputTokens ride the Anthropic + // usage block, whose field names the OpenAI parser cannot read. + VLLMMessagesInputTokens = 17 + VLLMMessagesOutputTokens = 3 +) + +// Token counts the streaming surface reports. They differ from the +// non-streaming ones on purpose: a test that asserts these numbers proves the +// SSE accumulator ran, rather than a buffered JSON body having been parsed. +// +// Input and cache-read arrive on message_start; output arrives on +// message_delta and supersedes the seed value message_start carries. Any +// parser that cannot read message_start reports zero input tokens — which is +// exactly the bug these counts exist to catch. +const ( + VLLMStreamInputTokens = 29 + VLLMStreamOutputTokens = 5 + VLLMStreamCacheReadTokens = 7 ) // vllmNginxConf emulates a vLLM OpenAI-compatible server over plain HTTP (vLLM's -// default: no TLS, port 8000). It answers /v1/models with a one-model list and -// any chat/completions path with a canned OpenAI-shaped chat completion carrying -// a non-zero usage block, so the proxy's OpenAI parser records real token -// consumption. Running actual vLLM in CI is infeasible (GPU + multi-GB model +// default: no TLS, port 8000), and additionally answers the wire shapes the +// other catalog surfaces speak so one mock can stand in for every provider the +// proxy routes to. Running actual vLLM in CI is infeasible (GPU + multi-GB model // download), so this stands in for the wire contract the proxy depends on. +// +// Each shape answers with its own vendor's usage block, so a response parsed +// under the wrong surface meters zero rather than passing by accident: +// +// - /v1/chat/completions (and any unmatched path): OpenAI chat completion. +// - /v1/messages: Anthropic Messages, snake_case usage plus a cache bucket. +// - /model/{id}/invoke: Bedrock InvokeModel, which carries the Anthropic body. +// - the token-counting endpoints: a count, with no usage block at all. +// +// The model listing advertises two models so a policy that authorises one +// produces an observably shorter list than the upstream's own. const vllmNginxConf = `pid /tmp/nginx.pid; events {} http { @@ -37,13 +82,75 @@ http { listen 8000; location = /v1/models { default_type application/json; - return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"}]}'; + return 200 '{"object":"list","data":[{"id":"Qwen/Qwen2.5-0.5B-Instruct","object":"model","owned_by":"vllm"},{"id":"Qwen/Qwen2.5-7B-Instruct","object":"model","owned_by":"vllm"}]}'; + } + location = /v1/messages { + default_type application/json; + return 200 '{"id":"msg_e2e","type":"message","role":"assistant","model":"claude-sonnet-5","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}'; + } + location = /v1/messages/count_tokens { + default_type application/json; + return 200 '{"input_tokens":7}'; + } + location ~ ^/model/.+/invoke$ { + default_type application/json; + return 200 '{"id":"msg_e2e_bedrock","type":"message","role":"assistant","content":[{"type":"text","text":"pong"}],"stop_reason":"end_turn","usage":{"input_tokens":17,"output_tokens":3,"cache_read_input_tokens":5}}'; + } + location ~ ^/model/.+/count-tokens$ { + default_type application/json; + return 200 '{"inputTokens":9}'; + } + location = /api/hello { + return 200; + } + location = /inference-profiles { + default_type application/json; + return 200 '{"inferenceProfileSummaries":[{"inferenceProfileId":"us.anthropic.claude-sonnet-5","status":"ACTIVE"}]}'; } location / { default_type application/json; return 200 '{"id":"chatcmpl-e2e-vllm","object":"chat.completion","created":1700000000,"model":"Qwen/Qwen2.5-0.5B-Instruct","choices":[{"index":0,"message":{"role":"assistant","content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":2,"total_tokens":13}}'; } } + + # The streaming surface, on its own port so the response content type is a + # property of the listener rather than of a per-request branch: nginx sets + # Content-Type from default_type, which cannot be varied inside an "if", and + # a second Content-Type via add_header would leave the proxy reading the + # wrong one. A provider record pointed at this port streams every answer. + # + # Input and cache-read tokens ride message_start, output rides message_delta + # — the split that makes a stream different from a buffered body, and the + # reason a parser that ignores message_start meters input as zero. + server { + listen 8001; + location = /v1/messages { + default_type text/event-stream; + return 200 'event: message_start +data: {"type":"message_start","message":{"id":"msg_e2e_stream","type":"message","role":"assistant","model":"claude-sonnet-5","content":[],"usage":{"input_tokens":29,"output_tokens":1,"cache_read_input_tokens":7}}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"pong"}} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":5}} + +event: message_stop +data: {"type":"message_stop"} + +'; + } + location / { + default_type text/event-stream; + return 200 'data: {"choices":[{"delta":{"content":"pong"}}]} + +data: {"choices":[],"usage":{"prompt_tokens":29,"completion_tokens":5,"total_tokens":34}} + +data: [DONE] + +'; + } + } } ` @@ -55,6 +162,10 @@ type VLLM struct { workDir string // URL is the upstream URL the vllm provider points at (http://:8000). URL string + // StreamURL is the same mock's streaming listener. A provider pointed here + // answers every request as SSE, so the proxy's streaming accumulator runs + // instead of its buffered-body parser. + StreamURL string } // StartVLLM runs the mock vLLM server on the shared network over plain HTTP. @@ -73,14 +184,17 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) { req := testcontainers.ContainerRequest{ Image: vllmImage, - ExposedPorts: []string{vllmPort}, + ExposedPorts: []string{vllmPort, vllmStreamPort}, Networks: []string{c.network.Name}, NetworkAliases: map[string][]string{c.network.Name: {vllmAlias}}, Cmd: []string{"nginx", "-c", "/conf/nginx.conf", "-g", "daemon off;"}, HostConfigModifier: func(hc *container.HostConfig) { hc.Binds = append(hc.Binds, workDir+":/conf:ro") }, - WaitingFor: wait.ForListeningPort(vllmPort).WithStartupTimeout(60 * time.Second), + WaitingFor: wait.ForAll( + wait.ForListeningPort(vllmPort), + wait.ForListeningPort(vllmStreamPort), + ).WithStartupTimeout(60 * time.Second), } ctr, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ @@ -92,7 +206,12 @@ func StartVLLM(ctx context.Context, c *Combined) (*VLLM, error) { return nil, fmt.Errorf("start vllm container: %w", err) } - return &VLLM{container: ctr, workDir: workDir, URL: "http://" + vllmAlias + ":8000"}, nil + return &VLLM{ + container: ctr, + workDir: workDir, + URL: "http://" + vllmAlias + ":8000", + StreamURL: "http://" + vllmAlias + ":8001", + }, nil } // Logs returns the vLLM container logs, for diagnostics on failure. diff --git a/management/internals/modules/agentnetwork/catalog/catalog.go b/management/internals/modules/agentnetwork/catalog/catalog.go index 2c4efd0b4..c534f9a85 100644 --- a/management/internals/modules/agentnetwork/catalog/catalog.go +++ b/management/internals/modules/agentnetwork/catalog/catalog.go @@ -296,6 +296,8 @@ var providers = []Provider{ // account to be on >= 30-day data retention or all requests // 400. Models: []Model{ + {ID: "claude-opus-5", Label: "Claude Opus 5", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, + {ID: "claude-sonnet-5", Label: "Claude Sonnet 5", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000}, {ID: "claude-fable-5", Label: "Claude Fable 5", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000}, {ID: "claude-opus-4-8", Label: "Claude Opus 4.8", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, {ID: "claude-opus-4-7", Label: "Claude Opus 4.7", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, @@ -355,6 +357,8 @@ var providers = []Provider{ // Llama 3.3 70B entry kept unchanged — LiteLLM tracks only // per-region Llama 3 entries; standalone 3.3 not yet listed. Models: []Model{ + {ID: "anthropic.claude-opus-5", Label: "Claude Opus 5 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, + {ID: "anthropic.claude-sonnet-5", Label: "Claude Sonnet 5 (Bedrock)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000}, {ID: "anthropic.claude-opus-4-8", Label: "Claude Opus 4.8 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, {ID: "anthropic.claude-opus-4-7", Label: "Claude Opus 4.7 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, {ID: "anthropic.claude-opus-4-6", Label: "Claude Opus 4.6 (Bedrock)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, @@ -406,6 +410,8 @@ var providers = []Provider{ // exists — the router denies unmeterable publishers rather than forward // them uncounted. Models: []Model{ + {ID: "claude-opus-5", Label: "Claude Opus 5 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, + {ID: "claude-sonnet-5", Label: "Claude Sonnet 5 (Vertex)", InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003, CacheCreationPer1k: 0.00375, ContextWindow: 1000000}, {ID: "claude-fable-5", Label: "Claude Fable 5 (Vertex)", InputPer1k: 0.010, OutputPer1k: 0.050, CacheReadPer1k: 0.001, CacheCreationPer1k: 0.0125, ContextWindow: 1000000}, {ID: "claude-opus-4-8", Label: "Claude Opus 4.8 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, {ID: "claude-opus-4-7", Label: "Claude Opus 4.7 (Vertex)", InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625, ContextWindow: 1000000}, diff --git a/management/internals/modules/agentnetwork/catalog/catalog_test.go b/management/internals/modules/agentnetwork/catalog/catalog_test.go new file mode 100644 index 000000000..e4e887e6f --- /dev/null +++ b/management/internals/modules/agentnetwork/catalog/catalog_test.go @@ -0,0 +1,36 @@ +package catalog + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestClaudeLineupSelectable pins the models Claude Code resolves to by +// default. A model absent from the lineup can't be ticked on a provider +// record, so llm_router denies it as not-routable and the operator has no +// way to authorise the client's own default. +func TestClaudeLineupSelectable(t *testing.T) { + for providerID, wanted := range map[string][]string{ + "anthropic_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"}, + "bedrock_api": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5", "anthropic.claude-haiku-4-5"}, + "vertex_ai_api": {"claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"}, + } { + provider, ok := Lookup(providerID) + require.True(t, ok, "catalog must define %s", providerID) + + selectable := make(map[string]Model, len(provider.Models)) + for _, m := range provider.Models { + selectable[m.ID] = m + } + for _, id := range wanted { + model, found := selectable[id] + require.True(t, found, "%s must offer %s", providerID, id) + assert.NotEmpty(t, model.Label, "%s/%s needs a label for the picker", providerID, id) + assert.Positive(t, model.InputPer1k, "%s/%s needs an input rate", providerID, id) + assert.Positive(t, model.OutputPer1k, "%s/%s needs an output rate", providerID, id) + assert.Positive(t, model.ContextWindow, "%s/%s needs a context window", providerID, id) + } + } +} diff --git a/management/internals/modules/agentnetwork/pricing/defaults.go b/management/internals/modules/agentnetwork/pricing/defaults.go index c690313bc..315cfe208 100644 --- a/management/internals/modules/agentnetwork/pricing/defaults.go +++ b/management/internals/modules/agentnetwork/pricing/defaults.go @@ -47,17 +47,11 @@ var supplementalDefaults = map[string]map[string]Entry{ "gpt-5-nano": {InputPer1k: 0.00005, OutputPer1k: 0.0004, CachedInputPer1k: 0.000005}, }, "anthropic": { - // claude-opus-5 is not yet in the catalog lineup but gateway / - // grandfathered traffic uses it; priced so it isn't skipped. - "claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625}, // "kimi-k3[1m]" is the 1M-context alias some Claude Code guides // configure against Moonshot's Anthropic-compatible endpoint; // priced identically to kimi-k3 so those requests aren't skipped. "kimi-k3[1m]": {InputPer1k: 0.003, OutputPer1k: 0.015, CacheReadPer1k: 0.0003}, }, - "bedrock": { - "anthropic.claude-opus-5": {InputPer1k: 0.005, OutputPer1k: 0.025, CacheReadPer1k: 0.0005, CacheCreationPer1k: 0.00625}, - }, } var ( diff --git a/management/internals/modules/agentnetwork/pricing/defaults_llm_pricing.example.yaml b/management/internals/modules/agentnetwork/pricing/defaults_llm_pricing.example.yaml index bb1cb09a8..78830ae3c 100644 --- a/management/internals/modules/agentnetwork/pricing/defaults_llm_pricing.example.yaml +++ b/management/internals/modules/agentnetwork/pricing/defaults_llm_pricing.example.yaml @@ -82,6 +82,11 @@ anthropic: output_per_1k: 0.015 cache_read_per_1k: 0.0003 cache_creation_per_1k: 0.00375 + claude-sonnet-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 kimi-k3: input_per_1k: 0.003 output_per_1k: 0.015 @@ -145,6 +150,11 @@ bedrock: output_per_1k: 0.015 cache_read_per_1k: 0.0003 cache_creation_per_1k: 0.00375 + anthropic.claude-sonnet-5: + input_per_1k: 0.003 + output_per_1k: 0.015 + cache_read_per_1k: 0.0003 + cache_creation_per_1k: 0.00375 meta.llama3-3-70b-instruct: input_per_1k: 0.00072 output_per_1k: 0.00072 diff --git a/management/internals/modules/agentnetwork/pricing/defaults_test.go b/management/internals/modules/agentnetwork/pricing/defaults_test.go index 99c965687..04b6de550 100644 --- a/management/internals/modules/agentnetwork/pricing/defaults_test.go +++ b/management/internals/modules/agentnetwork/pricing/defaults_test.go @@ -116,11 +116,13 @@ func TestDefaultTable_PinnedRates(t *testing.T) { assert.InDelta(t, 0.010, fable.InputPer1k, 1e-9, "claude-fable-5 input") assert.InDelta(t, 0.0125, fable.CacheCreationPer1k, 1e-9, "claude-fable-5 cache creation") - // Supplementals present on their surfaces. + // Every id below must stay priced whichever source provides it: the + // catalog lineup for the current Claude 5 family, supplementalDefaults + // for the ids the dashboard deliberately doesn't offer. for surface, ids := range map[string][]string{ "openai": {"gpt-5", "gpt-5-mini", "gpt-5-nano"}, - "anthropic": {"claude-opus-5", "kimi-k3[1m]", "kimi-k3"}, - "bedrock": {"anthropic.claude-opus-5"}, + "anthropic": {"claude-opus-5", "claude-sonnet-5", "kimi-k3[1m]", "kimi-k3"}, + "bedrock": {"anthropic.claude-opus-5", "anthropic.claude-sonnet-5"}, } { for _, id := range ids { _, ok := table[surface][id] diff --git a/proxy/internal/llm/model.go b/proxy/internal/llm/model.go index 76ccfeccf..2e056a57a 100644 --- a/proxy/internal/llm/model.go +++ b/proxy/internal/llm/model.go @@ -13,6 +13,14 @@ func NormalizeBedrockModel(modelID string) string { return sharedllm.NormalizeBedrockModel(modelID) } +// NormalizeAnthropicModel strips the trailing "-YYYYMMDD" release-date suffix +// from an Anthropic model id so a dated id a client pins matches the undated +// one the operator registered. Thin delegate to shared/llm for the same +// contract reason as the two below. +func NormalizeAnthropicModel(modelID string) string { + return sharedllm.NormalizeAnthropicModel(modelID) +} + // NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id // so it matches the catalog/pricing key. Thin delegate to shared/llm, kept // beside NormalizeBedrockModel for the same contract reason. diff --git a/proxy/internal/llm/pricing/pricing.go b/proxy/internal/llm/pricing/pricing.go index ce6e636cf..52cedb60e 100644 --- a/proxy/internal/llm/pricing/pricing.go +++ b/proxy/internal/llm/pricing/pricing.go @@ -10,6 +10,8 @@ package pricing import ( "fmt" "math" + + sharedllm "github.com/netbirdio/netbird/shared/llm" ) // Entry is a single model's input and output pricing, expressed in USD per @@ -92,7 +94,10 @@ func NewTable(raw map[string]map[string]EntryJSON) (*Table, error) { return &Table{entries: entries}, nil } -// Lookup returns the entry for the given provider surface and model. +// Lookup returns the entry for the given provider surface and model. A +// dated Anthropic id falls back to its undated form, so a client pinning +// "claude-sonnet-4-5-20250929" bills at the registered "claude-sonnet-4-5" +// rate instead of recording no cost at all. func (t *Table) Lookup(provider, model string) (Entry, bool) { if t == nil { return Entry{}, false @@ -101,7 +106,14 @@ func (t *Table) Lookup(provider, model string) (Entry, bool) { if !ok { return Entry{}, false } - e, ok := byModel[model] + if e, found := byModel[model]; found { + return e, true + } + undated := sharedllm.NormalizeAnthropicModel(model) + if undated == model { + return Entry{}, false + } + e, ok := byModel[undated] return e, ok } diff --git a/proxy/internal/llm/pricing/pricing_test.go b/proxy/internal/llm/pricing/pricing_test.go index b946faa7f..e7d339f06 100644 --- a/proxy/internal/llm/pricing/pricing_test.go +++ b/proxy/internal/llm/pricing/pricing_test.go @@ -175,3 +175,22 @@ func TestNewTable_NilAndEmpty(t *testing.T) { require.NoError(t, err) assert.Empty(t, entries, "nil in, empty (never-matching) map out for the per-record map") } + +// TestLookup_DatedAnthropicIDFallsBackToUndated covers a client pinning a +// release date on a model priced under its undated id. Without the +// fallback the request records no cost at all. +func TestLookup_DatedAnthropicIDFallsBackToUndated(t *testing.T) { + table, err := NewTable(map[string]map[string]EntryJSON{ + "anthropic": { + "claude-sonnet-4-5": {InputPer1K: 0.003, OutputPer1K: 0.015}, + }, + }) + require.NoError(t, err, "table must build from a valid defaults map") + + entry, ok := table.Lookup("anthropic", "claude-sonnet-4-5-20250929") + require.True(t, ok, "a dated id must resolve to the undated entry") + assert.InDelta(t, 0.003, entry.InputPer1K, 1e-9, "dated id must bill at the registered rate") + + _, ok = table.Lookup("anthropic", "claude-sonnet-9-9-20250929") + assert.False(t, ok, "an unknown family must stay unpriced") +} diff --git a/proxy/internal/middleware/builtin/cost_meter/middleware.go b/proxy/internal/middleware/builtin/cost_meter/middleware.go index 2ce706cda..8e2e0590c 100644 --- a/proxy/internal/middleware/builtin/cost_meter/middleware.go +++ b/proxy/internal/middleware/builtin/cost_meter/middleware.go @@ -11,6 +11,7 @@ import ( "fmt" "strconv" + "github.com/netbirdio/netbird/proxy/internal/llm" "github.com/netbirdio/netbird/proxy/internal/llm/pricing" "github.com/netbirdio/netbird/proxy/internal/middleware" ) @@ -175,13 +176,28 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar // Anthropic route still bills its cache buckets additively. func (m *Middleware) lookupCosts(md []middleware.KV, surface, model string, inTokens, outTokens, cachedTokens, cacheCreationTokens int64) (pricing.Costs, bool) { if recordID := lookupKV(md, middleware.KeyLLMResolvedProviderID); recordID != "" { - if entry, ok := m.perRecord[recordID][model]; ok { + if entry, ok := perRecordEntry(m.perRecord[recordID], model); ok { return pricing.EntryCosts(entry, surface, inTokens, outTokens, cachedTokens, cacheCreationTokens), true } } return m.defaults.Costs(surface, model, inTokens, outTokens, cachedTokens, cacheCreationTokens) } +// perRecordEntry resolves the operator's stored price for a model on one +// provider record, falling back to the undated form of a dated Anthropic id +// so a client that pins a release date still bills at the registered rate. +func perRecordEntry(byModel map[string]pricing.Entry, model string) (pricing.Entry, bool) { + if entry, ok := byModel[model]; ok { + return entry, true + } + undated := llm.NormalizeAnthropicModel(model) + if undated == model { + return pricing.Entry{}, false + } + entry, ok := byModel[undated] + return entry, ok +} + // usd renders a cost as the fixed-precision string every cost.usd_* key // carries, so the per-bucket values and the aggregates round identically. // diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go index 1863aff20..d2b14f265 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware.go @@ -84,8 +84,10 @@ func (m *Middleware) MutationsSupported() bool { return false } func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { model, modelPresent := lookupMetadata(in.Metadata, middleware.KeyLLMModel) providerID, _ := lookupMetadata(in.Metadata, middleware.KeyLLMResolvedProviderID) + surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider) + nonInference, _ := lookupMetadata(in.Metadata, middleware.KeyLLMNonInference) - if denial := m.evaluateAllowlist(providerID, model, modelPresent); denial != nil { + if denial := m.evaluateAllowlist(providerID, surface, model, modelPresent, nonInference == "true"); denial != nil { return denial, nil } @@ -114,7 +116,7 @@ func (m *Middleware) Close() error { return nil } // evaluateAllowlist denies when the resolved provider's allowlist rejects the // model; nil means proceed. Scoped to the provider llm_router resolved, so an // unrestricted provider (absent from config) is never caught by another's list. -func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bool) *middleware.Output { +func (m *Middleware) evaluateAllowlist(providerID, surface, model string, modelPresent, nonInference bool) *middleware.Output { if len(m.cfg.ProviderAllowlists) == 0 { return nil } @@ -122,7 +124,7 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo // if this request targets a restricted provider — fail closed. llm_router // normally stamps the provider first, so this is a defensive guard. if providerID == "" { - return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) + return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) } allowlist, restricted := m.cfg.ProviderAllowlists[providerID] if !restricted { @@ -133,18 +135,29 @@ func (m *Middleware) evaluateAllowlist(providerID, model string, modelPresent bo // Fail closed: with an allowlist in effect for this provider, a request whose // model the parser couldn't extract (absent/empty) is denied. This enforces // the allowlist for path-routed providers (Bedrock, Vertex) with no body model. + // + // The exception is a non-inference endpoint the router already authorised. + // The model listing and the connection-warming probe name no model + // anywhere — not in a body, not in the path — so failing closed here + // rejected model discovery for exactly the accounts that configured an + // allowlist, which is the outage this endpoint is meant to avoid. The + // per-model lookup does name one (the router stamps it from the path), so + // it still falls through to the allowlist check below. if !modelPresent || normaliseModel(model) == "" { - return denyModel("", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) + if nonInference { + return nil + } + return denyModel(surface, "", denyCodeModelUnknown, denyMessageModelUnknown, denyReasonModelUnknown) } if modelInAllowlist(allowlist, model) { return nil } - return denyModel(model, denyCodeModel, denyMessageModel, denyReasonModel) + return denyModel(surface, model, denyCodeModel, denyMessageModel, denyReasonModel) } // denyModel builds a 403 deny Output for a model-allowlist rejection. model is // included in the details only when non-empty. -func denyModel(model, code, message, reason string) *middleware.Output { +func denyModel(surface, model, code, message, reason string) *middleware.Output { details := map[string]string{} if model != "" { details["model"] = model @@ -156,6 +169,7 @@ func denyModel(model, code, message, reason string) *middleware.Output { Code: code, Message: message, Details: details, + Surface: surface, }, Metadata: []middleware.KV{ {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, diff --git a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go index 5f35fefd3..19d8473fe 100644 --- a/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_guardrail/middleware_test.go @@ -343,3 +343,52 @@ func TestFactoryNormalisesAllowlist(t *testing.T) { require.NoError(t, err) assert.Equal(t, middleware.DecisionAllow, out2.Decision, "trimmed entry must still match") } + +// TestAllowlistSkipsNonInferenceWithoutModel covers the reported regression: +// GET /v1/models carries no model anywhere, so the fail-closed rule above +// denied model discovery for exactly the accounts that configured a provider +// allowlist — the clients that read a 403 here render an empty model picker. +// The router authorises those endpoints by path before the guardrail sees +// them, so an absent model there is expected rather than undeterminable. +func TestAllowlistSkipsNonInferenceWithoutModel(t *testing.T) { + mw := New(providerCfg("gpt-4o")) + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"}, + )) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "model discovery must not be refused because it names no model") +} + +// TestAllowlistStillAppliesToNonInferenceWithModel pins that the exemption is +// scoped to requests that genuinely name nothing. The per-model lookup +// (GET /v1/models/{id}) is non-inference too, but the router stamps the model +// from its path, so the allowlist must still decide it — otherwise the +// exemption becomes a way to confirm a model the policy blocks. +func TestAllowlistStillAppliesToNonInferenceWithModel(t *testing.T) { + mw := New(providerCfg("gpt-4o")) + + t.Run("model in the allowlist", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"}, + middleware.KV{Key: middleware.KeyLLMModel, Value: "gpt-4o"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "an allowlisted model must stay reachable") + }) + + t.Run("model outside the allowlist", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), newInputProvider(testProvider, + middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"}, + middleware.KV{Key: middleware.KeyLLMModel, Value: "claude-opus-5"}, + )) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "non-inference must not become a way past the allowlist") + require.NotNil(t, out.DenyReason) + assert.Equal(t, "llm_policy.model_blocked", out.DenyReason.Code, + "a named but blocked model is blocked, not unknown") + }) +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go index 722588a15..60b99e194 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go @@ -217,6 +217,32 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut return mutations } +// bodyInjectableSurfaces are the request-body dialects that accept the +// OpenAI-standard identity fields this middleware writes. A surface +// outside this set gets header-only stamping: "user" and "metadata.tags" +// are not part of the Anthropic Messages schema, which rejects unknown +// top-level fields and permits only "user_id" under metadata, so writing +// them into an Anthropic-shaped body turns a working request into a 400. +// Claude Code speaks that shape through gateway records pinned to the +// OpenAI parser, so the check keys on the detected surface rather than +// on the provider record. +var bodyInjectableSurfaces = map[string]struct{}{ + "openai": {}, + // An empty surface means no parser claimed the path (a custom gateway + // base). Those upstreams are OpenAI-compatible by convention, so keep + // the long-standing behaviour rather than silently dropping identity. + "": {}, +} + +// bodyAcceptsOpenAIIdentity reports whether the request body may carry the +// OpenAI-standard identity fields, read from the surface llm_request_parser +// resolved from the request path. +func bodyAcceptsOpenAIIdentity(in *middleware.Input) bool { + surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider) + _, ok := bodyInjectableSurfaces[surface] + return ok +} + // injectIntoBody parses the request body and writes the supplied // identity dimensions into it. Tags land at metadata.tags (creating // the metadata object when absent); the user identity lands at the @@ -225,6 +251,8 @@ func applyHeaderPair(rule *HeaderPairRule, in *middleware.Input) *middleware.Mut // was written. Returns ok=false (no mutation) when: // // - both inputs are empty (nothing to write); +// - the body speaks a dialect without these fields (see +// bodyInjectableSurfaces); // - the body is empty or truncated (we don't have the full document // to safely round-trip); // - the body isn't a JSON object (skip silently — this middleware @@ -245,6 +273,9 @@ func injectIntoBody(in *middleware.Input, tags []string, userID string) ([]byte, if in == nil || len(in.Body) == 0 || in.BodyTruncated { return nil, false } + if !bodyAcceptsOpenAIIdentity(in) { + return nil, false + } var doc map[string]any if err := json.Unmarshal(in.Body, &doc); err != nil { return nil, false diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go index 8ec0930b5..f602f5c33 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go @@ -704,3 +704,57 @@ func TestInject_ExtraHeaders_EmptyValueSkipped(t *testing.T) { "empty extra value must not be stamped") } } + +// TestInject_AnthropicBodyIsNotRewritten pins the shape gate. Claude Code +// reaches a LiteLLM record on /v1/messages, where "user" is not a +// permitted top-level field and metadata accepts only "user_id", so +// writing the OpenAI-standard fields would turn a working request into a +// 400 naming a field the client never sent. Header stamping still runs, so +// spend tracking and per-end-user budgets keep working. +func TestInject_AnthropicBodyIsNotRewritten(t *testing.T) { + rule := liteLLMRuleWithBody() + rule.HeaderPair.EndUserIDInBody = true + mw := New(Config{Providers: []ProviderInjection{rule}}) + + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.UserEmail = "alice@example.com" + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"}) + in.Body = []byte(`{"model":"claude-sonnet-5","messages":[]}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + assert.Empty(t, out.Mutations.BodyReplace, + "an Anthropic-shaped body must reach the upstream unmodified") + + var endUser string + for _, kv := range out.Mutations.HeadersAdd { + if kv.Key == "x-litellm-end-user-id" { + endUser = kv.Value + } + } + assert.Equal(t, "alice@example.com", endUser, + "header stamping must still carry identity when body inject is skipped") +} + +// TestInject_OpenAIBodyStillRewritten guards the gate against +// over-reaching: the OpenAI surface must keep its body-level identity, +// which is the only path LiteLLM's tag-budget check reads. +func TestInject_OpenAIBodyStillRewritten(t *testing.T) { + mw := New(Config{Providers: []ProviderInjection{liteLLMRuleWithBody()}}) + + in := newInput(litellmProvider, "alice", []string{"grp-eng"}) + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "openai"}) + in.Body = []byte(`{"model":"gpt-4o-mini","messages":[]}`) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.NotEmpty(t, out.Mutations.BodyReplace, "the OpenAI surface still gets body tags") + + var doc map[string]any + require.NoError(t, json.Unmarshal(out.Mutations.BodyReplace, &doc)) + meta, ok := doc["metadata"].(map[string]any) + require.True(t, ok, "metadata must be an object") + assert.NotEmpty(t, meta["tags"], "metadata.tags must still be written") +} diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go index 42ac56b9b..1e7edcf42 100644 --- a/proxy/internal/middleware/builtin/llm_limit_check/middleware.go +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware.go @@ -84,6 +84,15 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew return allowNoAttribution(), nil } + // Model-listing and other non-inference endpoints carry no model, and + // management's per-model allowlist fails closed on an empty one. The + // router has already authorised the route against the caller's groups + // and the request consumes no tokens, so gating it on a model that + // cannot exist would only break gateway model discovery. + if lookupKV(in.Metadata, middleware.KeyLLMNonInference) == "true" { + return allowNoAttribution(), nil + } + providerID := lookupKV(in.Metadata, middleware.KeyLLMResolvedProviderID) if providerID == "" { // llm_router didn't emit a resolved provider id — usually @@ -117,7 +126,7 @@ func (m *Middleware) Invoke(ctx context.Context, in *middleware.Input) (*middlew } if resp.GetDecision() == "deny" { - return denyFromManagement(resp), nil + return denyFromManagement(resp, lookupKV(in.Metadata, middleware.KeyLLMProvider)), nil } return allowFromManagement(resp), nil } @@ -161,7 +170,7 @@ func allowFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.O // envelope. The deny code surfaces verbatim through the framework's // fixed JSON template; arbitrary middleware bytes can't reach the // wire. -func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Output { +func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse, surface string) *middleware.Output { code := resp.GetDenyCode() if code == "" { code = "llm_policy.cap_exceeded" @@ -176,6 +185,7 @@ func denyFromManagement(resp *proto.CheckLLMPolicyLimitsResponse) *middleware.Ou DenyReason: &middleware.DenyReason{ Code: code, Message: denyMessageForCode(code), + Surface: surface, }, Metadata: []middleware.KV{ {Key: middleware.KeyLLMPolicyDecision, Value: "deny"}, diff --git a/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go index 87aa8e9e9..7754998ee 100644 --- a/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_limit_check/middleware_test.go @@ -224,3 +224,35 @@ func TestMetadataKeys_Allowlist(t *testing.T) { } assert.ElementsMatch(t, want, keys) } + +// TestInvoke_NonInferenceSkipsPreflight covers gateway model discovery: +// GET /v1/models carries no model, and management's per-model allowlist +// fails closed on an empty one, so a pre-flight would deny discovery for +// exactly the accounts that use the model allowlist. The router marks the +// request non-inference after authorising the route, and the gate must +// then allow without calling management at all. +func TestInvoke_NonInferenceSkipsPreflight(t *testing.T) { + mgmt := &fakeMgmt{ + checkResp: &proto.CheckLLMPolicyLimitsResponse{ + Decision: "deny", + DenyCode: "llm_policy.model_blocked", + }, + } + m := New(mgmt, nil) + + out := runInvoke(t, m, &middleware.Input{ + AccountID: "acc-1", + UserID: "user-bob", + UserGroups: []string{"grp-engineers"}, + Metadata: []middleware.KV{ + {Key: middleware.KeyLLMResolvedProviderID, Value: "prov-1"}, + {Key: middleware.KeyLLMNonInference, Value: "true"}, + }, + }) + + assert.Equal(t, middleware.DecisionAllow, out.Decision, "model-less endpoints must not be gated on a model") + assert.Nil(t, mgmt.checkReq, "no pre-flight may be sent for a non-inference request") + + assert.Empty(t, lookupKV(out.Metadata, middleware.KeyLLMSelectedPolicyID), + "no policy is attributed when nothing was metered") +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go index d8cd81437..82f44cb50 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/bedrock_test.go @@ -1,9 +1,13 @@ package llm_request_parser import ( + "context" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" ) func TestParseBedrockPath(t *testing.T) { @@ -36,3 +40,25 @@ func TestParseBedrockPath(t *testing.T) { } } } + +// TestInvoke_BedrockCountTokens covers the dedicated token-counting +// endpoint. Denying it does not break the client, it just pushes context +// counting back onto the inference endpoint, which is billable. +func TestInvoke_BedrockCountTokens(t *testing.T) { + mw := newMiddleware(t) + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens", + Body: []byte(`{"input":{"converse":{"messages":[]}}}`), + }) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + + model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel) + require.True(t, ok, "count-tokens carries a model in the path and must emit it") + assert.Equal(t, "anthropic.claude-sonnet-4-5", model, "model must be normalized like any other action") + + stream, _ := metaValue(t, out.Metadata, middleware.KeyLLMStream) + assert.Equal(t, "false", stream, "count-tokens never streams") +} diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go index b4d1e16d4..7129c2298 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware.go @@ -61,6 +61,8 @@ func (middlewareImpl) MetadataKeys() []string { middleware.KeyLLMRequestPromptRaw, middleware.KeyLLMCaptureTruncated, middleware.KeyLLMSessionID, + middleware.KeyLLMAgentID, + middleware.KeyLLMParentAgentID, } } @@ -72,9 +74,9 @@ func (middlewareImpl) Close() error { return nil } // Invoke detects the LLM provider, parses request facts, and emits // metadata. Always returns DecisionAllow; never errors. Provider -// selection prefers the configured providerID (synthesiser-stamped on -// agent-network targets) so requests routed to a custom upstream URL -// still resolve. Falls back to URL sniffing when no providerID is set. +// selection prefers the request path, falling back to the configured +// providerID (synthesiser-stamped on agent-network targets) so requests +// routed to a custom upstream URL still resolve. func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { out := &middleware.Output{Decision: middleware.DecisionAllow} if in == nil { @@ -92,9 +94,14 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle return m.invokeBedrock(in, br), nil } - parser, ok := llm.ParserByName(m.providerID) + // A path that names an API surface wins over the configured providerID: + // a gateway record pinned to "openai" still serves Claude Code on + // /v1/messages, and reading that body with the OpenAI parser loses the + // Anthropic usage block and prices the request on the wrong surface. + // providerID stays the fallback for upstreams whose path says nothing. + parser, ok := llm.DetectParser(extractPath(in.URL)) if !ok { - parser, ok = llm.DetectParser(extractPath(in.URL)) + parser, ok = llm.ParserByName(m.providerID) } if !ok { return out, nil @@ -116,9 +123,9 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle } appendSessionID := func(md []middleware.KV) []middleware.KV { if sessionID != "" { - return append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) + md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } - return md + return appendAgentIDs(md, in.Headers) } facts, err := parser.ParseRequest(in.Body) @@ -160,6 +167,41 @@ func (m middlewareImpl) Invoke(_ context.Context, in *middleware.Input) (*middle return out, nil } +// agentIDHeader and parentAgentIDHeader carry sub-agent attribution: a +// coding agent that spawns helpers stamps the spawned agent's id, plus the +// spawning agent's when that helper is itself nested. Both are opaque +// identifiers rather than content, so they're emitted regardless of the +// prompt-collection toggle, the same way the session id is. +const ( + agentIDHeader = "x-claude-code-agent-id" + parentAgentIDHeader = "x-claude-code-parent-agent-id" +) + +// appendAgentIDs stamps the sub-agent attribution headers onto the metadata +// bag, skipping either one the request doesn't carry. +func appendAgentIDs(md []middleware.KV, headers []middleware.KV) []middleware.KV { + for _, pair := range []struct{ key, header string }{ + {middleware.KeyLLMAgentID, agentIDHeader}, + {middleware.KeyLLMParentAgentID, parentAgentIDHeader}, + } { + if v := headerValue(headers, pair.header); v != "" { + md = append(md, middleware.KV{Key: pair.key, Value: v}) + } + } + return md +} + +// headerValue returns the first non-empty value for the named header. +// Headers arrive in canonical form, so the match is case-insensitive. +func headerValue(headers []middleware.KV, want string) string { + for _, kv := range headers { + if strings.EqualFold(kv.Key, want) && kv.Value != "" { + return kv.Value + } + } + return "" +} + // sessionIDHeaders are request header names that may carry a client // session identifier, checked in order, case-insensitively. Matching is // against Go's canonical header form, so use the hyphenated names the @@ -173,10 +215,8 @@ var sessionIDHeaders = []string{"x-claude-code-session-id", "session-id", "x-ses // canonical form, so the match is case-insensitive. func sessionIDFromHeaders(headers []middleware.KV) string { for _, want := range sessionIDHeaders { - for _, kv := range headers { - if strings.EqualFold(kv.Key, want) && kv.Value != "" { - return kv.Value - } + if v := headerValue(headers, want); v != "" { + return v } } return "" @@ -252,6 +292,12 @@ func parseVertexPath(reqPath string) (vertexRequest, bool) { if c := strings.LastIndex(rest, ":"); c >= 0 { model, action = rest[:c], rest[c+1:] } + // Token counting hangs off the model as its own path segment + // (".../models/{model}/count-tokens:rawPredict"), so anything past the + // first "/" belongs to the method rather than the model id. + if slash := strings.Index(model, "/"); slash >= 0 { + model = model[:slash] + } model = llm.NormalizeVertexModel(model) if model == "" { return vertexRequest{}, false @@ -298,6 +344,7 @@ func (m middlewareImpl) invokeVertex(in *middleware.Input, vx vertexRequest) *mi if sessionID != "" { md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } + md = appendAgentIDs(md, in.Headers) promptTruncated := false if parser != nil && m.capturePrompt { @@ -345,7 +392,9 @@ func trimBedrockNamespace(reqPath string) string { // // /model/{modelId}/{action} // -// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream}. +// action ∈ {invoke, invoke-with-response-stream, converse, converse-stream, +// count-tokens}. Token counting carries a model and no usage, so it routes +// like any other action and meters to zero. // The modelId may be URL-encoded and may carry a cross-region inference-profile // prefix and a version suffix; normalizeBedrockModel strips both so the model // matches catalog pricing. @@ -369,7 +418,7 @@ func parseBedrockPath(reqPath string) (bedrockRequest, bool) { return bedrockRequest{}, false } switch action { - case "invoke", "converse": + case "invoke", "converse", "count-tokens": return bedrockRequest{model: model}, true case "invoke-with-response-stream", "converse-stream": return bedrockRequest{model: model, stream: true}, true @@ -397,6 +446,7 @@ func (m middlewareImpl) invokeBedrock(in *middleware.Input, br bedrockRequest) * if sessionID != "" { md = append(md, middleware.KV{Key: middleware.KeyLLMSessionID, Value: sessionID}) } + md = appendAgentIDs(md, in.Headers) promptTruncated := false if parser != nil && m.capturePrompt { diff --git a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go index bc185b295..8d8517860 100644 --- a/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_request_parser/middleware_test.go @@ -45,6 +45,8 @@ func TestMiddleware_StaticSurface(t *testing.T) { middleware.KeyLLMRequestPromptRaw, middleware.KeyLLMCaptureTruncated, middleware.KeyLLMSessionID, + middleware.KeyLLMAgentID, + middleware.KeyLLMParentAgentID, } assert.Equal(t, expected, keys, "metadata key allowlist must match the spec") } @@ -230,6 +232,31 @@ func TestInvoke_ProviderIDConfigBypassesURLSniff(t *testing.T) { assert.Equal(t, "gpt-4o-mini", model) } +func TestInvoke_PathSurfaceBeatsProviderIDConfig(t *testing.T) { + // Gateway records (LiteLLM, Portkey, OpenRouter) pin provider_id + // "openai", but the same record serves Claude Code on /v1/messages. + // Parsing that body as OpenAI reads no usage off the Anthropic + // response and prices the request on a surface where no claude-* + // model exists, so the path has to win. + mw, err := Factory{}.New([]byte(`{"provider_id":"openai"}`)) + require.NoError(t, err, "factory must accept provider_id config") + + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","stream":true,"messages":[{"role":"user","content":"Hi"}]}`), + }) + require.NoError(t, err) + require.NotNil(t, out) + + provider, ok := metaValue(t, out.Metadata, middleware.KeyLLMProvider) + require.True(t, ok, "provider must be emitted") + assert.Equal(t, "anthropic", provider, "the /v1/messages path selects the Anthropic surface") + + model, ok := metaValue(t, out.Metadata, middleware.KeyLLMModel) + require.True(t, ok, "model must be extracted") + assert.Equal(t, "claude-sonnet-5", model) +} + func TestInvoke_UnknownProviderIDFallsBackToURL(t *testing.T) { mw, err := Factory{}.New([]byte(`{"provider_id":"not-a-real-parser"}`)) require.NoError(t, err, "factory must accept any provider_id string") @@ -416,3 +443,81 @@ func TestInvoke_NilInputAllows(t *testing.T) { assert.Equal(t, middleware.DecisionAllow, out.Decision, "nil input still allows") assert.Empty(t, out.Metadata, "nil input emits no metadata") } + +// TestParseVertexPath_CountTokensKeepsModel covers Vertex token counting, +// where the method hangs off the model as its own path segment. Splitting +// only on the final colon swallowed "/count-tokens" into the model id, so +// the router saw a model no route could claim. +func TestParseVertexPath_CountTokensKeepsModel(t *testing.T) { + cases := map[string]struct { + model string + stream bool + }{ + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:rawPredict": {model: "claude-sonnet-5"}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5:streamRawPredict": {model: "claude-sonnet-5", stream: true}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5/count-tokens:rawPredict": {model: "claude-sonnet-5"}, + "/v1/projects/p/locations/global/publishers/anthropic/models/claude-sonnet-5@20250929/count-tokens:rawPredict": {model: "claude-sonnet-5"}, + } + for path, want := range cases { + vx, ok := parseVertexPath(path) + require.True(t, ok, "must parse %q", path) + assert.Equal(t, want.model, vx.model, "model for %q", path) + assert.Equal(t, want.stream, vx.stream, "stream flag for %q", path) + assert.Equal(t, "anthropic", vx.publisher, "publisher for %q", path) + } +} + +// TestInvoke_EmitsAgentIDs covers sub-agent attribution: several agents run +// in parallel inside one session, and without their ids every request in +// the session attributes to the session alone. +func TestInvoke_EmitsAgentIDs(t *testing.T) { + mw := newMiddleware(t) + + t.Run("spawned agent", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + Headers: []middleware.KV{ + {Key: "X-Claude-Code-Session-Id", Value: "sess-1"}, + {Key: "X-Claude-Code-Agent-Id", Value: "agent-7"}, + }, + }) + require.NoError(t, err) + + agent, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + require.True(t, ok, "the spawned agent's id must be emitted") + assert.Equal(t, "agent-7", agent) + + _, ok = metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID) + assert.False(t, ok, "a top-level agent has no parent to emit") + }) + + t.Run("nested agent", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + Headers: []middleware.KV{ + {Key: "X-Claude-Code-Agent-Id", Value: "agent-9"}, + {Key: "X-Claude-Code-Parent-Agent-Id", Value: "agent-7"}, + }, + }) + require.NoError(t, err) + + agent, _ := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + assert.Equal(t, "agent-9", agent) + parent, ok := metaValue(t, out.Metadata, middleware.KeyLLMParentAgentID) + require.True(t, ok, "a nested agent must carry the spawning agent's id") + assert.Equal(t, "agent-7", parent) + }) + + t.Run("absent on a plain request", func(t *testing.T) { + out, err := mw.Invoke(context.Background(), &middleware.Input{ + URL: "/v1/messages", + Body: []byte(`{"model":"claude-sonnet-5","messages":[]}`), + }) + require.NoError(t, err) + + _, ok := metaValue(t, out.Metadata, middleware.KeyLLMAgentID) + assert.False(t, ok, "no key is emitted when the client sends no agent id") + }) +} diff --git a/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go b/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go index 40cbcb6bd..badd358c5 100644 --- a/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go +++ b/proxy/internal/middleware/builtin/llm_router/bedrock_route_test.go @@ -1,9 +1,13 @@ package llm_router import ( + "context" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" ) // TestRouteClaimsModel_BedrockNormalizesCandidate guards the fix for the native @@ -28,3 +32,86 @@ func TestRouteClaimsModel_BedrockNormalizesCandidate(t *testing.T) { assert.False(t, routeClaimsModel(openai, "us.gpt-4o"), "non-Bedrock routes must not strip a us. prefix") } + +// TestRouter_BedrockCountTokensRoutes pins that the token-counting action +// reaches the Bedrock route instead of denying as not-routable. +func TestRouter_BedrockCountTokensRoutes(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "bedrock-prod", + Bedrock: true, + Models: []string{"anthropic.claude-sonnet-4-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + }}}) + + in := newInputWithModelAndURL("anthropic.claude-sonnet-4-5", + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/count-tokens") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "bedrock"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "count-tokens must route, not deny") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host) +} + +// TestRouter_BedrockInferenceProfilesRoutes covers the startup lookups a +// client makes to resolve a configured inference profile. They carry no +// model, so before they were recognised they denied and wrote a policy +// rejection into the access log on every session start. +func TestRouter_BedrockInferenceProfilesRoutes(t *testing.T) { + bedrock := ProviderRoute{ + ID: "bedrock-prod", + Bedrock: true, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + } + openai := ProviderRoute{ + ID: "openai-prod", + Models: []string{"gpt-4o"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.openai.com", + } + mw := New(Config{Providers: []ProviderRoute{openai, bedrock}}) + + for _, path := range []string{ + "/inference-profiles?type=SYSTEM_DEFINED", + "/inference-profiles/us.anthropic.claude-sonnet-5", + } { + out, err := mw.Invoke(context.Background(), newModellessInput(path)) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "%s must route", path) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host, + "%s must reach the Bedrock provider, not the first authorised one", path) + + nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + assert.Equal(t, "true", nonInference, "%s carries no model to gate on", path) + } +} + +// TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix pins that the +// optional gateway namespace is removed before the request goes upstream. +func TestRouter_BedrockNamespacedInferenceProfilesStripsPrefix(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "bedrock-prod", + Bedrock: true, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + }}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/bedrock/inference-profiles")) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "/bedrock", out.Mutations.RewriteUpstream.StripPathPrefix, + "the namespace prefix must not reach the real Bedrock endpoint") +} diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 2d987eef6..e6ad332fc 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -109,6 +109,10 @@ func (m *Middleware) MetadataKeys() []string { middleware.KeyLLMAuthorisingGroups, middleware.KeyLLMPolicyDecision, middleware.KeyLLMPolicyReason, + middleware.KeyLLMNonInference, + // Emitted only for the per-model lookup, whose model lives in the path + // rather than a body the parser could read. + middleware.KeyLLMModel, } } @@ -137,29 +141,26 @@ const ( // known to a provider that no policy authorises for the caller deny // with no_authorised_provider. func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middleware.Output, error) { + reqPath := requestPath(in.URL) + // The caller's API dialect, used to mirror a denial in the vendor's own + // error shape so the client can explain it to the user. + surface, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider) + model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel) + // Vertex AI carries the model in the URL path, not the body, and is // selected by path rather than by the model/vendor table. Route it before // the model lookup so a model the parser extracted from the path can't be // claimed by a same-vendor direct provider (e.g. claude-* on api.anthropic.com). - reqPath := requestPath(in.URL) if isVertexPath(reqPath) { - model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel) // The request parser emits no llm.provider for a Vertex publisher it // can't parse (e.g. google/gemini). Forwarding such a request would // bypass token/budget metering, so deny it rather than serve it // unmetered. - if vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider); vendor == "" { - return denyUnmeterable(), nil + if surface == "" { + return denyUnmeterable(surface), nil } route, outcome := m.matchVertex(reqPath, model, in.UserGroups) - switch outcome { - case matchOutcomeFound: - return m.allowWithRoute(route, in.UserGroups), nil - case matchOutcomeUnauthorised: - return denyNoAuthorisedRoute(model), nil - default: - return denyUnknownModel(model), nil - } + return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil } // Bedrock likewise carries the model in the URL path (/model/{id}/{action}), @@ -167,52 +168,120 @@ func (m *Middleware) Invoke(_ context.Context, in *middleware.Input) (*middlewar // before the model lookup; when the prefix is present, strip it from the // forwarded path so the real Bedrock endpoint receives its native path. if isBedrockPath(reqPath) { - model, _ := lookupMetadata(in.Metadata, middleware.KeyLLMModel) native, hadPrefix := splitBedrockNamespace(reqPath) route, outcome := m.matchBedrock(native, model, in.UserGroups) - switch outcome { - case matchOutcomeFound: - out := m.allowWithRoute(route, in.UserGroups) - if hadPrefix && out.Mutations != nil && out.Mutations.RewriteUpstream != nil { - out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix + return m.decide(route, outcome, surface, model, in.UserGroups, func(out *middleware.Output) { + if hadPrefix { + stripBedrockNamespace(out) } - return out, nil - case matchOutcomeUnauthorised: - return denyNoAuthorisedRoute(model), nil - default: - return denyUnknownModel(model), nil - } + }), nil } - model, ok := lookupMetadata(in.Metadata, middleware.KeyLLMModel) - if !ok || model == "" { - // Non-inference endpoints (model listing) carry no model but still - // need rewriting from the synth placeholder to a real upstream; - // clients such as Codex call GET /v1/models at startup to enumerate - // availability and read a 403 as "model unavailable". - route, outcome := m.matchModelless(requestPath(in.URL), in.UserGroups) - switch outcome { - case matchOutcomeFound: - return m.allowWithRoute(route, in.UserGroups), nil - case matchOutcomeUnauthorised: - // A recognised model-less endpoint exists but no provider - // authorises the caller — deny as an authorisation failure - // rather than masking it as a missing model. - return denyNoAuthorisedRoute(model), nil - default: - return denyMissingModel(), nil - } + // GET /v1/models/{id} carries no body, so no model reaches the router in + // metadata — but the path names one, and answering it confirms a model + // exists and is reachable. Authorise it against the model table like any + // other per-model request, then mark it non-inference so it still skips + // the token pre-flight it would otherwise charge nothing against. + if detail, isDetail := modelDetailID(reqPath); isDetail && isNonInferenceMethod(in.Method) { + route, outcome := m.matchRoute(detail, surface, reqPath, in.UserGroups) + return m.decide(route, outcome, surface, detail, in.UserGroups, func(out *middleware.Output) { + markNonInference(out) + // The parser reads models from JSON bodies only, and this request + // has none, so stamp the one the path names. Without it the + // guardrail's own allowlist — a separate, possibly narrower list + // than the route's — never sees a model to check. + out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMModel, Value: detail}) + }), nil } - vendor, _ := lookupMetadata(in.Metadata, middleware.KeyLLMProvider) - route, outcome := m.matchRoute(model, vendor, requestPath(in.URL), in.UserGroups) + if model == "" { + return m.routeModelless(reqPath, surface, in.Method, in.UserGroups), nil + } + + route, outcome := m.matchRoute(model, surface, reqPath, in.UserGroups) + return m.decide(route, outcome, surface, model, in.UserGroups, nil), nil +} + +// decide turns a per-model match result into the middleware's decision. Every +// surface that routes by model shares the same two denial arms — a model no +// route claims is not routable, one that some route claims but none authorises +// for this caller is an authorisation failure — so they live here once. +// decorate, when non-nil, adjusts the allow with whatever that surface needs. +func (m *Middleware) decide( + route ProviderRoute, + outcome matchOutcome, + surface, model string, + userGroups []string, + decorate func(*middleware.Output), +) *middleware.Output { switch outcome { case matchOutcomeFound: - return m.allowWithRoute(route, in.UserGroups), nil + out := m.allowWithRoute(route, surface, userGroups) + if decorate != nil { + decorate(out) + } + return out case matchOutcomeUnauthorised: - return denyNoAuthorisedRoute(model), nil + return denyNoAuthorisedRoute(surface, model) default: - return denyUnknownModel(model), nil + return denyUnknownModel(surface, model) + } +} + +// routeModelless serves the endpoints that name no model at all: the model +// listing, the connection-warming probe, and the Bedrock inference-profile +// lookup. They still need rewriting from the synth placeholder to a real +// upstream — clients such as Codex call GET /v1/models at startup to enumerate +// availability and read a 403 as "model unavailable". +func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups []string) *middleware.Output { + route, outcome := m.matchModelless(reqPath, method, userGroups) + switch outcome { + case matchOutcomeFound: + out := m.allowWithRoute(route, surface, userGroups) + markNonInference(out) + 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...) + } + return out + case matchOutcomeUnauthorised: + // A recognised model-less endpoint exists but no provider authorises + // the caller — deny as an authorisation failure rather than masking it + // as a missing model. + return denyNoAuthorisedRoute(surface, "") + default: + return denyMissingModel(surface) + } +} + +// isNonInferenceMethod reports whether a request method is one the +// non-inference endpoints actually use: the listing and the per-model lookup +// are GET, the connection-warming probe is HEAD or GET. The method is the only +// thing separating "GET /v1/models/{id}" from a POST to the same path carrying +// an inference body, and the non-inference mark exempts a request from the +// token pre-flight — so anything else falls through to normal per-model +// routing, which denies when the request names no model. +func isNonInferenceMethod(method string) bool { + return method == http.MethodGet || method == http.MethodHead +} + +// 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) { + out.Metadata = append(out.Metadata, middleware.KV{Key: middleware.KeyLLMNonInference, Value: "true"}) +} + +// stripBedrockNamespace tells the rewrite to drop the optional "/bedrock" +// gateway namespace so the upstream receives its native Bedrock path. +func stripBedrockNamespace(out *middleware.Output) { + if out.Mutations != nil && out.Mutations.RewriteUpstream != nil { + out.Mutations.RewriteUpstream.StripPathPrefix = bedrockNamespacePrefix } } @@ -300,12 +369,60 @@ func (m *Middleware) matchRoute(model, vendor, reqPath string, userGroups []stri return best, matchOutcomeFound } -// isModelLessPath reports whether reqPath is a known OpenAI-shaped -// non-inference endpoint that legitimately carries no model in its -// request (the model-listing endpoints). These must route to an upstream -// rather than deny, so model enumeration works end to end. +// connectionWarmPath is the probe Anthropic clients send before their first +// inference request to open the upstream connection early. Forwarding it +// warms the connection the request will actually use; denying it only fills +// the access log with rejections at every session start. +const connectionWarmPath = "/api/hello" + +// modelListingPath is the endpoint clients read at startup to populate +// their model picker. Its response is a list the proxy can bound; the +// per-model "/v1/models/{id}" lookup returns a single object and is left +// alone. +const modelListingPath = "/v1/models" + +// isModelLessPath reports whether reqPath is a known non-inference endpoint +// that legitimately carries no model at all: the model listing and the +// connection-warming probe. These must route to an upstream rather than +// deny, so model enumeration works end to end. The per-model +// "/v1/models/{id}" lookup is deliberately excluded — it names a model, so +// it is authorised against the model table instead (see modelDetailID). func isModelLessPath(reqPath string) bool { - return reqPath == "/v1/models" || strings.HasPrefix(reqPath, "/v1/models/") + return reqPath == modelListingPath || reqPath == connectionWarmPath +} + +// modelDetailID returns the model id named by a "/v1/models/{id}" lookup. +// reqPath comes from url.URL.Path, which is already percent-decoded, so an +// id carrying a "/" (a self-hosted "Qwen/Qwen2.5-0.5B-Instruct" sent as +// "Qwen%2FQwen2.5-...") arrives whole and everything after the prefix is the +// id, separators included. +func modelDetailID(reqPath string) (string, bool) { + if !strings.HasPrefix(reqPath, modelListingPath+"/") { + return "", false + } + id := strings.TrimPrefix(reqPath, modelListingPath+"/") + if id == "" { + return "", false + } + return id, true +} + +// isBedrockModelLessPath reports whether reqPath is a Bedrock +// inference-profile lookup, optionally behind the "/bedrock" gateway +// namespace. Clients read these at startup to resolve a configured profile +// to its underlying model. They carry no model of their own, so they route +// by path to a Bedrock provider rather than through the model table. +// +// On native AWS these live on the control plane ("bedrock.") while a +// provider's upstream is normally the runtime host ("bedrock-runtime."), +// so forwarding yields a 404 there. That is deliberate: a client has one base +// URL, so pointing it straight at the runtime host 404s identically, and +// forwarding keeps the proxy transparent instead of inventing a policy denial +// the client would never otherwise see. Operators whose Bedrock upstream is a +// gateway that does serve the lookup get a working answer. +func isBedrockModelLessPath(reqPath string) bool { + native, _ := splitBedrockNamespace(reqPath) + return native == "/inference-profiles" || strings.HasPrefix(native, "/inference-profiles/") } // isVertexPath reports whether reqPath is a Google Vertex AI publisher @@ -332,20 +449,33 @@ func splitBedrockNamespace(reqPath string) (string, bool) { return reqPath, false } +// bedrockActions are the runtime actions that follow the model id in a +// Bedrock path. count-tokens is here so a client can price its context +// against the dedicated endpoint; denying it pushes that work back onto +// the inference endpoint, which bills for it. +var bedrockActions = []string{ + "/invoke", + "/invoke-with-response-stream", + "/converse", + "/converse-stream", + "/count-tokens", +} + // isBedrockPath reports whether reqPath is an AWS Bedrock runtime model -// endpoint: /model/{modelId}/{action} where action is invoke, -// invoke-with-response-stream, converse, or converse-stream — optionally behind -// a "/bedrock" gateway-namespace prefix. The model lives in the path, so these -// requests are routed by path to the Bedrock provider. +// endpoint: /model/{modelId}/{action} — optionally behind a "/bedrock" +// gateway-namespace prefix. The model lives in the path, so these requests +// are routed by path to the Bedrock provider. func isBedrockPath(reqPath string) bool { native, _ := splitBedrockNamespace(reqPath) if !strings.HasPrefix(native, "/model/") { return false } - return strings.HasSuffix(native, "/invoke") || - strings.HasSuffix(native, "/invoke-with-response-stream") || - strings.HasSuffix(native, "/converse") || - strings.HasSuffix(native, "/converse-stream") + for _, action := range bedrockActions { + if strings.HasSuffix(native, action) { + return true + } + } + return false } // matchVertex selects the Vertex provider authorised for the caller's groups @@ -425,19 +555,26 @@ func (m *Middleware) matchPathRoute(reqPath, model string, userGroups []string, // declaration order), matchOutcomeUnauthorised when no provider authorises // the caller, or matchOutcomeUnknownModel when the path isn't a recognised // model-less endpoint. -func (m *Middleware) matchModelless(reqPath string, userGroups []string) (ProviderRoute, matchOutcome) { - if !isModelLessPath(reqPath) { +func (m *Middleware) matchModelless(reqPath, method string, userGroups []string) (ProviderRoute, matchOutcome) { + if !isNonInferenceMethod(method) { return ProviderRoute{}, matchOutcomeUnknownModel } - var candidates []ProviderRoute - for _, route := range m.cfg.Providers { + var eligible func(ProviderRoute) bool + switch { + case isBedrockModelLessPath(reqPath): + eligible = func(r ProviderRoute) bool { return r.Bedrock } + case isModelLessPath(reqPath): // Vertex/Bedrock are path-routed and don't serve OpenAI-style // model-listing endpoints; including them here could rewrite a // GET /v1/models to an upstream that 404s it. - if route.Vertex || route.Bedrock { - continue - } - if routeAuthorisesGroups(route, userGroups) { + eligible = func(r ProviderRoute) bool { return !r.Vertex && !r.Bedrock } + default: + return ProviderRoute{}, matchOutcomeUnknownModel + } + + var candidates []ProviderRoute + for _, route := range m.cfg.Providers { + if eligible(route) && routeAuthorisesGroups(route, userGroups) { candidates = append(candidates, route) } } @@ -564,6 +701,16 @@ func routeClaimsModel(route ProviderRoute, model string) bool { if route.Bedrock && llm.NormalizeBedrockModel(candidate) == model { return true } + // A client may pin a dated Anthropic id ("claude-sonnet-4-5-20250929") + // where the operator registered the undated one. Only an undated + // registration absorbs a dated request: normalising both sides would + // let a route pinned to one dated release claim a different one, so an + // operator who deliberately pinned a build would silently serve + // another — and with several such routes, ordering would decide which. + if candidate == llm.NormalizeAnthropicModel(candidate) && + candidate == llm.NormalizeAnthropicModel(model) { + return true + } } return false } @@ -612,7 +759,7 @@ func requestPath(raw string) string { // provider id so identity-stamping middlewares (llm_identity_inject) // tag the request with ONLY the groups that authorised this specific // route — not every group the peer happens to be in. -func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *middleware.Output { +func (m *Middleware) allowWithRoute(route ProviderRoute, surface string, userGroups []string) *middleware.Output { rewrite := &middleware.UpstreamRewrite{ Scheme: route.UpstreamScheme, Host: route.UpstreamHost, @@ -634,7 +781,7 @@ func (m *Middleware) allowWithRoute(route ProviderRoute, userGroups []string) *m // request time (cached + auto-refreshed) instead of a static value. bearer, err := m.gcpBearer(route.GCPServiceAccountKeyB64) if err != nil { - return denyUpstreamAuth() + return denyUpstreamAuth(surface) } authValue = bearer } @@ -704,11 +851,12 @@ func (m *Middleware) gcpTokenSource(saKeyB64 string) (oauth2.TokenSource, error) // denyUpstreamAuth is returned when the router cannot obtain the upstream // credential (e.g. a malformed service-account key or an unreachable token // endpoint). It surfaces as a 502 — an upstream problem, not a policy denial. -func denyUpstreamAuth() *middleware.Output { +func denyUpstreamAuth(surface string) *middleware.Output { return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 502, DenyReason: &middleware.DenyReason{ + Surface: surface, Code: denyCodeUpstreamAuth, Message: "could not obtain upstream credential", }, @@ -722,11 +870,12 @@ func denyUpstreamAuth() *middleware.Output { // denyUnmeterable returns the deny envelope for a path-routed request whose // publisher has no parser surface, so its usage can't be metered. Serving it // would bypass token/budget caps, so it is rejected with a 403. -func denyUnmeterable() *middleware.Output { +func denyUnmeterable(surface string) *middleware.Output { return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 403, DenyReason: &middleware.DenyReason{ + Surface: surface, Code: denyCodeUnmeterable, Message: "request publisher is not supported for metering", }, @@ -739,11 +888,12 @@ func denyUnmeterable() *middleware.Output { // denyMissingModel returns the deny envelope for a request whose // envelope has no llm.model metadata. -func denyMissingModel() *middleware.Output { +func denyMissingModel(surface string) *middleware.Output { return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 403, DenyReason: &middleware.DenyReason{ + Surface: surface, Code: denyCodeNotRoutable, Message: "missing llm.model on request envelope", }, @@ -756,11 +906,12 @@ func denyMissingModel() *middleware.Output { // denyUnknownModel returns the deny envelope for a model that no // configured provider claims. -func denyUnknownModel(model string) *middleware.Output { +func denyUnknownModel(surface, model string) *middleware.Output { return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 403, DenyReason: &middleware.DenyReason{ + Surface: surface, Code: denyCodeNotRoutable, Message: fmt.Sprintf("no provider configured for model %s", model), Details: map[string]string{"model": model}, @@ -775,11 +926,12 @@ func denyUnknownModel(model string) *middleware.Output { // denyNoAuthorisedRoute returns the deny envelope for a model that one // or more providers claim, but where no policy authorises the caller's // groups for any of those providers. -func denyNoAuthorisedRoute(model string) *middleware.Output { +func denyNoAuthorisedRoute(surface, model string) *middleware.Output { return &middleware.Output{ Decision: middleware.DecisionDeny, DenyStatus: 403, DenyReason: &middleware.DenyReason{ + Surface: surface, Code: denyCodeNoAuthorisedRoute, Message: fmt.Sprintf("no policy authorises model %s for the caller's groups", model), Details: map[string]string{"model": model}, diff --git a/proxy/internal/middleware/builtin/llm_router/middleware_test.go b/proxy/internal/middleware/builtin/llm_router/middleware_test.go index 425c383c1..336cdb9fe 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware_test.go @@ -2,6 +2,7 @@ package llm_router import ( "context" + "net/http" "testing" "github.com/stretchr/testify/assert" @@ -60,6 +61,8 @@ func TestMiddlewareIdentity(t *testing.T) { []string{ middleware.KeyLLMResolvedProviderID, middleware.KeyLLMAuthorisingGroups, + middleware.KeyLLMNonInference, + middleware.KeyLLMModel, middleware.KeyLLMPolicyDecision, middleware.KeyLLMPolicyReason, }, @@ -171,8 +174,12 @@ func TestRouter_MissingModel(t *testing.T) { // from which a model could be parsed). UserGroups matches defaultTestGroup. func newModellessInput(reqURL string) *middleware.Input { return &middleware.Input{ - Slot: middleware.SlotOnRequest, - URL: reqURL, + Slot: middleware.SlotOnRequest, + URL: reqURL, + // The non-inference endpoints are read requests; the method is what + // separates them from an inference body posted to the same path, so + // state it rather than leaning on the zero value. + Method: http.MethodGet, UserGroups: []string{defaultTestGroup}, } } @@ -197,6 +204,12 @@ func TestRouter_ModelLessPath_RoutesToAuthorisedProvider(t *testing.T) { provider, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) assert.Equal(t, "openai-prod", provider, "resolved provider must be the authorised route") + + // The limits gate reads this to tell "no model applies here" from + // "the model could not be determined", which fails closed. + nonInference, ok := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + require.True(t, ok, "model-less allow must mark the request non-inference") + assert.Equal(t, "true", nonInference) } func TestRouter_ModelLessPath_MultiProviderDeclarationOrder(t *testing.T) { @@ -873,3 +886,262 @@ func TestRouter_EmptyModelsClaimsAnyModel(t *testing.T) { resolved, _ := metaValue(t, out.Metadata, middleware.KeyLLMResolvedProviderID) assert.Equal(t, "litellm", resolved) } + +// TestRouter_DatedAnthropicModelRoutes covers a client pinning a release +// date on a model the operator registered undated. Exact matches still win, +// so an operator who registers both dated releases keeps them distinct. +func TestRouter_DatedAnthropicModelRoutes(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "anthropic-prod", + Vendor: "anthropic", + Models: []string{"claude-sonnet-4-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + }}}) + + in := newInputWithModelAndURL("claude-sonnet-4-5-20250929", "/v1/messages") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "a dated id must route to the undated registration") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host) +} + +// TestRouter_ConnectionWarmProbeRoutes covers the HEAD /api/hello probe an +// Anthropic client sends before its first request. Forwarding it warms the +// connection that request will use; denying it only wrote a rejection into +// the access log at every session start. +func TestRouter_ConnectionWarmProbeRoutes(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{{ + ID: "anthropic-prod", + Vendor: "anthropic", + Models: []string{"claude-sonnet-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + }}}) + + in := newModellessInput("/api/hello") + in.Method = http.MethodHead + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.NotNil(t, out) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "the warm-up probe must reach the upstream") + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host) + + nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + assert.Equal(t, "true", nonInference, "the probe carries no model to gate on") +} + +// TestRouter_ModelListingCarriesAuthorisedModels pins the list the proxy +// bounds the discovery response with. A catch-all route enumerates nothing, +// so it must not bound the upstream's list at all. +func TestRouter_ModelListingCarriesAuthorisedModels(t *testing.T) { + enumerated := ProviderRoute{ + ID: "anthropic-prod", + Models: []string{"claude-sonnet-5", "claude-haiku-4-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + + t.Run("enumerated route bounds the listing", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{enumerated}}) + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models")) + require.NoError(t, err) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, + out.Mutations.RewriteUpstream.DiscoveryModels, + "the picker must be bounded by what the route authorises") + }) + + t.Run("catch-all route leaves the listing alone", func(t *testing.T) { + catchAll := enumerated + catchAll.Models = nil + mw := New(Config{Providers: []ProviderRoute{catchAll}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models")) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels, + "a route that claims every model cannot bound the upstream's list") + }) + + t.Run("per-model lookup is not a listing", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{enumerated}}) + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5")) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Empty(t, out.Mutations.RewriteUpstream.DiscoveryModels, + "the single-object lookup has no data array to filter") + }) +} + +// TestRouter_ModelDetailHonoursAllowlist pins that GET /v1/models/{id} is +// authorised against the model table. It carries no body model, so treating +// it as a model-less endpoint would let a caller confirm a model the route +// does not list — the listing itself is bounded to the allowlist, so the +// detail lookup must be too. +func TestRouter_ModelDetailHonoursAllowlist(t *testing.T) { + enumerated := ProviderRoute{ + ID: "anthropic-prod", + Models: []string{"claude-sonnet-5"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "api.anthropic.com", + } + + t.Run("allowlisted model routes and skips metering", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{enumerated}}) + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + assert.Equal(t, "api.anthropic.com", out.Mutations.RewriteUpstream.Host) + + nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + assert.Equal(t, "true", nonInference, "a detail lookup spends no tokens") + }) + + t.Run("model outside the allowlist denies", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{enumerated}}) + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-opus-5")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "a model no route lists must not be confirmed by the detail lookup") + }) + + t.Run("dated id matches its undated registration", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{enumerated}}) + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/claude-sonnet-5-20250929")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "a pinned release of an allowlisted family stays reachable") + }) + + t.Run("catch-all route still answers every lookup", func(t *testing.T) { + catchAll := enumerated + catchAll.Models = nil + mw := New(Config{Providers: []ProviderRoute{catchAll}}) + + out, err := mw.Invoke(context.Background(), newModellessInput("/v1/models/anything-at-all")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "a gateway that enumerates nothing cannot refuse a lookup") + }) +} + +// TestRouter_NonInferenceRequiresReadMethod pins that the non-inference mark — +// which exempts a request from the token pre-flight — is reachable only by the +// read methods these endpoints actually use. A POST to the same path could +// carry an inference body, so it must not buy the exemption; it falls through +// to normal per-model routing instead, which denies when no model is named. +func TestRouter_NonInferenceRequiresReadMethod(t *testing.T) { + route := ProviderRoute{ + ID: "gateway", + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "gateway.example.com", + } + + for _, path := range []string{"/v1/models", "/v1/models/claude-sonnet-5", "/api/hello"} { + t.Run("POST "+path, func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{route}}) + + in := newModellessInput(path) + in.Method = http.MethodPost + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "a write to a non-inference path must not route unmetered") + + nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + assert.NotEqual(t, "true", nonInference, + "only a read method may skip the token pre-flight") + }) + } + + t.Run("HEAD keeps the warm probe working", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{route}}) + + in := newModellessInput(connectionWarmPath) + in.Method = http.MethodHead + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "the HEAD warm probe must still reach the upstream") + + nonInference, _ := metaValue(t, out.Metadata, middleware.KeyLLMNonInference) + assert.Equal(t, "true", nonInference, + "the HEAD warm probe carries no model to meter") + }) +} + +// TestRouter_PinnedDatedModelStaysDistinct pins that a route registered +// against one dated Anthropic release does not claim another. Normalising +// both sides of the comparison made every dated build of a family +// interchangeable, so an operator who deliberately pinned a build would have +// served a different one — and with several such routes, declaration or path +// order would have decided which. +func TestRouter_PinnedDatedModelStaysDistinct(t *testing.T) { + pinned := ProviderRoute{ + ID: "anthropic-pinned", + Vendor: "anthropic", + Models: []string{"claude-sonnet-4-5-20250101"}, + AllowedGroupIDs: []string{defaultTestGroup}, + UpstreamScheme: "https", + UpstreamHost: "pinned.example.com", + } + + t.Run("a different dated release is not claimed", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{pinned}}) + in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, out.Decision, + "a route pinned to one dated build must not serve another") + }) + + t.Run("its own dated release still routes", func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{pinned}}) + in := newInputWithModelAndURL("claude-sonnet-4-5-20250101", "/v1/messages") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"}) + + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, "the exact match must still route") + }) + + t.Run("two pinned builds each route to their own provider", func(t *testing.T) { + other := pinned + other.ID = "anthropic-pinned-newer" + other.Models = []string{"claude-sonnet-4-5-20250202"} + other.UpstreamHost = "newer.example.com" + mw := New(Config{Providers: []ProviderRoute{pinned, other}}) + + in := newInputWithModelAndURL("claude-sonnet-4-5-20250202", "/v1/messages") + in.Metadata = append(in.Metadata, middleware.KV{Key: middleware.KeyLLMProvider, Value: "anthropic"}) + + 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) + assert.Equal(t, "newer.example.com", out.Mutations.RewriteUpstream.Host, + "declaration order must not decide between two deliberately pinned builds") + }) +} diff --git a/proxy/internal/middleware/decision.go b/proxy/internal/middleware/decision.go index 0970bdea4..97dca4af5 100644 --- a/proxy/internal/middleware/decision.go +++ b/proxy/internal/middleware/decision.go @@ -11,11 +11,78 @@ var codeRegex = regexp.MustCompile(`^[a-z][a-z0-9._-]{0,63}$`) // denyResponse is the on-wire shape rendered by RenderDenyResponse. // Keeping this as a typed struct ensures we never leak // middleware-supplied bytes outside known fields. +// +// Type and Error mirror the denial in the vendor's own error shape when +// the request reached a known LLM surface. LLM clients only parse their +// provider's envelope, so without the mirror a budget stop reaches the +// user as an unexplained API error. The NetBird fields stay where they +// were, so the body is a superset and existing consumers are unaffected. type denyResponse struct { Code string `json:"code"` Message string `json:"message,omitempty"` Details map[string]string `json:"details,omitempty"` Middleware string `json:"middleware,omitempty"` + Type string `json:"type,omitempty"` + Error *providerError `json:"error,omitempty"` +} + +// providerError is the nested error object both vendor envelopes carry. +type providerError struct { + Type string `json:"type"` + Message string `json:"message,omitempty"` + Code string `json:"code,omitempty"` +} + +// Vendor error types keyed by HTTP status, per each provider's published +// error reference. +const ( + anthropicErrInvalidRequest = "invalid_request_error" + anthropicErrPermission = "permission_error" + anthropicErrRateLimit = "rate_limit_error" + anthropicErrAPI = "api_error" + openAIErrInvalidRequest = "invalid_request_error" + openAIErrRateLimit = "rate_limit_error" +) + +// providerEnvelope returns the vendor-shaped mirror for a denial on the +// given surface, or nil when the surface has no envelope we can speak. +// message is the already-redacted public message. +func providerEnvelope(surface, code, message string, status int) (string, *providerError) { + switch surface { + case "anthropic": + return "error", &providerError{ + Type: anthropicErrorType(status), + Message: message, + } + case "openai": + return "", &providerError{ + Type: openAIErrorType(status), + Message: message, + Code: code, + } + default: + return "", nil + } +} + +func anthropicErrorType(status int) string { + switch status { + case http.StatusForbidden: + return anthropicErrPermission + case http.StatusTooManyRequests: + return anthropicErrRateLimit + case http.StatusBadRequest: + return anthropicErrInvalidRequest + default: + return anthropicErrAPI + } +} + +func openAIErrorType(status int) string { + if status == http.StatusTooManyRequests { + return openAIErrRateLimit + } + return openAIErrInvalidRequest } // RenderDenyResponse writes a structured JSON deny body. Status is @@ -36,6 +103,7 @@ func RenderDenyResponse(w http.ResponseWriter, middlewareID string, reason *Deny Message: truncate(Scan(reason.Message), 256), Middleware: truncate(Scan(middlewareID), 64), } + resp.Type, resp.Error = providerEnvelope(reason.Surface, resp.Code, resp.Message, status) if n := len(reason.Details); n > 0 { resp.Details = make(map[string]string, min(n, 8)) for k, v := range reason.Details { diff --git a/proxy/internal/middleware/decision_test.go b/proxy/internal/middleware/decision_test.go new file mode 100644 index 000000000..cf14c86ff --- /dev/null +++ b/proxy/internal/middleware/decision_test.go @@ -0,0 +1,92 @@ +package middleware + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// decodeDeny renders a denial and returns the parsed body plus the status. +func decodeDeny(t *testing.T, reason *DenyReason, status int) (map[string]any, int) { + t.Helper() + rec := httptest.NewRecorder() + RenderDenyResponse(rec, "llm_limit_check", reason, status) + + var body map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body), "deny body must be valid JSON") + return body, rec.Code +} + +// TestRenderDeny_AnthropicSurfaceMirrorsVendorShape covers a budget stop +// reaching Claude Code. The client only parses the Anthropic envelope, so +// without the mirror the user sees an unexplained API error instead of the +// reason their request was refused. +func TestRenderDeny_AnthropicSurfaceMirrorsVendorShape(t *testing.T) { + body, status := decodeDeny(t, &DenyReason{ + Code: "llm_policy.budget_cap_exceeded", + Message: "LLM policy limit exceeded", + Surface: "anthropic", + }, http.StatusForbidden) + + assert.Equal(t, http.StatusForbidden, status) + assert.Equal(t, "error", body["type"], "Anthropic errors carry type=error at the top level") + + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "error must be an object") + assert.Equal(t, "permission_error", errObj["type"], "403 maps to permission_error") + assert.Equal(t, "LLM policy limit exceeded", errObj["message"]) + + // The NetBird fields stay put so existing consumers keep working. + assert.Equal(t, "llm_policy.budget_cap_exceeded", body["code"]) + assert.Equal(t, "LLM policy limit exceeded", body["message"]) + assert.Equal(t, "llm_limit_check", body["middleware"]) +} + +// TestRenderDeny_OpenAISurfaceMirrorsVendorShape pins the OpenAI envelope, +// which nests the code and carries no top-level type. +func TestRenderDeny_OpenAISurfaceMirrorsVendorShape(t *testing.T) { + body, _ := decodeDeny(t, &DenyReason{ + Code: "llm_policy.model_blocked", + Message: "model is not in the policy allowlist", + Surface: "openai", + }, http.StatusForbidden) + + assert.NotContains(t, body, "type", "OpenAI errors have no top-level type") + + errObj, ok := body["error"].(map[string]any) + require.True(t, ok, "error must be an object") + assert.Equal(t, "invalid_request_error", errObj["type"]) + assert.Equal(t, "llm_policy.model_blocked", errObj["code"], "the NetBird code rides in the vendor code field") + assert.Equal(t, "model is not in the policy allowlist", errObj["message"]) +} + +// TestRenderDeny_RateLimitStatusMapsToVendorRateLimit pins the mapping a +// client's backoff keys on. +func TestRenderDeny_RateLimitStatusMapsToVendorRateLimit(t *testing.T) { + body, status := decodeDeny(t, &DenyReason{ + Code: "llm_policy.token_cap_exceeded", + Message: "LLM policy limit exceeded", + Surface: "anthropic", + }, http.StatusTooManyRequests) + + assert.Equal(t, http.StatusTooManyRequests, status, "429 must survive the status clamp") + errObj := body["error"].(map[string]any) + assert.Equal(t, "rate_limit_error", errObj["type"]) +} + +// TestRenderDeny_NoSurfaceKeepsLegacyShape guards non-LLM middlewares and +// denials raised before a surface is known. +func TestRenderDeny_NoSurfaceKeepsLegacyShape(t *testing.T) { + body, _ := decodeDeny(t, &DenyReason{ + Code: "llm_policy.model_not_routable", + Message: "no provider configured for model x", + }, http.StatusForbidden) + + assert.NotContains(t, body, "type", "no surface means no vendor mirror") + assert.NotContains(t, body, "error", "no surface means no vendor mirror") + assert.Equal(t, "llm_policy.model_not_routable", body["code"]) +} diff --git a/proxy/internal/middleware/keys.go b/proxy/internal/middleware/keys.go index 336bed19f..eff3fa756 100644 --- a/proxy/internal/middleware/keys.go +++ b/proxy/internal/middleware/keys.go @@ -22,6 +22,15 @@ const ( // body. Empty for clients that don't send one. KeyLLMSessionID = "llm.session_id" + // Sub-agent attribution (emitted by llm_request_parser from the + // client's request headers). A coding agent that spawns helpers + // stamps the spawned agent's id, and the spawning agent's id when + // the helper is itself nested, so cost within one session can be + // split across the agents that ran in parallel. These identify an + // agent, not a person or a device: never treat them as a user id. + KeyLLMAgentID = "llm.agent_id" + KeyLLMParentAgentID = "llm.parent_agent_id" + // LLM response-side metadata (emitted by llm_response_parser). //nolint:gosec // metadata key name, not a credential KeyLLMInputTokens = "llm.input_tokens" @@ -66,6 +75,14 @@ const ( // downstream gateways' spend logs. KeyLLMAuthorisingGroups = "llm.authorising_groups" + // LLM non-inference marker (emitted by llm_router on the allow path + // for endpoints that legitimately carry no model, such as model + // listing). The router still authorises these against the caller's + // groups; the marker only tells the limits gate that a per-model + // allowlist has nothing to evaluate, so an empty model must not be + // read as an undetermined one. Never derived from client input. + KeyLLMNonInference = "llm.non_inference" + // LLM policy attribution (emitted by llm_limit_check on the allow // path). Names the policy that paid for this request and the // dimension counters the post-flight llm_limit_record middleware diff --git a/proxy/internal/middleware/types.go b/proxy/internal/middleware/types.go index 1ed5c9d88..3c0ac0ab6 100644 --- a/proxy/internal/middleware/types.go +++ b/proxy/internal/middleware/types.go @@ -179,6 +179,12 @@ type DenyReason struct { Code string Message string Details map[string]string + // Surface names the LLM API dialect the caller speaks (the + // llm.provider value), so the rendered body can mirror the denial in + // that vendor's error shape alongside the NetBird fields. Empty for + // non-LLM middlewares and for denials raised before a surface was + // resolved; the body then carries the NetBird fields alone. + Surface string } // Output is the value each middleware returns to the dispatcher. The @@ -247,6 +253,12 @@ type UpstreamRewrite struct { // without verifying its TLS certificate. Set by llm_router from the // provider's skip_tls_verification for self-hosted / internal gateways. SkipTLSVerify bool + // DiscoveryModels, when non-empty, is the set of model ids the resolved + // route authorises, and the proxy drops everything else from the + // model-listing response. Empty leaves the upstream's list untouched, + // which is what a route that claims every model wants. Set by + // llm_router on a model-listing request only. + DiscoveryModels []string } // AuthHeader is a single name/value pair the proxy injects on the diff --git a/proxy/internal/proxy/discovery_filter.go b/proxy/internal/proxy/discovery_filter.go new file mode 100644 index 000000000..c9d606970 --- /dev/null +++ b/proxy/internal/proxy/discovery_filter.go @@ -0,0 +1,215 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "strconv" + "strings" + + sharedllm "github.com/netbirdio/netbird/shared/llm" +) + +// maxDiscoveryBodyBytes bounds the model-listing response the filter will +// buffer. A listing is a few kilobytes of ids; anything larger is not a +// listing we recognise, and buffering it to rewrite would cost more than +// the filtering is worth. +const maxDiscoveryBodyBytes = 1 << 20 + +// modelDiscoveryFilter returns a ModifyResponse hook that drops models the +// caller's policy does not authorise from a model-listing response, then +// delegates to next (which may be nil). +// +// Clients populate their model picker from this endpoint, so an unfiltered +// list offers models the very next request denies. The filter is +// best-effort: a response it cannot safely rewrite passes through +// untouched rather than reaching the client corrupted. +func modelDiscoveryFilter(allowed []string, next func(*http.Response) error) func(*http.Response) error { + permitted := make(map[string]struct{}, len(allowed)*2) + for _, id := range allowed { + permitted[id] = struct{}{} + permitted[sharedllm.NormalizeAnthropicModel(id)] = struct{}{} + } + + return func(resp *http.Response) error { + if err := filterModelListing(resp, permitted); err != nil { + return err + } + if next == nil { + return nil + } + return next(resp) + } +} + +// filterModelListing rewrites the response body in place, keeping only the +// entries whose id the policy authorises. Responses that are not a plain +// JSON listing are left alone. +func filterModelListing(resp *http.Response, permitted map[string]struct{}) error { + if !isPlainJSONListing(resp) { + return nil + } + + // One byte past the cap, so an oversized body is detectable without + // buffering all of it. + body, err := io.ReadAll(io.LimitReader(resp.Body, maxDiscoveryBodyBytes+1)) + if err != nil { + _ = resp.Body.Close() + return err + } + if len(body) > maxDiscoveryBodyBytes { + // Too large to filter. Put the bytes already read back in front of the + // unread remainder and forward the response exactly as the upstream + // sent it, headers included. Buffering what was read and closing here + // would truncate the body at the cap and hand the client a short, + // invalid listing — worse than not filtering at all. + resp.Body = spliceBody(body, resp.Body) + return nil + } + if err := resp.Body.Close(); err != nil { + return err + } + + filtered, ok := filterListingBody(body, permitted) + if !ok { + restoreBody(resp, body) + return nil + } + restoreBody(resp, filtered) + return nil +} + +// isPlainJSONListing reports whether the response is a JSON body the filter +// can parse. A content-encoded body is skipped: the transport only +// transparently decompresses what it negotiated itself, and the client +// negotiates its own encoding on this request. +func isPlainJSONListing(resp *http.Response) bool { + if resp == nil || resp.Body == nil { + return false + } + if resp.StatusCode != http.StatusOK { + return false + } + if enc := resp.Header.Get("Content-Encoding"); enc != "" && !strings.EqualFold(enc, "identity") { + return false + } + return strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") +} + +// filterListingBody returns the listing with unauthorised entries removed. +// ok is false when the body is not a listing shape, in which case the +// caller must forward the original bytes. +func filterListingBody(body []byte, permitted map[string]struct{}) ([]byte, bool) { + var doc map[string]json.RawMessage + if err := json.Unmarshal(body, &doc); err != nil { + return nil, false + } + raw, present := doc["data"] + if !present { + return nil, false + } + var entries []map[string]json.RawMessage + if err := json.Unmarshal(raw, &entries); err != nil { + return nil, false + } + + kept := make([]map[string]json.RawMessage, 0, len(entries)) + for _, entry := range entries { + if entryPermitted(entry, permitted) { + kept = append(kept, entry) + } + } + + encoded, err := json.Marshal(kept) + if err != nil { + return nil, false + } + doc["data"] = encoded + out, err := json.Marshal(doc) + if err != nil { + return nil, false + } + return out, true +} + +// entryPermitted reports whether a listing entry names a model the policy +// authorises, trying every form the same model is written in. +func entryPermitted(entry map[string]json.RawMessage, permitted map[string]struct{}) bool { + raw, ok := entry["id"] + if !ok { + return false + } + var id string + if err := json.Unmarshal(raw, &id); err != nil { + return false + } + for _, candidate := range modelIDForms(id) { + if _, ok := permitted[candidate]; ok { + return true + } + } + 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 — 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.Index(id, "/"); slash > 0 { + if _, ok := gatewayNamespaces[id[:slash]]; ok { + tail := id[slash+1:] + forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail)) + } + } + return forms +} + +// restoreBody puts body back on the response and fixes the length headers +// so the client reads exactly what is there. +// spliceBody returns a ReadCloser that yields prefix followed by whatever is +// left in rest, closing rest when closed. It lets the filter put back bytes it +// consumed while deciding, without owning the rest of the stream. +func spliceBody(prefix []byte, rest io.ReadCloser) io.ReadCloser { + return struct { + io.Reader + io.Closer + }{ + Reader: io.MultiReader(bytes.NewReader(prefix), rest), + Closer: rest, + } +} + +func restoreBody(resp *http.Response, body []byte) { + resp.Body = io.NopCloser(bytes.NewReader(body)) + resp.ContentLength = int64(len(body)) + resp.Header.Set("Content-Length", strconv.Itoa(len(body))) +} diff --git a/proxy/internal/proxy/discovery_filter_test.go b/proxy/internal/proxy/discovery_filter_test.go new file mode 100644 index 000000000..103eac594 --- /dev/null +++ b/proxy/internal/proxy/discovery_filter_test.go @@ -0,0 +1,235 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// jsonListingResponse builds a 200 model-listing response with the given +// body, as an upstream would return it. +func jsonListingResponse(body string) *http.Response { + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(body)), + ContentLength: int64(len(body)), + } + resp.Header.Set("Content-Type", "application/json") + return resp +} + +// listedIDs runs the filter and returns the ids left in the response. +func listedIDs(t *testing.T, allowed []string, body string) []string { + t.Helper() + resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body, replaced by the filter + require.NoError(t, modelDiscoveryFilter(allowed, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter + + raw, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + var doc struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(raw, &doc), "filtered body must stay valid JSON") + + ids := make([]string, 0, len(doc.Data)) + for _, entry := range doc.Data { + ids = append(ids, entry.ID) + } + return ids +} + +// TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels covers the picker a +// developer sees: an unfiltered upstream list offers every model the shared +// key can reach, and each one the policy excludes is a request the chain +// denies a moment later. +func TestModelDiscoveryFilter_KeepsOnlyAuthorisedModels(t *testing.T) { + ids := listedIDs(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, `{ + "data": [ + {"id": "claude-opus-5", "display_name": "Claude Opus 5"}, + {"id": "claude-sonnet-5", "display_name": "Claude Sonnet 5"}, + {"id": "claude-haiku-4-5"} + ], + "has_more": false + }`) + + assert.Equal(t, []string{"claude-sonnet-5", "claude-haiku-4-5"}, ids, + "only the models the route authorises may reach the picker") +} + +// TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs pins the two id forms +// a gateway returns for a model the operator registered plainly. +func TestModelDiscoveryFilter_MatchesDatedAndPrefixedIDs(t *testing.T) { + ids := listedIDs(t, []string{"claude-sonnet-4-5", "anthropic.claude-opus-5"}, `{ + "data": [ + {"id": "claude-sonnet-4-5-20250929"}, + {"id": "bedrock/anthropic.claude-opus-5"}, + {"id": "gpt-4o"} + ] + }`) + + assert.Equal(t, []string{"claude-sonnet-4-5-20250929", "bedrock/anthropic.claude-opus-5"}, ids, + "a dated or provider-prefixed id must match its registered form") +} + +// TestModelDiscoveryFilter_PreservesEnvelopeFields guards the rest of the +// document: clients read paging fields alongside data. +func TestModelDiscoveryFilter_PreservesEnvelopeFields(t *testing.T) { + resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}],"has_more":true,"first_id":"x"}`) //nolint:bodyclose // in-memory body, replaced by the filter + require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter + + raw, err := io.ReadAll(resp.Body) + require.NoError(t, err) + + var doc map[string]any + require.NoError(t, json.Unmarshal(raw, &doc)) + assert.Equal(t, true, doc["has_more"], "paging fields must survive the rewrite") + assert.Equal(t, "x", doc["first_id"]) + assert.Equal(t, strconv.Itoa(len(raw)), resp.Header.Get("Content-Length"), + "Content-Length must match the rewritten body") +} + +// TestModelDiscoveryFilter_PassesThroughUnfilterable covers the responses +// the filter must not touch: a compressed body it cannot parse, a non-JSON +// body, an error status, and a document with no data array. +func TestModelDiscoveryFilter_PassesThroughUnfilterable(t *testing.T) { + cases := map[string]func() *http.Response{ + "compressed": func() *http.Response { + resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`) + resp.Header.Set("Content-Encoding", "gzip") + return resp + }, + "not json": func() *http.Response { + resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`) + resp.Header.Set("Content-Type", "text/html") + return resp + }, + "error status": func() *http.Response { + resp := jsonListingResponse(`{"data":[{"id":"gpt-4o"}]}`) + resp.StatusCode = http.StatusInternalServerError + return resp + }, + "no data array": func() *http.Response { + return jsonListingResponse(`{"object":"list"}`) + }, + } + + for name, build := range cases { + t.Run(name, func(t *testing.T) { + resp := build() //nolint:bodyclose // in-memory body, replaced by the filter + original, err := io.ReadAll(resp.Body) + require.NoError(t, err) + resp.Body = io.NopCloser(bytes.NewReader(original)) + + require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, nil)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter + + got, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Equal(t, string(original), string(got), "an unfilterable response must reach the client unchanged") + }) + } +} + +// TestModelDiscoveryFilter_RunsNextHook pins that an existing +// ModifyResponse hook still runs after filtering. +func TestModelDiscoveryFilter_RunsNextHook(t *testing.T) { + called := false + next := func(*http.Response) error { + called = true + return nil + } + + resp := jsonListingResponse(`{"data":[{"id":"claude-sonnet-5"}]}`) //nolint:bodyclose // in-memory body, replaced by the filter + require.NoError(t, modelDiscoveryFilter([]string{"claude-sonnet-5"}, next)(resp)) //nolint:bodyclose // in-memory body, replaced by the filter + assert.True(t, called, "the chained hook must still run") +} + +// TestModelDiscoveryFilter_KeepsSlashBearingIDs covers self-hosted backends +// whose model ids carry a slash of their own. Treating the slash as a +// gateway prefix and keeping only the tail dropped every such model from +// the picker even though the policy named it exactly. +func TestModelDiscoveryFilter_KeepsSlashBearingIDs(t *testing.T) { + ids := listedIDs(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, `{ + "object": "list", + "data": [ + {"id": "Qwen/Qwen2.5-0.5B-Instruct"}, + {"id": "Qwen/Qwen2.5-7B-Instruct"} + ] + }`) + + assert.Equal(t, []string{"Qwen/Qwen2.5-0.5B-Instruct"}, ids, + "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 +// at exactly 1 MiB — valid-looking, short, and unparseable as JSON. The bytes +// already read must be spliced back in front of the unread remainder so the +// response reaches the client exactly as the upstream sent it. +func TestModelDiscoveryFilter_ForwardsOversizedBodyIntact(t *testing.T) { + // A well-formed listing whose single entry pads the body past the cap. + padding := strings.Repeat("x", maxDiscoveryBodyBytes) + body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}` + require.Greater(t, len(body), maxDiscoveryBodyBytes+1, + "the fixture must exceed the cap by more than the one-byte probe") + + resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body + require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body + + got, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Equal(t, len(body), len(got), + "an oversized listing must reach the client whole, not truncated at the cap") + assert.Equal(t, body, string(got), "the forwarded bytes must be the upstream's own") + + var doc map[string]json.RawMessage + assert.NoError(t, json.Unmarshal(got, &doc), + "the forwarded body must still parse as JSON") +} + +// TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders pins that the +// oversized path leaves the response metadata alone. Rewriting Content-Length +// to the truncated prefix is what made the corruption invisible to the client +// until it tried to parse. +func TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders(t *testing.T) { + padding := strings.Repeat("x", maxDiscoveryBodyBytes) + body := `{"object":"list","data":[{"id":"gpt-4o","note":"` + padding + `"}]}` + + resp := jsonListingResponse(body) //nolint:bodyclose // in-memory body + resp.Header.Set("Content-Length", strconv.Itoa(len(body))) + require.NoError(t, modelDiscoveryFilter(nil, nil)(resp)) //nolint:bodyclose // in-memory body + + assert.Equal(t, int64(len(body)), resp.ContentLength, + "ContentLength must keep describing the body the client receives") + assert.Equal(t, strconv.Itoa(len(body)), resp.Header.Get("Content-Length"), + "the Content-Length header must not be rewritten to the truncated prefix") +} diff --git a/proxy/internal/proxy/reverseproxy.go b/proxy/internal/proxy/reverseproxy.go index 9150c0329..7c9e21261 100644 --- a/proxy/internal/proxy/reverseproxy.go +++ b/proxy/internal/proxy/reverseproxy.go @@ -363,6 +363,9 @@ func (p *ReverseProxy) forwardUpstream(respWriter http.ResponseWriter, r *http.R if result.rewriteRedirects { rp.ModifyResponse = p.rewriteLocationFunc(effectiveURL, rewriteMatchedPath, r) //nolint:bodyclose } + if upstreamRewrite != nil && len(upstreamRewrite.DiscoveryModels) > 0 { + rp.ModifyResponse = modelDiscoveryFilter(upstreamRewrite.DiscoveryModels, rp.ModifyResponse) //nolint:bodyclose // the hook replaces the body and closes the original + } rp.ServeHTTP(respWriter, r.WithContext(ctx)) } diff --git a/proxy/server.go b/proxy/server.go index bd70b7e70..aee748339 100644 --- a/proxy/server.go +++ b/proxy/server.go @@ -2074,9 +2074,17 @@ func (s *Server) updateMapping(ctx context.Context, mapping *proto.ProxyMapping) return fmt.Errorf("auth setup for domain %s: %w", mapping.GetDomain(), err) } m := s.protoToMapping(ctx, mapping) - s.proxy.AddMapping(m) + // The chain is published before the route that leads to it. A request + // arriving at a target whose chain has not been rebuilt yet is served + // straight through, so a provider update that added the route first left a + // window in which an inference could complete unrouted and unmetered. + // Rebuilding first inverts that: the worst a request in the window meets is + // the new chain in front of the previous target, which is still counted. + if err := s.rebuildMiddlewareChains(svcID, m); err != nil { + return err + } s.meter.AddMapping(m) - s.rebuildMiddlewareChains(svcID, m) + s.proxy.AddMapping(m) return nil } @@ -2114,15 +2122,21 @@ func (s *Server) initMiddlewareManager(ctx context.Context) error { } // rebuildMiddlewareChains converts m into per-path bindings and calls -// Manager.Rebuild. Short-circuits when the middleware manager is unset. -func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) { +// Manager.Rebuild. Short-circuits when the middleware manager is unset, which +// is a deployment without middleware rather than a failure to install it. +// +// A rebuild that fails is reported rather than logged: the caller publishes +// the route once this returns, and a route published over chains that were +// not installed serves requests with no policy enforcement and no metering. +func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) error { if s.middlewareManager == nil { - return + return nil } bindings := buildMiddlewareBindings(svcID, m) if err := s.middlewareManager.Rebuild(string(svcID), bindings); err != nil { - s.Logger.WithError(err).WithField("service_id", svcID).Error("failed to rebuild middleware chains") + return fmt.Errorf("rebuild middleware chains for service %s: %w", svcID, err) } + return nil } // isLiveService reports whether svcID is currently present in the live diff --git a/shared/llm/model.go b/shared/llm/model.go index 08e42e5a4..4fb631520 100644 --- a/shared/llm/model.go +++ b/shared/llm/model.go @@ -46,6 +46,27 @@ func NormalizeBedrockModel(modelID string) string { return bedrockVersionSuffix.ReplaceAllString(m, "") } +// anthropicDatedModel matches a Claude model id carrying the trailing +// "-YYYYMMDD" release-date suffix Anthropic appends to a pinned release, +// capturing the id without it. The "claude" anchor is load-bearing: pricing +// looks every model up through this helper regardless of surface, and an +// operator may register a custom id with any shape at all, so an unanchored +// "-\d{8}$" would let "internal-llm-20250101" silently inherit the rate +// registered for "internal-llm". The anchor also covers the vendor-prefixed +// forms ("anthropic.claude-...", "us.anthropic.claude-..."). +var anthropicDatedModel = regexp.MustCompile(`(?i)^(.*claude.*)-\d{8}$`) + +// NormalizeAnthropicModel strips the trailing release-date suffix from a +// Claude model id, e.g. "claude-sonnet-4-5-20250929" -> "claude-sonnet-4-5", +// so a dated id a client pins matches the undated one the operator +// registered. Ids that are not Claude-family are returned untouched. +// Callers try the verbatim id first and fall back to this, so two dated +// releases of the same family stay distinct wherever both are registered +// explicitly. +func NormalizeAnthropicModel(modelID string) string { + return anthropicDatedModel.ReplaceAllString(modelID, "$1") +} + // NormalizeVertexModel strips the "@version" suffix from a Vertex AI model id // (e.g. "claude-sonnet-4-5@20250929" -> "claude-sonnet-4-5") so it matches // the catalog/pricing key. Vertex publisher models are priced under their diff --git a/shared/llm/model_test.go b/shared/llm/model_test.go index 42f2e9ca5..5ce2ff497 100644 --- a/shared/llm/model_test.go +++ b/shared/llm/model_test.go @@ -34,3 +34,29 @@ func TestNormalizeVertexModel(t *testing.T) { require.Equal(t, want, NormalizeVertexModel(in), "normalize %q", in) } } + +func TestNormalizeAnthropicModel(t *testing.T) { + cases := map[string]string{ + "claude-sonnet-4-5-20250929": "claude-sonnet-4-5", + "claude-3-5-haiku-20241022": "claude-3-5-haiku", + "claude-sonnet-5": "claude-sonnet-5", + "claude-opus-4-8": "claude-opus-4-8", + "anthropic.claude-haiku-4-5": "anthropic.claude-haiku-4-5", + "anthropic.claude-sonnet-4-5-20250929": "anthropic.claude-sonnet-4-5", + "us.anthropic.claude-opus-4-8-20250101": "us.anthropic.claude-opus-4-8", + // Non-Claude ids must survive untouched even when they end in eight + // consecutive digits: an operator can register a custom model under + // any id, and pricing looks every one of them up through this helper. + "gpt-4o": "gpt-4o", + "gpt-4o-2024-08-06": "gpt-4o-2024-08-06", + "gpt-4o-20240806": "gpt-4o-20240806", + "internal-llm-20250101": "internal-llm-20250101", + "deepseek-r1-20250120": "deepseek-r1-20250120", + "Qwen/Qwen2.5-20250101": "Qwen/Qwen2.5-20250101", + "gemini-2-5-pro-20250101": "gemini-2-5-pro-20250101", + "": "", + } + for in, want := range cases { + require.Equal(t, want, NormalizeAnthropicModel(in), "normalize %q", in) + } +}