From 8afb903227e2c9dd8dca736030de36aaf9850d5b Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Fri, 21 Aug 2026 19:13:36 +0100 Subject: [PATCH] fix(mcp): backend timeouts, response validity, session handshake fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review round 2 fix batch for the federated /mcp endpoint, driven by a live probe against all 11 production MCP servers behind this gateway through the byte-pipe proxy (session-less tools/list): 8 answer bare, readitall and fetch reject it and need a session handshake, playwright is a legacy HTTP+SSE-transport server (400 Missing sessionId) that no handshake fixes. MF1: every backend call now runs under its own context.WithTimeout derived from r.Context() — toolsListBackendTimeout (20s) per fan-out call, toolsCallBackendTimeout (120s) for tools/call, named constants. MF2: doBackendJSONRPC rejects any non-2xx upstream status as an error (previously unchecked — a non-JSON-RPC error body only failed by accident of not parsing). parseBackendJSONRPC additionally rejects a validly-JSON body that carries neither "result" nor "error" — the exact silent-failure shape that would otherwise decode into an all-zero jsonrpcResponse and hand the client a bare {"jsonrpc":"2.0","id":1}. mcpFederatedToolsCall carries a defensive final guard against the same shape before relaying. MF3: mcpBackendCall is now bare-first-with-fallback — attempt the bare call; on a JSON-RPC error OR HTTP 4xx (never a network failure, timeout, or 5xx), retry once via initialize -> capture Mcp-Session-Id -> retry with the header -> best-effort DELETE close. A legacy HTTP+SSE-transport server (URL path ending "/sse") is excluded from federation entirely before ever reaching this — isLegacySSETransportURL, folded into allowedMCPServerNames — since no handshake fixes a different transport; it stays reachable via the per-server proxy. When every attempted server fails, tools/list returns a JSON-RPC internal error ("no MCP server reachable") instead of a silently-empty list. SF6: defaultMCPProtocolVersion -> "2025-11-25"; a caller's own protocolVersion is echoed back only from a small explicit allowlist (knownMCPProtocolVersions), falling back to the default otherwise — never echoes an arbitrary caller string as this gateway's own claimed capability. SF8: GET/DELETE/etc. on exactly /mcp is now 405 with Allow: POST, checked before auth (llmgateway.go). Folded minors: countTargetRequests batches per-target counter writes for tools/list's whole fan-out into ONE storeIncrMulti call after wg.Wait() (was one call per server); every "actually contacted" doc comment now says "attempted" (attribution is attempt-, not success-, accounting); lastSSEDataLine documents its no-cross-message- id-correlation limitation; routes_unified.go's stale buildLimitScopes/ withTotalScope comments (which claimed handleAdminAPI calls buildLimitScopes directly — it never has) are corrected; jsonrpcRequest documents why its "jsonrpc" field is decoded but never validated. Tests: ~30 new/updated cases covering per-backend timeout derivation (via the shorter-of-two-context-deadlines technique, no real 20s/120s waits), non-2xx rejection, the empty-envelope guard end-to-end, handshake-trigger classification (table-driven: JSON-RPC error and HTTP 4xx trigger, HTTP 500 and network failure do not), the full handshake sequence (bare -> initialize -> retry -> close, asserted call-by-call), handshake-gives-up-too, all-servers-failed, legacy-SSE exclusion (never contacted, in both tools/list and tools/call), protocol-version allowlist (all 4 versions + one rejected), 405/Allow, and best-effort session-close failure never surfacing to the caller. --- limits.go | 68 ++-- llmgateway.go | 17 +- mcp_federation.go | 518 +++++++++++++++++++++------ mcp_federation_test.go | 777 +++++++++++++++++++++++++++++++++++++++++ routes_unified.go | 29 +- 5 files changed, 1277 insertions(+), 132 deletions(-) diff --git a/limits.go b/limits.go index 3561b99..fee20e1 100644 --- a/limits.go +++ b/limits.go @@ -1011,14 +1011,16 @@ func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) { l.storeIncrMulti(entries) // no error return by contract (doc comment above); ok is intentionally discarded } -// countTargetRequest increments one MCP-server's or A2A agent's per-target -// request counters (Feature B, v0.21) at min, hour, day, AND month — the -// only scope kind this package tracks a month window of REQUESTS for. -// kind is targetScopeKind's own output ("mcp" or scopeKindAgent — -// mcp_a2a.go), id is the configured target name; both come embedded in -// windowKey's own key string, so a target scope can never collide with a -// user/group/total scope of the same id even by coincidence — kind is -// part of the key, not just a struct field. +// countTargetRequests increments MULTIPLE target scopes' per-target request +// counters (Feature B, v0.21) — each at min, hour, day, AND month, the only +// scope kind this package tracks a month window of REQUESTS for — in ONE +// storeIncrMulti call. kind is targetScopeKind's own output ("mcp" or +// scopeKindAgent — mcp_a2a.go), id is the configured target name; both come +// embedded in windowKey's own key string, so a target scope can never +// collide with a user/group/total scope of the same id even by coincidence +// — kind is part of the key, not just a struct field. A scope's own +// `limits` field is ignored entirely (never read); callers pass plain +// {kind, id} pairs. // // This is deliberately NOT folded into checkAndCount: a target scope // carries no limits by design (this round adds no config surface for @@ -1033,20 +1035,44 @@ func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) { // rejected as scope creep touching user/group accounting nothing else in // this round asked to change. // -// Callers: handleTargetProxy (mcp_a2a.go, once per proxied request, after -// its own user/group/total admission check passes) and the federated /mcp -// endpoint (mcp_federation.go, once per actual backend server contacted — -// tools/list may contact several in one incoming client request, tools/call -// exactly one), so a call attributed to a specific target is counted the -// same way regardless of which route reached it. -func (l *limiter) countTargetRequest(kind, id string) { +// Batching (review round 2, 2026-08-21): the federated /mcp endpoint's own +// tools/list fan-out can attempt several backend servers for ONE incoming +// client request — the earlier per-server countTargetRequest call inside +// each fan-out goroutine paid one round trip per server; this collects +// every scope from the whole fan-out into a single call after the fan-out's +// own wg.Wait(), the same "one request in, one round trip out" discipline +// checkAndCount/account already apply per scope-slice. +// +// Callers: handleTargetProxy (mcp_a2a.go, a 1-element slice, once per +// proxied request, after its own user/group/total admission check passes) +// and the federated /mcp endpoint (mcp_federation.go: a 1-element slice for +// tools/call's single resolved target, or one element per server ATTEMPTED +// — not necessarily reached or succeeded — for tools/list's fan-out; see +// mcpFederatedToolsList's own doc comment for why "attempted" is the +// correct word here). +func (l *limiter) countTargetRequests(scopes []limitScope) { + if len(scopes) == 0 { + return + } now := l.now() - l.storeIncrMulti([]counterIncr{ - newCounterIncr(kind, id, metricReq, windowMin, now, 1, minWindowTTL), - newCounterIncr(kind, id, metricReq, windowHour, now, 1, hourWindowTTL), - newCounterIncr(kind, id, metricReq, windowDay, now, 1, dayWindowTTL), - newCounterIncr(kind, id, metricReq, windowMonth, now, 1, monthWindowTTL), - }) + entries := make([]counterIncr, 0, len(scopes)*4) + for _, sc := range scopes { + entries = append(entries, + newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL), + newCounterIncr(sc.kind, sc.id, metricReq, windowMonth, now, 1, monthWindowTTL), + ) + } + l.storeIncrMulti(entries) +} + +// countTargetRequest is countTargetRequests for exactly one target scope — +// handleTargetProxy's own single-target shape (mcp_a2a.go), where a batch +// of one buys nothing but keeps the call site a plain two-argument call +// instead of a one-element slice literal. +func (l *limiter) countTargetRequest(kind, id string) { + l.countTargetRequests([]limitScope{{kind: kind, id: id}}) } // targetCounterKeysPerScope is the number of windowKey strings diff --git a/llmgateway.go b/llmgateway.go index 781cc95..3425a5c 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -455,7 +455,22 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { // there is no ambiguity between the two, but this ordering keeps every // exact-path route grouped together (mirroring "/v1/mcp/servers" and // "/v1/agents" immediately above) ahead of the prefix-matched one. - if r.Method == http.MethodPost && r.URL.Path == federatedMCPPath { + // + // Checked on r.URL.Path alone, BEFORE the method check: a GET or + // DELETE to exactly "/mcp" is a real, well-known route addressed with + // the wrong verb (SF8, review round 2, 2026-08-21) — 405 with an + // Allow header naming the one method this route accepts, matching + // RFC 9110 §15.5.6, rather than falling through to passthroughUnknown/ + // 404 the way a path this handler has never heard of would. No + // auth check runs first: which HTTP method a route accepts is + // unauthenticated request-shape information, not something worth + // gating behind a valid API key. + if r.URL.Path == federatedMCPPath { + if r.Method != http.MethodPost { + sw.Header().Set("Allow", http.MethodPost) + writeOAIError(sw, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed") + return + } u, grp, ok := g.auth.identify(r) g.logAuthEvent(ok, authEventUserName(u), r) if !ok { diff --git a/mcp_federation.go b/mcp_federation.go index b49cf41..d2386e7 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -4,11 +4,14 @@ import ( "bytes" "context" "encoding/json" + "fmt" "io" "net/http" + "net/url" "sort" "strings" "sync" + "time" ) // federatedMCPPath is the bare "/mcp" route this file implements — @@ -34,6 +37,35 @@ const ( jsonrpcInternalError = -32603 ) +// mcpSessionHeader is the MCP Streamable HTTP transport's session-id +// header — set by a backend server's "initialize" response to open a +// session, and echoed by a client (this gateway, on the handshake-fallback +// retry) on every subsequent call that session covers. See +// mcpBackendCall's own doc comment for when this gateway ever sends one. +const mcpSessionHeader = "Mcp-Session-Id" + +// Per-request budgets a federated call to a single backend server is +// allowed: the WHOLE per-server operation — a bare attempt, and, only when +// that signals a session is required, one initialize + one retry + a +// best-effort session close — must fit inside this window, not each +// individual HTTP request separately (mcpBackendCall's own doc comment). +// tools/list's fan-out runs one of these per allowed server, concurrently, +// so the shorter budget bounds a single client request's worst case to +// toolsListBackendTimeout regardless of how many servers are configured; +// tools/call talks to exactly one server and gets the longer budget a +// slower tool invocation (a web search, a browser action) may need. +const ( + toolsListBackendTimeout = 20 * time.Second + toolsCallBackendTimeout = 120 * time.Second + // mcpSessionCloseTimeout bounds the best-effort DELETE that ends a + // handshake-fallback session (mcpBackendCall). It runs on its own + // context.Background()-derived deadline, deliberately NOT the parent + // call's own (possibly already near-exhausted) toolsList/ + // toolsCallBackendTimeout budget: cleanup must never race, or get + // starved by, the real call it is cleaning up after. + mcpSessionCloseTimeout = 5 * time.Second +) + // jsonrpcRequest is one JSON-RPC 2.0 request/notification, decoded from // the client's POST body and re-encoded (with a synthetic id) for this // gateway's own outbound calls to a backend MCP server (mcpBackendCall). @@ -42,6 +74,15 @@ const ( // Only a single JSON-RPC object is supported, never a batch array: MCP's // own Streamable HTTP transport is the target wire format here, and // nothing in this round's brief calls for JSON-RPC batch support. +// +// JSONRPC itself is decoded but never validated against the literal +// string "2.0" — deliberately permissive: real MCP clients this gateway +// has actually seen (and the mock backends this repo's own tests drive) +// never send anything else, rejecting a request purely for carrying a +// different (or absent) "jsonrpc" field would reject a client this +// gateway can otherwise serve correctly for a purity check no consumer +// has ever needed, and the field is not read for any routing or safety +// decision this file makes. type jsonrpcRequest struct { Params json.RawMessage `json:"params,omitempty"` JSONRPC string `json:"jsonrpc"` @@ -78,10 +119,10 @@ type jsonrpcResponse struct { // minimalism — a JSON-RPC error is still carried at HTTP 200 (JSON-RPC // errors are a payload-level concept, not a transport-level one; every // genuinely transport-level failure this route can hit — an unreadable -// body, an unauthenticated/forbidden caller, a rate limit — is handled -// before this function is ever reached, via the gateway's own -// writeOAIError/writeLimitViolation envelope, matching every other route -// in this package). +// body, an unauthenticated/forbidden caller, a rate limit, a disallowed +// HTTP method — is handled before this function is ever reached, via the +// gateway's own writeOAIError/writeLimitViolation envelope, matching +// every other route in this package). func writeJSONRPCEnvelope(w http.ResponseWriter, resp jsonrpcResponse) { resp.JSONRPC = jsonrpcVersion w.Header().Set("Content-Type", "application/json") @@ -107,18 +148,55 @@ func writeJSONRPCErrorResponse(w http.ResponseWriter, id json.RawMessage, code i writeJSONRPCEnvelope(w, jsonrpcResponse{ID: id, Error: &jsonrpcError{Code: code, Message: message}}) } -// allowedMCPServerNames returns the sorted names of every MCP server -// grp.allowsMCP permits — the SAME per-server authorization check -// handleTargetProxy enforces (mcp_a2a.go), applied here across the whole -// catalog: tools/list aggregation only ever fans out to these names, -// and tools/call prefix resolution (resolveFederatedTool) only ever -// searches them — so a group restricted to N servers can see or invoke -// only those N servers' tools through POST /mcp too, exactly as it can -// only reach them through POST /mcp/{name} directly. +// isLegacySSETransportURL reports whether rawURL's path ends in "/sse" — +// the MCP prior HTTP+SSE transport's own URL convention (superseded by +// Streamable HTTP), where the CLIENT — not an intermediary — must drive a +// persistent SSE session: a GET establishing the event stream, POSTs +// correlated against a server-assigned endpoint the stream itself +// announces. mcpBackendCall's simple "POST a JSON-RPC object, read one +// response" model cannot speak that transport at all; there is no +// handshake or retry that fixes this, unlike the session-required case +// MF3 handles. Verified live against production traffic (review round 2, +// 2026-08-21): a real server on this transport (URL ending "/sse") +// answers a Streamable-HTTP-shaped POST with HTTP 400 "Missing sessionId" +// — a different failure MODE than the session-required 4xx/JSON-RPC-error +// case mcpBackendCall's handshake retry targets, and one no retry of any +// kind resolves. Such a server is excluded from federation's tools/list +// fan-out and tools/call resolution (allowedMCPServerNames) before ever +// reaching mcpBackendCall — it stays fully reachable via the per-server +// proxy (/mcp/{name}/...), where the real client drives the transport +// end to end, exactly as it always has. +// +// A rawURL that fails to parse returns false (not excluded) rather than +// panicking or erroring: every configured target URL already passed +// validateTargetURLs at construction (mcp_a2a.go), so this branch is +// unreachable in practice, and "don't exclude" is the safe default for an +// unreachable branch — the same convention sanitizeBaseURL applies to its +// own unparsable-URL case (admin.go). +func isLegacySSETransportURL(rawURL string) bool { + u, err := url.Parse(rawURL) + if err != nil { + return false + } + return strings.HasSuffix(u.Path, "/sse") +} + +// allowedMCPServerNames returns the sorted names of every MCP server BOTH +// grp.allowsMCP permits AND that speaks a transport federation can +// actually use (isLegacySSETransportURL excludes the rest) — the SAME +// per-server authorization check handleTargetProxy enforces (mcp_a2a.go), +// applied here across the whole catalog: tools/list aggregation only ever +// fans out to these names, and tools/call prefix resolution +// (resolveFederatedTool) only ever searches them — so a group restricted +// to N servers can see or invoke only those N servers' tools through +// POST /mcp too, exactly as it can only reach them through POST +// /mcp/{name} directly, and a legacy-SSE-transport server is invisible to +// federation regardless of group access (it is not a federation-eligible +// server for ANY caller, not an authorization decision). func allowedMCPServerNames(cfg *Config, grp *group) []string { names := make([]string, 0, len(cfg.MCPServers)) - for name := range cfg.MCPServers { - if grp.allowsMCP(name) { + for name, tc := range cfg.MCPServers { + if grp.allowsMCP(name) && !isLegacySSETransportURL(tc.URL) { names = append(names, name) } } @@ -128,7 +206,8 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string { // handleMCPFederated implements POST /mcp (federatedMCPPath): the // federated JSON-RPC endpoint this file's own doc comment explains the -// need for. +// need for. ServeHTTP (llmgateway.go) answers any OTHER method on this +// exact path with 405 (Allow: POST) before this function is ever called. // // Unlike handleTargetProxy, this is not a byte-level reverse proxy — it // cannot be: aggregating N servers' tools/list results into one response, @@ -139,8 +218,10 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string { // needs no backend at all) and, for "tools/list"/"tools/call", issues its // OWN outbound JSON-RPC call(s) to the relevant backend server(s) as part // of handling THIS one incoming request — no session state is ever -// persisted across separate client requests (mcpBackendCall's own doc -// comment), matching how the per-server proxy today carries no +// persisted across separate CLIENT requests (mcpBackendCall's own doc +// comment covers the one-call, per-backend-server session a handshake +// fallback DOES open and close, entirely inside a single incoming +// request), matching how the per-server proxy today carries no // gateway-side session state of its own either (it is a pure byte pipe; // whatever MCP session semantics exist are negotiated directly between // the client and that one backend, via headers proxyUpstream forwards @@ -153,11 +234,9 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string { // // checkAndCount runs once per incoming request, before method dispatch — // covering the caller's own user/group/total scopes exactly like every -// other metered route — with per-target attribution (countTargetRequest) -// added separately, per backend server actually contacted, inside -// mcpFederatedToolsList/mcpFederatedToolsCall: a single tools/list call -// may contact several servers, so "one request in, one target scope -// incremented" does not hold here the way it does for handleTargetProxy. +// other metered route — with per-target attribution added separately +// inside mcpFederatedToolsList/mcpFederatedToolsCall (their own doc +// comments cover exactly what "attributed" means for each). func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *user, grp *group) { if r.Header.Get("Upgrade") != "" { writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported") @@ -210,29 +289,55 @@ type mcpInitializeParams struct { ProtocolVersion string `json:"protocolVersion"` } -// defaultMCPProtocolVersion is reported only when the caller's own -// "initialize" request omits protocolVersion entirely. This is a -// fallback default this plugin picks for its own local response, not a -// verified compatibility claim about any downstream MCP server's -// supported version — see mcpFederatedInitialize's own doc comment for -// why the caller's own value is preferred whenever it sends one. -const defaultMCPProtocolVersion = "2025-06-18" +// defaultMCPProtocolVersion is reported when the caller's own "initialize" +// request omits protocolVersion, or names one outside +// knownMCPProtocolVersions. This is a fallback default this plugin picks +// for its own local response, not a verified compatibility claim about any +// downstream MCP server's supported version — see +// mcpFederatedInitialize's own doc comment for why the caller's own value +// is preferred whenever it names a recognized one. +const defaultMCPProtocolVersion = "2025-11-25" + +// knownMCPProtocolVersions is the small, explicit allowlist of protocol +// revisions mcpFederatedInitialize will ever echo back verbatim — every +// dated revision this plugin's own reading of the MCP spec's version +// history recognizes at the time this constant was last reviewed. +// Echoing an ARBITRARY caller-supplied string back as this gateway's own +// declared protocolVersion would let a malformed or malicious client make +// this response claim compatibility with a revision this plugin has never +// actually been evaluated against; the allowlist keeps the "echo the +// caller's choice" behavior (mcpFederatedInitialize's own doc comment) +// bounded to values this plugin actually knows about, falling back to +// defaultMCPProtocolVersion for anything else — including a genuinely +// newer, valid revision this constant simply predates; that is a +// known, accepted staleness risk of hardcoding a version list at all, +// not a correctness bug. +var knownMCPProtocolVersions = map[string]bool{ + "2024-11-05": true, + "2025-03-26": true, + "2025-06-18": true, + "2025-11-25": true, +} // mcpFederatedInitialize answers "initialize" locally — no backend server // is ever contacted for it (handleMCPFederated's own doc comment). -// protocolVersion echoes the caller's own requested value when present: -// version negotiation is the caller's choice to make, and echoing it back -// avoids this plugin hardcoding a claim about which dated MCP protocol -// revision it actually tracks. capabilities.tools is present (empty -// object, no sub-fields) to declare tool-calling support without -// overclaiming listChanged notifications this gateway never sends. +// protocolVersion echoes the caller's own requested value when it names a +// revision in knownMCPProtocolVersions: version negotiation is normally +// the caller's choice to make, but echoing an arbitrary unrecognized +// string back as this gateway's own declared capability would overclaim +// compatibility this plugin cannot back up (knownMCPProtocolVersions' own +// doc comment) — falling back to defaultMCPProtocolVersion in that case, +// and whenever the caller omits protocolVersion entirely. +// capabilities.tools is present (empty object, no sub-fields) to declare +// tool-calling support without overclaiming listChanged notifications +// this gateway never sends. func (g *Gateway) mcpFederatedInitialize(w http.ResponseWriter, req jsonrpcRequest) { var params mcpInitializeParams _ = json.Unmarshal(req.Params, ¶ms) // best-effort; empty/malformed params falls back to defaultMCPProtocolVersion below - protocolVersion := params.ProtocolVersion - if protocolVersion == "" { - protocolVersion = defaultMCPProtocolVersion + protocolVersion := defaultMCPProtocolVersion + if knownMCPProtocolVersions[params.ProtocolVersion] { + protocolVersion = params.ProtocolVersion } writeJSONRPCResult(w, req.ID, map[string]any{ @@ -262,39 +367,43 @@ type mcpToolsListResult struct { Tools []mcpTool `json:"tools"` } -// mcpBackendCall issues one JSON-RPC 2.0 request to targetURL — an MCP +// doBackendJSONRPC issues one JSON-RPC 2.0 POST to targetURL — an MCP // server's own configured base URL, the same single endpoint -// handleTargetProxy reverse-proxies every method to — and returns its -// parsed response. It carries a synthetic id ("federated"), never the -// client's own: tools/list may issue this call several times over for one -// incoming client request (once per allowed server), so reusing the -// client's id across N independent outbound calls would be meaningless, -// and the id on the way back is discarded by every caller of this -// function either way — only Result/Error is read. +// handleTargetProxy reverse-proxies every method to — carrying sessionID +// as the Mcp-Session-Id request header when non-empty. It is the one place +// that actually builds and sends an HTTP request for this file; every +// other function in the handshake-fallback machinery (mcpBackendCall, +// mcpBackendHandshake, mcpBackendCloseSession's own DELETE aside) goes +// through this. // -// No outbound "initialize" handshake precedes this call, and none of its -// own session/cookie state is cached across separate mcpBackendCall -// invocations — deliberately: per-request, throwaway upstream calls are -// what "keep upstream sessions per-request" (handleMCPFederated's own doc -// comment) means in practice here, matching the per-server proxy's own -// complete absence of gateway-side session tracking. A backend MCP server -// that hard-requires a prior initialize handshake before honoring -// tools/list or tools/call will reject this call; that is a real, -// documented trade-off of this design (see README's federation section), -// not an oversight. +// Return values: resp is nil whenever err is non-nil. status is the +// backend's real HTTP status when one was received at all (0 for a +// request-build or network failure that never got a response), reported +// even alongside a non-nil err — mcpBackendCall's own handshake-fallback +// decision (a JSON-RPC error OR an HTTP 4xx) needs the status specifically +// in the case where a non-2xx response has no valid JSON-RPC body at all +// (err is then the "upstream returned HTTP %d" error below, not a +// parse/validity error). respSessionID is the backend's own +// Mcp-Session-Id RESPONSE header — read on every call, not only +// "initialize", so mcpBackendHandshake needs no special-cased second +// code path for the one call that actually cares about it. // -// A backend may answer either a bare JSON object (Content-Type: -// application/json) or an MCP Streamable HTTP SSE stream (Content-Type: -// text/event-stream) carrying one or more "data:" events — both are -// accepted (parseBackendJSONRPC). -func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) { +// A non-2xx status is always an error (MF2, review round 2, 2026-08-21): +// the earlier version of this file had no such check at all, which let a +// non-JSON-RPC-shaped error page (a plain-text 500, an nginx 502 HTML +// page) fall through to parseBackendJSONRPC and fail there anyway, but +// only by accident of that body not happening to be valid JSON — a body +// that WAS valid JSON but not a real JSON-RPC response (see +// parseBackendJSONRPC's own validity check) would have silently produced +// an all-zero-fields response otherwise. +func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string, params any, sessionID string) (resp *jsonrpcResponse, status int, respSessionID string, err error) { paramsRaw, err := json.Marshal(params) if err != nil { - return nil, err + return nil, 0, "", err } reqBody, err := json.Marshal(jsonrpcRequest{JSONRPC: jsonrpcVersion, ID: json.RawMessage(`"federated"`), Method: method, Params: paramsRaw}) if err != nil { - return nil, err + return nil, 0, "", err } // gosec G704 (SSRF via taint analysis): targetURL is an @@ -307,22 +416,152 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, // site. httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(reqBody)) //nolint:gosec // operator-configured target URL, no client-supplied path/query if err != nil { - return nil, err + return nil, 0, "", err } httpReq.Header.Set("Content-Type", "application/json") httpReq.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + httpReq.Header.Set(mcpSessionHeader, sessionID) + } - resp, err := g.targetClient.Do(httpReq) + httpResp, err := g.targetClient.Do(httpReq) if err != nil { + return nil, 0, "", err + } + defer httpResp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure + + status = httpResp.StatusCode + respSessionID = httpResp.Header.Get(mcpSessionHeader) + + respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, maxRequestBytes)) + if err != nil { + return nil, status, respSessionID, err + } + if status < 200 || status >= 300 { + return nil, status, respSessionID, fmt.Errorf("upstream returned HTTP %d", status) + } + + parsed, err := parseBackendJSONRPC(httpResp.Header.Get("Content-Type"), respBody) + if err != nil { + return nil, status, respSessionID, err + } + return parsed, status, respSessionID, nil +} + +// mcpBackendHandshake issues a fresh, session-less "initialize" call +// against targetURL — minimal clientInfo, this gateway's own default +// protocol version — for mcpBackendCall's handshake-fallback retry. +// sessionID is "" when the backend's response carries no Mcp-Session-Id +// header of its own (a legitimate, spec-permitted case for a server that +// turns out to be stateless-tolerant after all); err is non-nil only when +// the initialize call itself failed outright (network error, non-2xx, a +// JSON-RPC error response to initialize itself) — there is no retry of a +// failed handshake. +func (g *Gateway) mcpBackendHandshake(ctx context.Context, targetURL string) (sessionID string, err error) { + initParams := map[string]any{ + "protocolVersion": defaultMCPProtocolVersion, + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "traefik-llmgateway", "version": pluginVersion}, + } + resp, _, respSessionID, err := g.doBackendJSONRPC(ctx, targetURL, "initialize", initParams, "") + if err != nil { + return "", err + } + if resp.Error != nil { + return "", fmt.Errorf("initialize: %s (code %d)", resp.Error.Message, resp.Error.Code) + } + return respSessionID, nil +} + +// mcpBackendCloseSession best-effort DELETEs targetURL with sessionID +// attached (the MCP Streamable HTTP transport's own session-end +// convention) — failure is logged, never surfaced: this runs after the +// caller's own real response has already been read and returned, so +// nothing downstream is waiting on it, and a backend that does not support +// session termination at all (silently ignoring the DELETE, or answering +// 404/405 for it) must never turn into a federation-visible error over +// pure cleanup. +func (g *Gateway) mcpBackendCloseSession(ctx context.Context, targetURL, sessionID string) { + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, targetURL, nil) //nolint:gosec // same operator-configured target URL as doBackendJSONRPC + if err != nil { + return + } + req.Header.Set(mcpSessionHeader, sessionID) + + resp, err := g.targetClient.Do(req) + if err != nil { + g.logf("federated mcp: best-effort session close failed for %q: %v", targetURL, err) + return + } + _ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure +} + +// mcpBackendCall issues a JSON-RPC 2.0 request for method to targetURL — +// bare first, matching the per-request "no persisted CLIENT session" +// design (handleMCPFederated's own doc comment) — and falls back to a +// ONE-SHOT initialize/retry/close handshake ONLY when the bare attempt +// signals the server actually requires a session: a JSON-RPC-level error +// response, or an HTTP 4xx status. +// +// This ruling is backed by a live probe against all 11 production servers +// behind this gateway (review round 2, 2026-08-21): 8 answer a bare, +// session-less call outright (brave-search, captcha-solver, google-maps, +// lightpanda, time, weather, web-search, wikipedia); readitall and fetch +// both reject the bare call — readitall: "method tools/list is invalid +// during..."; fetch: -32602 with or without params — and need exactly +// this handshake to work; a third class (a legacy HTTP+SSE-transport +// server, verified as playwright, HTTP 400 "Missing sessionId") is +// excluded from federation entirely before ever reaching this function +// (isLegacySSETransportURL, allowedMCPServerNames), since no handshake +// fixes a transport this function was never built to speak. +// +// A network failure, timeout, or 5xx never triggers the fallback: +// retrying those would spend a second timeout budget against a server +// that was never going to answer either way, bare or not — only a +// response that specifically signals "you're missing a session" is worth +// a second attempt. ctx governs the WHOLE operation (bare attempt, +// optional handshake, optional retry, best-effort close all share it, +// except the close itself — see mcpSessionCloseTimeout's own doc +// comment); callers derive it from context.WithTimeout(r.Context(), ...) +// with toolsListBackendTimeout or toolsCallBackendTimeout (MF1, review +// round 2, 2026-08-21). +func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) { + resp, status, _, err := g.doBackendJSONRPC(ctx, targetURL, method, params, "") + if err == nil && resp.Error == nil { + return resp, nil + } + + retryEligible := (err == nil && resp.Error != nil) || (status >= 400 && status < 500) + if !retryEligible { return nil, err } - defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure - respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxRequestBytes)) - if err != nil { - return nil, err + sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL) + if initErr != nil { + // The handshake itself failed: surface the bare attempt's own + // outcome when it has one (a JSON-RPC error response is more + // informative to a caller than a generic "handshake failed"), + // falling back to initErr only when the bare attempt produced no + // usable response at all (its own err was non-nil — the HTTP-4xx- + // with-no-JSON-RPC-body case). + if err != nil { + return nil, err + } + return resp, nil } - return parseBackendJSONRPC(resp.Header.Get("Content-Type"), respBody) + if sessionID != "" { + defer func() { + closeCtx, cancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout) + defer cancel() + g.mcpBackendCloseSession(closeCtx, targetURL, sessionID) + }() + } + + retryResp, _, _, retryErr := g.doBackendJSONRPC(ctx, targetURL, method, params, sessionID) + if retryErr != nil { + return nil, retryErr + } + return retryResp, nil } // parseBackendJSONRPC decodes body as a backend MCP server's JSON-RPC @@ -330,8 +569,21 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, // allows a POST response to take: a bare JSON object (any other // Content-Type is treated as this default case), or an SSE stream // (Content-Type: text/event-stream) of one or more "data:" lines — the -// LAST such line is parsed, matching a minimal backend that emits exactly -// one data event carrying its final response. +// LAST such line is parsed (lastSSEDataLine's own doc comment covers what +// this does and does not correlate). +// +// A body that parses as valid JSON but carries NEITHER "result" NOR +// "error" is also rejected (MF2, review round 2, 2026-08-21): a backend +// answering with some other JSON shape entirely (a plain `{}`, a +// non-compliant `{"status":"ok"}`) would otherwise decode into a +// jsonrpcResponse with every field at its zero value and LOOK like a +// legitimate, empty-but-successful JSON-RPC response — silently producing +// a client-visible `{"jsonrpc":"2.0","id":1}` with neither a result to +// read nor an error to report. Every caller of this function treats the +// returned error identically to a network failure (fails the call, +// mcpFederatedToolsList skips the server rather than crediting it with an +// empty success), so this is caught here once rather than at every call +// site. func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, error) { if strings.Contains(strings.ToLower(contentType), "text/event-stream") { body = lastSSEDataLine(body) @@ -340,6 +592,9 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err if err := json.Unmarshal(body, &out); err != nil { return nil, err } + if out.Result == nil && out.Error == nil { + return nil, fmt.Errorf("invalid JSON-RPC response: neither result nor error present") + } return &out, nil } @@ -348,6 +603,19 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err // the caller's own json.Unmarshal fails informatively on whatever body // actually was, rather than this function silently swallowing a // malformed stream into an empty byte slice). +// +// This does NOT implement per-request id correlation across multiple +// concurrent SSE-delivered messages the way a real MCP client's +// Streamable-HTTP transport layer would (matching a specific response to +// the request that triggered it by id, ignoring unrelated +// server-initiated notifications interleaved on the same stream) — it +// takes the textually LAST "data:" line, full stop. This is a deliberate, +// documented limitation, acceptable for mcpBackendCall's own one-request- +// in-flight-at-a-time usage of this stream (it never has two outstanding +// calls sharing one SSE response to correlate between): a backend that +// interleaves an unrelated notification AFTER its real response on the +// same stream would be misread by this function. No production server +// probed for this round exhibited that behavior. func lastSSEDataLine(body []byte) []byte { var last []byte for _, line := range bytes.Split(body, []byte("\n")) { @@ -366,22 +634,28 @@ func lastSSEDataLine(body []byte) []byte { // mcpFederatedToolsList implements the federated "tools/list": it fans // out one mcpBackendCall per server allowedMCPServerNames(grp) permits, -// concurrently, then merges every reachable server's own tools into one -// result — each tool's name prefixed "_", the exact -// underscore-separator convention the old agentgateway this plugin -// replaces used, and the one pugbot's mcpclient and agentkit have already -// persisted into their own tool-id databases (e.g. -// "brave-search_brave_web_search" — verified live). A server that errors, -// times out, or returns an unparsable result is logged and skipped, not -// surfaced as a whole-call failure: a federated aggregate degrading to -// "every OTHER server's tools" when one upstream is down is the useful -// behavior an aggregator should have, matching how a client would want a -// federated view to behave in practice. +// concurrently — each under its own toolsListBackendTimeout budget +// (context.WithTimeout off r.Context(), MF1) — then merges every +// reachable server's own tools into one result, each tool's name prefixed +// "_", the exact underscore-separator convention the +// old agentgateway this plugin replaces used, and the one pugbot's +// mcpclient and agentkit have already persisted into their own tool-id +// databases (e.g. "brave-search_brave_web_search" — verified live). A +// server that errors, times out, or returns an unparsable/invalid result +// is logged and skipped, not surfaced as a whole-call failure — UNLESS +// every single attempted server failed, in which case an empty tools list +// would misleadingly look like "this caller's group has no MCP access" — +// see the loud-failure branch below. // -// countTargetRequest attributes one request to EVERY server actually -// contacted here (handleMCPFederated's own doc comment) — a single +// countTargetRequests (limits.go) attributes one request to EVERY server +// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not +// "reached" or "succeeded": a server this gateway dialed and got a +// response (or a timeout, or a connection refusal) from still had a real +// request sent to it and a real slot of this gateway's outbound capacity +// spent on it, which is what these counters exist to track. A single // federated tools/list call can move several targets' own counters, not -// just one. +// just one — unlike handleTargetProxy's always-exactly-one-target shape +// (mcp_a2a.go). func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) { names := allowedMCPServerNames(g.cfg, grp) @@ -395,9 +669,10 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, wg.Add(1) go func(name string) { defer wg.Done() + ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout) + defer cancel() targetURL := g.cfg.MCPServers[name].URL - resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/list", struct{}{}) - g.limiter.countTargetRequest(targetKindMCP, name) + resp, err := g.mcpBackendCall(ctx, targetURL, "tools/list", struct{}{}) mu.Lock() defer mu.Unlock() @@ -417,12 +692,32 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, } wg.Wait() + if len(names) > 0 { + attempted := make([]limitScope, len(names)) + for i, name := range names { + attempted[i] = limitScope{kind: targetKindMCP, id: name} + } + g.limiter.countTargetRequests(attempted) + } + if len(failed) > 0 { sort.Strings(failed) g.logf("federated tools/list: %d of %d allowed MCP server(s) failed and were skipped: %v", len(failed), len(names), failed) } - sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name }) + // Every attempted server failed: an empty tools list here would read + // as "this caller's group has legitimate no-access", indistinguishable + // from the case where names itself was empty (no allowed servers at + // all — that IS a correct empty result, left alone below). A total + // outage is a real infrastructure problem the caller needs to see as + // one, not a silently degraded empty success (MF3, review round 2, + // 2026-08-21). + if len(names) > 0 && len(failed) == len(names) { + writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "no MCP server reachable") + return + } + + sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name }) writeJSONRPCResult(w, req.ID, mcpToolsListResult{Tools: merged}) } @@ -441,9 +736,10 @@ type mcpToolCallParams struct { // "foo_bar_lookup" must resolve to server "foo_bar"'s tool "lookup", not // server "foo"'s tool "bar_lookup"). ok is false when no allowed server's // prefix matches at all — including a tool that belongs to a real, -// configured server the caller's group simply cannot reach, since -// allowedNames (allowedMCPServerNames) already excludes it; this is what -// makes tools/call authz-aware without a second, separate access check. +// configured server the caller's group simply cannot reach, or one that +// speaks a transport federation excludes entirely, since allowedNames +// (allowedMCPServerNames) already excludes both; this is what makes +// tools/call authz- and transport-aware without a second, separate check. func resolveFederatedTool(fullName string, allowedNames []string) (serverName, toolName string, ok bool) { best := "" for _, name := range allowedNames { @@ -458,13 +754,20 @@ func resolveFederatedTool(fullName string, allowedNames []string) (serverName, t } // mcpFederatedToolsCall resolves req's "name" param to one allowed MCP -// server (resolveFederatedTool), forwards the call there with the server -// prefix stripped, and relays its result or error back to the client -// under the CLIENT's own id — never the synthetic one mcpBackendCall's -// outbound call used. An unresolvable prefix is a JSON-RPC invalid-params -// error, not an HTTP 404: the request reached a real method on a real -// route, it just named a tool nothing federation could route (an -// unconfigured server, or one this caller's group cannot reach). +// server (resolveFederatedTool), forwards the call there — under its own +// toolsCallBackendTimeout budget (context.WithTimeout off r.Context(), +// MF1) — with the server prefix stripped, and relays its result or error +// back to the client under the CLIENT's own id — never the synthetic one +// mcpBackendCall's outbound call used. An unresolvable prefix is a +// JSON-RPC invalid-params error, not an HTTP 404: the request reached a +// real method on a real route, it just named a tool nothing federation +// could route (an unconfigured server, one this caller's group cannot +// reach, or one on a transport federation excludes entirely). +// +// countTargetRequest (limits.go) attributes exactly one request to the +// single resolved server, once the call was actually attempted against +// it — the same "attempted, not necessarily succeeded" accounting +// mcpFederatedToolsList's own doc comment explains for its own fan-out. func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) { var params mcpToolCallParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { @@ -478,13 +781,28 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request, return } + ctx, cancel := context.WithTimeout(r.Context(), toolsCallBackendTimeout) + defer cancel() targetURL := g.cfg.MCPServers[serverName].URL - resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}) + resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}) g.limiter.countTargetRequest(targetKindMCP, serverName) if err != nil { g.logf("federated tools/call: server %q: %v", serverName, err) writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "upstream error") return } + + // Defensive final guard (MF2, review round 2, 2026-08-21): mcpBackendCall + // can only return a nil-err resp whose Result/Error are both nil if some + // future code path bypasses doBackendJSONRPC's own parseBackendJSONRPC + // validity check (which already rejects this shape as an error) — this + // is unreachable today, kept as the explicit guard the review round + // asked for rather than trusting that invariant silently. + if resp.Result == nil && resp.Error == nil { + g.logf("federated tools/call: server %q returned a JSON-RPC response with neither result nor error", serverName) + writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "invalid upstream response") + return + } + writeJSONRPCEnvelope(w, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error}) } diff --git a/mcp_federation_test.go b/mcp_federation_test.go index 4b938a8..7ee55a2 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -4,8 +4,11 @@ import ( "bytes" "context" "encoding/json" + "io" "net/http" "net/http/httptest" + "slices" + "sync" "sync/atomic" "testing" "time" @@ -181,6 +184,84 @@ func newMockJSONRPCServer(t *testing.T, tools []mcpTool) *mockJSONRPCServer { return m } +// sessionRequiredMockServer simulates a backend that rejects a bare +// (session-less) call with a JSON-RPC error but honors a proper +// initialize -> session -> retry handshake — modeling readitall/fetch's +// real, live-probed behavior (mcpBackendCall's own doc comment). It +// records the exact sequence of calls it received (method plus whichever +// session header, if any, accompanied it) so a test can assert the +// handshake shape precisely: bare attempt, initialize, retried real call +// carrying the session header, best-effort DELETE. +type sessionRequiredMockServer struct { + srv *httptest.Server + sessionID string + tools []mcpTool + calls []string + mu sync.Mutex + failInit bool +} + +func newSessionRequiredMockServer(t *testing.T, tools []mcpTool) *sessionRequiredMockServer { + t.Helper() + m := &sessionRequiredMockServer{tools: tools, sessionID: "sess-1"} + m.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSession := r.Header.Get(mcpSessionHeader) + + if r.Method == http.MethodDelete { + m.mu.Lock() + m.calls = append(m.calls, "DELETE:session="+gotSession) + m.mu.Unlock() + w.WriteHeader(http.StatusNoContent) + return + } + + var req jsonrpcRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("session-required mock: decode request: %v", err) + } + m.mu.Lock() + m.calls = append(m.calls, req.Method+":session="+gotSession) + m.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + if req.Method == "initialize" { + if m.failInit { + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "initialize refused"}}) + return + } + w.Header().Set(mcpSessionHeader, m.sessionID) + result, _ := json.Marshal(map[string]any{"protocolVersion": defaultMCPProtocolVersion, "capabilities": map[string]any{}}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + return + } + + if gotSession != m.sessionID { + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required: call initialize first"}}) + return + } + switch req.Method { + case "tools/list": + result, _ := json.Marshal(mcpToolsListResult{Tools: m.tools}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + case "tools/call": + var params mcpToolCallParams + _ = json.Unmarshal(req.Params, ¶ms) + echo, _ := json.Marshal(map[string]any{"calledName": params.Name}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: echo}) + default: + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: jsonrpcMethodNotFound, Message: "method not found"}}) + } + })) + t.Cleanup(m.srv.Close) + return m +} + +func (m *sessionRequiredMockServer) callSequence() []string { + m.mu.Lock() + defer m.mu.Unlock() + return append([]string(nil), m.calls...) +} + // newFederationTestConfig builds a Config with two MCP servers ("alpha", // "beta"), a group whose mcpServers glob restricts it to alphaAllowed // only when alphaOnly is true (otherwise unrestricted), and one inline @@ -703,3 +784,699 @@ func TestHandleMCPFederated_RateLimited_Returns429(t *testing.T) { t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String()) } } + +// --- MF1: per-backend-call timeouts --- + +func TestBackendTimeoutConstants(t *testing.T) { + if toolsListBackendTimeout != 20*time.Second { + t.Errorf("toolsListBackendTimeout = %v, want 20s", toolsListBackendTimeout) + } + if toolsCallBackendTimeout != 120*time.Second { + t.Errorf("toolsCallBackendTimeout = %v, want 120s", toolsCallBackendTimeout) + } +} + +// TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext +// proves the per-server context.WithTimeout(r.Context(), toolsListBackendTimeout) +// (MF1) actually derives its deadline from the INCOMING request's own +// context, not a freestanding timer: context.WithTimeout always resolves +// to the EARLIER of its parent's existing deadline and its own duration, +// so giving the incoming request a context with a much shorter deadline +// than toolsListBackendTimeout (20s) and confirming the whole call +// returns in well under a second — against a backend that hangs +// indefinitely — proves the wiring without ever waiting anywhere near the +// real 20s budget in this test suite. +func TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext(t *testing.T) { + block := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-block + })) + t.Cleanup(srv.Close) + t.Cleanup(func() { close(block) }) // must unblock the handler BEFORE srv.Close (which waits for it) — registered after, so it runs first (t.Cleanup is LIFO) + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) // alphaOnly: only the hanging server is reachable + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}) + shortCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + req = req.WithContext(shortCtx) + + start := time.Now() + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + elapsed := time.Since(start) + + if elapsed > 3*time.Second { + t.Fatalf("elapsed = %v, want well under toolsListBackendTimeout (20s) — the per-server timeout must derive from r.Context(), not ignore it", elapsed) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInternalError { + t.Errorf("error = %+v, want internal error (the only allowed server timed out, degrading to 'no MCP server reachable')", got.Error) + } +} + +// TestHandleMCPFederated_ToolsCall_SlowBackend_BoundedByRequestContext +// mirrors the tools/list timeout test above for tools/call's own, +// separate context.WithTimeout(r.Context(), toolsCallBackendTimeout) call +// site. +func TestHandleMCPFederated_ToolsCall_SlowBackend_BoundedByRequestContext(t *testing.T) { + block := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + <-block + })) + t.Cleanup(srv.Close) + t.Cleanup(func() { close(block) }) + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams}) + shortCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + req = req.WithContext(shortCtx) + + start := time.Now() + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + elapsed := time.Since(start) + + if elapsed > 3*time.Second { + t.Fatalf("elapsed = %v, want well under toolsCallBackendTimeout (120s) — the per-call timeout must derive from r.Context(), not ignore it", elapsed) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInternalError { + t.Errorf("error = %+v, want internal error (the resolved server timed out)", got.Error) + } +} + +// --- MF2: non-2xx status and empty-envelope guards --- + +// TestHandleMCPFederated_ToolsCall_BackendNon2xxStatus_IsError proves a +// non-2xx upstream response is always an error, even when its body +// happens to be a validly-shaped JSON-RPC success object — a 5xx (or any +// non-2xx) status is a transport-level failure the body's own content can +// never override. +func TestHandleMCPFederated_ToolsCall_BackendNon2xxStatus_IsError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","result":{}}`)) + })) + defer srv.Close() + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInternalError { + t.Errorf("error = %+v, want internal error — a 5xx status must never be rescued by a well-shaped body", got.Error) + } +} + +// TestHandleMCPFederated_ToolsCall_BackendEmptyEnvelope_ConvertsToInternalError +// is the reviewer's named silent-failure scenario (MF2): a backend answers +// 200 with valid JSON that is NOT a real JSON-RPC response (neither +// "result" nor "error") — without the guard, this decodes into an +// all-zero-fields jsonrpcResponse and the client would silently receive +// {"jsonrpc":"2.0","id":1} with no indication anything went wrong. +func TestHandleMCPFederated_ToolsCall_BackendEmptyEnvelope_ConvertsToInternalError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","status":"ok"}`)) + })) + defer srv.Close() + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil { + t.Fatal("want a JSON-RPC error, not a silent empty success — this is the reviewer's exact silent-failure scenario") + } + if got.Error.Code != jsonrpcInternalError { + t.Errorf("error code = %d, want %d", got.Error.Code, jsonrpcInternalError) + } + if len(got.Result) != 0 { + t.Errorf(`result = %s, want empty — must never emit {"jsonrpc":"2.0","id":1} silently`, got.Result) + } +} + +// --- MF3: bare-first-with-fallback handshake retry --- + +// TestMcpBackendCall_HandshakeFallbackTrigger proves the fallback fires +// ONLY on a JSON-RPC-level error or an HTTP 4xx status from the bare +// attempt — never on a 5xx — table-driven over exactly the two trigger +// shapes and the one non-trigger shape (a network failure is covered +// separately below, since it needs a different mock shape entirely: no +// server to answer at all). +func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) { + cases := []struct { + bareHandler http.HandlerFunc + name string + wantHandshake bool + }{ + { + name: "bare JSON-RPC error triggers the handshake", + bareHandler: func(w http.ResponseWriter, r *http.Request) { + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}}) + }, + wantHandshake: true, + }, + { + name: "HTTP 403 with no JSON-RPC body triggers the handshake", + bareHandler: func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + }, + wantHandshake: true, + }, + { + name: "HTTP 500 never triggers the handshake", + bareHandler: func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }, + wantHandshake: false, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var callCount atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if callCount.Add(1) == 1 { + tc.bareHandler(w, r) + return + } + if r.Method == http.MethodDelete { + w.WriteHeader(http.StatusNoContent) + return + } + body, _ := io.ReadAll(r.Body) + var req jsonrpcRequest + _ = json.Unmarshal(body, &req) + w.Header().Set("Content-Type", "application/json") + result, _ := json.Marshal(map[string]any{}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + })) + defer srv.Close() + + cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + _, _ = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}) + + gotCalls := callCount.Load() + if tc.wantHandshake && gotCalls < 2 { + t.Errorf("calls = %d, want >=2 (bare attempt + initialize handshake)", gotCalls) + } + if !tc.wantHandshake && gotCalls != 1 { + t.Errorf("calls = %d, want exactly 1 (bare attempt only, no handshake retry against a 5xx)", gotCalls) + } + }) + } +} + +// TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake proves a +// connection-level failure (nothing listening at all) never triggers the +// handshake fallback either — the same "not retry-eligible" rule as a +// 5xx, verified against a genuinely different failure shape (doBackendJSONRPC +// never even gets an HTTP status back). +func TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake(t *testing.T) { + cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + gw, ok := h.(*Gateway) + if !ok { + t.Fatal("handler is not *Gateway") + } + + _, err = gw.mcpBackendCall(context.Background(), "http://127.0.0.1:1", "tools/list", struct{}{}) + if err == nil { + t.Fatal("want an error (nothing listening on 127.0.0.1:1)") + } +} + +// TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry +// is MF3's core proof: a server that rejects a bare tools/list call +// (readitall/fetch's real, live-probed behavior — sessionRequiredMockServer's +// own doc comment) still ends up contributing its tools to the federated +// aggregate, via exactly the documented bare -> initialize -> retry -> +// close sequence. +func TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry(t *testing.T) { + strict := newSessionRequiredMockServer(t, []mcpTool{{Name: "search"}}) + cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true) // alphaOnly: only "alpha" (the strict server) is reachable + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error != nil { + t.Fatalf("error = %+v, want nil (the handshake fallback should have recovered)", got.Error) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_search" { + t.Fatalf("tools = %+v, want exactly alpha_search", result.Tools) + } + + seq := strict.callSequence() + if len(seq) != 4 { + t.Fatalf("call sequence = %v, want exactly 4 calls (bare attempt, initialize, retried real call, best-effort close — the DELETE runs synchronously inside mcpBackendCall's own defer, before it returns)", seq) + } + if seq[0] != "tools/list:session=" { + t.Errorf("call[0] = %q, want a bare tools/list attempt with no session header", seq[0]) + } + if seq[1] != "initialize:session=" { + t.Errorf("call[1] = %q, want an initialize handshake with no session header", seq[1]) + } + if seq[2] != "tools/list:session=sess-1" { + t.Errorf("call[2] = %q, want the retried tools/list carrying the session header the handshake returned", seq[2]) + } + if seq[3] != "DELETE:session=sess-1" { + t.Errorf("call[3] = %q, want the best-effort session close carrying the same session header", seq[3]) + } +} + +// TestHandleMCPFederated_ToolsCall_HandshakeFallback_SucceedsAfterSessionRetry +// mirrors the tools/list handshake test above for tools/call's own +// single-target path, additionally asserting the best-effort session +// close (DELETE) actually happened. +func TestHandleMCPFederated_ToolsCall_HandshakeFallback_SucceedsAfterSessionRetry(t *testing.T) { + strict := newSessionRequiredMockServer(t, nil) + cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_search"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage(`"client-7"`), Params: callParams})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error != nil { + t.Fatalf("error = %+v, want nil (the handshake fallback should have recovered)", got.Error) + } + if string(got.ID) != `"client-7"` { + t.Errorf("id = %s, want the client's own id echoed back", got.ID) + } + + seq := strict.callSequence() + if !slices.Contains(seq, "tools/call:session=") { + t.Errorf("call sequence = %v, want a bare tools/call attempt with no session header", seq) + } + if !slices.Contains(seq, "initialize:session=") { + t.Errorf("call sequence = %v, want an initialize handshake", seq) + } + if !slices.Contains(seq, "tools/call:session=sess-1") { + t.Errorf("call sequence = %v, want the retried tools/call carrying the session header", seq) + } + if !slices.Contains(seq, "DELETE:session=sess-1") { + t.Errorf("call sequence = %v, want a best-effort DELETE closing the session", seq) + } +} + +// TestHandleMCPFederated_ToolsList_HandshakeFallback_GivesUpWhenInitializeAlsoFails +// proves a server that rejects the handshake's own initialize call too +// ends up in the failed/skipped list (not a crash, not a hang) — and, +// since it is the only allowed server, the aggregate degrades all the way +// to the "no MCP server reachable" loud failure (MF3), not an empty +// success. +func TestHandleMCPFederated_ToolsList_HandshakeFallback_GivesUpWhenInitializeAlsoFails(t *testing.T) { + strict := newSessionRequiredMockServer(t, []mcpTool{{Name: "search"}}) + strict.failInit = true + cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInternalError || got.Error.Message != "no MCP server reachable" { + t.Errorf("error = %+v, want internal error %q", got.Error, "no MCP server reachable") + } +} + +// TestHandleMCPFederated_ToolsList_AllServersFail_ReturnsInternalError +// proves that when every attempted server fails, the response is a loud +// JSON-RPC error, not a silently-empty tools list that would look +// identical to "this caller's group has no MCP access at all". +func TestHandleMCPFederated_ToolsList_AllServersFail_ReturnsInternalError(t *testing.T) { + cfg := newFederationTestConfig("http://127.0.0.1:1", "http://127.0.0.1:2", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInternalError || got.Error.Message != "no MCP server reachable" { + t.Errorf("error = %+v, want internal error %q", got.Error, "no MCP server reachable") + } +} + +// --- MF3: legacy HTTP+SSE transport exclusion --- + +func TestIsLegacySSETransportURL(t *testing.T) { + cases := []struct { + url string + want bool + }{ + {"http://playwright.internal:8080/sse", true}, + {"http://playwright.internal:8080/sse?foo=bar", true}, + {"http://mcp.internal:8080", false}, + {"http://mcp.internal:8080/mcp", false}, + {"http://mcp.internal:8080/sse/extra", false}, + {"://not a valid url", false}, // malformed -> safe default: not excluded + } + for _, tc := range cases { + t.Run(tc.url, func(t *testing.T) { + if got := isLegacySSETransportURL(tc.url); got != tc.want { + t.Errorf("isLegacySSETransportURL(%q) = %v, want %v", tc.url, got, tc.want) + } + }) + } +} + +func TestAllowedMCPServerNames_ExcludesLegacySSETransport(t *testing.T) { + cfg := CreateConfig() + cfg.MCPServers = map[string]*TargetConfig{ + "alpha": {URL: "http://alpha.internal"}, + "playwright": {URL: "http://playwright.internal/sse"}, + "beta": {URL: "http://beta.internal"}, + } + grp := &group{name: "g"} // unrestricted mcpServers: matches everything except the transport exclusion + got := allowedMCPServerNames(cfg, grp) + want := []string{"alpha", "beta"} + if !slices.Equal(got, want) { + t.Errorf("allowedMCPServerNames = %v, want %v (playwright excluded: legacy HTTP+SSE transport)", got, want) + } +} + +// TestHandleMCPFederated_ToolsList_ExcludesLegacySSEServer_NeverContacted +// proves a legacy-SSE-transport server is not merely omitted from the +// aggregate's RESULT, but never contacted by the fan-out at all. +func TestHandleMCPFederated_ToolsList_ExcludesLegacySSEServer_NeverContacted(t *testing.T) { + alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}}) + legacy := newMockJSONRPCServer(t, []mcpTool{{Name: "browse"}}) + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{ + "alpha": {URL: alpha.srv.URL}, + "playwright": {URL: legacy.srv.URL + "/sse"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_lookup" { + t.Fatalf("tools = %+v, want exactly alpha_lookup (playwright excluded)", result.Tools) + } + if legacy.called.Load() != 0 { + t.Error("the legacy-SSE-transport server must never be contacted by tools/list fan-out") + } +} + +// TestHandleMCPFederated_ToolsCall_LegacySSEServer_UnresolvableLikeUnconfigured +// proves tools/call cannot route to a legacy-SSE-transport server either — +// unresolvable, same as an unconfigured name, and never contacted. +func TestHandleMCPFederated_ToolsCall_LegacySSEServer_UnresolvableLikeUnconfigured(t *testing.T) { + legacy := newMockJSONRPCServer(t, nil) + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"playwright": {URL: legacy.srv.URL + "/sse"}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + callParams, _ := json.Marshal(mcpToolCallParams{Name: "playwright_click"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams})) + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInvalidParams { + t.Errorf("error = %+v, want invalid params (legacy-SSE server must be unresolvable via tools/call, same as unconfigured)", got.Error) + } + if legacy.called.Load() != 0 { + t.Error("the legacy-SSE-transport server must never be contacted") + } +} + +// --- SF6: protocol-version allowlist --- + +func TestHandleMCPFederated_Initialize_AllowlistedVersionsEchoedVerbatim(t *testing.T) { + for version := range knownMCPProtocolVersions { + t.Run(version, func(t *testing.T) { + cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + params, _ := json.Marshal(map[string]any{"protocolVersion": version}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1"), Params: params})) + + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + var result struct { + ProtocolVersion string `json:"protocolVersion"` + } + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if result.ProtocolVersion != version { + t.Errorf("protocolVersion = %q, want the allowlisted caller-supplied version %q echoed verbatim", result.ProtocolVersion, version) + } + }) + } +} + +func TestHandleMCPFederated_Initialize_UnrecognizedVersion_FallsBackToDefault(t *testing.T) { + cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + params, _ := json.Marshal(map[string]any{"protocolVersion": "1999-01-01"}) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1"), Params: params})) + + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + var result struct { + ProtocolVersion string `json:"protocolVersion"` + } + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if result.ProtocolVersion != defaultMCPProtocolVersion { + t.Errorf("protocolVersion = %q, want the fallback default %q for an unrecognized caller-supplied version", result.ProtocolVersion, defaultMCPProtocolVersion) + } +} + +// --- SF8: explicit 405 for a disallowed method on exactly /mcp --- + +func TestHandleMCPFederated_DisallowedMethod_Returns405WithAllowHeader(t *testing.T) { + cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for _, method := range []string{http.MethodGet, http.MethodDelete, http.MethodPut} { + t.Run(method, func(t *testing.T) { + req := httptest.NewRequest(method, federatedMCPPath, nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusMethodNotAllowed { + t.Fatalf("status = %d, want 405, body=%s", rec.Code, rec.Body.String()) + } + if got := rec.Header().Get("Allow"); got != http.MethodPost { + t.Errorf("Allow header = %q, want %q", got, http.MethodPost) + } + }) + } + + t.Run("no auth still 405 (method check precedes auth)", func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, federatedMCPPath, nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusMethodNotAllowed { + t.Fatalf("status = %d, want 405 even without a valid API key, body=%s", rec.Code, rec.Body.String()) + } + }) +} + +// TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller proves +// mcpBackendCloseSession's own "best-effort... failure logged, never +// surfaced" contract: the server accepts the handshake and the retried +// real call normally, but drops the connection outright on the +// best-effort DELETE — the caller must still see the real call's own +// successful result, not a failure manufactured by cleanup. +func TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller(t *testing.T) { + const sessionID = "sess-close-fail" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodDelete { + // A network-level failure on close, not merely a non-2xx + // status: mcpBackendCloseSession never even inspects the + // response status, only whether Do() itself returned an + // error — hijacking and dropping the raw connection is what + // actually exercises that path. + hj, ok := w.(http.Hijacker) + if !ok { + t.Fatal("mock server: ResponseWriter does not support Hijack") + } + conn, _, err := hj.Hijack() + if err != nil { + t.Fatalf("hijack: %v", err) + } + _ = conn.Close() + return + } + + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + w.Header().Set("Content-Type", "application/json") + switch { + case req.Method == "initialize": + w.Header().Set(mcpSessionHeader, sessionID) + result, _ := json.Marshal(map[string]any{}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + case req.Method == "tools/list" && r.Header.Get(mcpSessionHeader) == "": + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}}) + default: + result, _ := json.Marshal(mcpToolsListResult{Tools: []mcpTool{{Name: "search"}}}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + } + })) + defer srv.Close() + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + rec := httptest.NewRecorder() + h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode: %v", err) + } + if got.Error != nil { + t.Fatalf("error = %+v, want nil — a failed best-effort session close must never surface as a call failure", got.Error) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_search" { + t.Fatalf("tools = %+v, want alpha_search despite the session-close failure", result.Tools) + } +} diff --git a/routes_unified.go b/routes_unified.go index bd68374..99fe9f6 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -244,10 +244,15 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr // contract (admin.go) that the dashboard's usage table already promised. // // Callers metering actual LLM traffic wrap this result in withTotalScope -// before passing it to checkAndCount/account; handleAdminAPI (admin.go) -// calls this directly, without withTotalScope, so an admin request's own -// req/min-req/day accounting never contributes to the total-scope -// LLM-traffic series (v0.2 data-layer task). +// before passing it to checkAndCount/account. handleAdminAPI (admin.go) +// never calls buildLimitScopes at all — its own buildAdminUsage builds +// {kind, id, limits} literals directly from authStore.snapshot's user/group +// listing, for currentUsage's read-only purposes — and, separately, never +// calls checkAndCount either (handleAdminAPI's own doc comment: admin +// traffic must never move req/min-req/day statistics). Fixed a stale +// version of this comment (review round 2, 2026-08-21) that claimed +// handleAdminAPI "calls this directly, without withTotalScope" — it never +// called this function at all, on any version of this file. func buildLimitScopes(u *user, grp *group) []limitScope { return []limitScope{ {limits: u.limits, kind: "user", id: u.name}, @@ -259,12 +264,16 @@ func buildLimitScopes(u *user, grp *group) []limitScope { // (totalScopeKind/totalScopeID, limits.go) appended — the single helper // every metered route (runUnified above; resolveMediaRequest, // routes_media.go; handlePassthrough, routes_passthrough.go; -// handleTargetProxy, mcp_a2a.go) calls around its own buildLimitScopes -// result, so none of them can forget it and none of them duplicate the -// scope literal. Deliberately not folded into buildLimitScopes itself: -// that helper is also called directly by handleAdminAPI (admin.go) for -// its own req/min-req/day accounting, and admin traffic must never -// contribute to the total LLM-traffic scope. +// handleTargetProxy, mcp_a2a.go; handleMCPFederated, mcp_federation.go) +// calls around its own buildLimitScopes result, so none of them can forget +// it and none of them duplicate the scope literal. Deliberately not folded +// into buildLimitScopes itself: buildLimitScopes' own doc comment now +// explains why admin traffic needs no special-casing here at all — +// handleAdminAPI never calls either function, so there was never a real +// "admin traffic must not contribute to total" case for this split to +// guard against; the split is kept anyway because buildLimitScopes' +// {user, group} pair and the synthetic total scope are conceptually +// different additions, worth two names. func withTotalScope(scopes []limitScope) []limitScope { return append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil}) }