[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.
2026-10-09 15:09:08 +02:00 · 2026-08-23 20:02:33 +02:00
parent ee253feddf
commit 766fcae3f8
39 changed files with 3094 additions and 154 deletions
@@ -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.
@@ -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://<alias>: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.