[proxy,management] Conform the Agent Network endpoint to the LLM gateway protocol (#7154)

[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.

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)))
+}

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")
+}

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))
 }