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

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 {

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