diff --git a/mcp_federation.go b/mcp_federation.go index 37dae6c..b916bd1 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -174,38 +174,140 @@ type jsonrpcResponse struct { ID json.RawMessage `json:"id"` } +// mcpResponseFormat is this route's negotiated response framing, decided +// once per request (mcpNegotiateFormat, called from handleMCPFederated) +// and threaded explicitly as a parameter through every function that can +// end in a call to writeJSONRPCEnvelope — handleMCPFederated itself, +// mcpFederatedInitialize, mcpFederatedToolsList, mcpFederatedToolsCall — +// down to that one chokepoint. Threaded explicitly, not carried by +// wrapping w in a ResponseWriter that rewrites Content-Type and reframes +// bytes after the fact: a wrapper like that has to buffer or intercept +// every Write to retroactively reshape output it already committed, +// exactly the kind of implicit, hard-to-test indirection this package +// avoids elsewhere (see writeJSONRPCEnvelope's own "single chokepoint" +// shape, which this preserves). +type mcpResponseFormat int + +const ( + // mcpResponseJSON is today's only framing — a bare JSON document, + // Content-Type: application/json — and mcpNegotiateFormat's default + // for every request that does not explicitly ask for SSE. + mcpResponseJSON mcpResponseFormat = iota + // mcpResponseSSE frames the identical JSON-RPC envelope as one + // text/event-stream event: a single "data:" line carrying the + // compact JSON body, terminated by the blank line the SSE wire + // format requires (sseWriter.writeData, sse.go, reused as-is here — + // see writeJSONRPCEnvelope's own doc comment for why). + mcpResponseSSE +) + +// mcpNegotiateFormat decides handleMCPFederated's response framing from +// r's own Accept header — the ONLY signal that ever produces +// mcpResponseSSE. Every other case, including no Accept header at all, +// "application/json" alone, and "*/*", yields mcpResponseJSON: the +// byte-identical-to-today default this route must never change unasked. +// +// "*/*" is a deliberate part of that default, not an oversight: it +// states no SPECIFIC preference for either framing, and a caller whose +// HTTP client stack sets it automatically (a common default for generic +// tooling that never touches this header by hand) would otherwise see +// this route's output silently change shape it never asked to change — +// exactly the regression this negotiation must not cause. Only a caller +// that names the concrete "text/event-stream" token — the real, reported +// client's own "Accept: application/json, text/event-stream" — gets the +// new framing. +func mcpNegotiateFormat(r *http.Request) mcpResponseFormat { + if mcpAcceptsSSE(r) { + return mcpResponseSSE + } + return mcpResponseJSON +} + +// mcpAcceptsSSE reports whether any "Accept" header on r names +// text/event-stream as an acceptable response media type. Accept is a +// comma-separated list of media ranges, each optionally followed by +// ";q=..." or other parameters (RFC 9110 §12.5.1) — this checks every +// element of every "Accept" header LINE present (net/http folds repeated +// header instances into one Header entry under the same key, but a +// client may still send more than one "Accept:" line; Header.Values +// returns each separately, covering both shapes), stripping parameters +// and surrounding whitespace before an EXACT, case-insensitive +// comparison against "text/event-stream" — deliberately never a +// substring or prefix check, so a near-miss media type like +// "application/x-text/event-stream-foo" cannot false-positive the way a +// naive strings.Contains over the whole header would. +func mcpAcceptsSSE(r *http.Request) bool { + for _, header := range r.Header.Values("Accept") { + for _, part := range strings.Split(header, ",") { + mediaType, _, _ := strings.Cut(part, ";") + if strings.EqualFold(strings.TrimSpace(mediaType), "text/event-stream") { + return true + } + } + } + return false +} + // writeJSONRPCEnvelope writes resp as the response body, matching // setAdminJSONHeaders' own "declared Content-Type, nothing else asserted" -// 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, 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) { +// minimalism for its default JSON framing — 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, 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). +// +// format (mcpNegotiateFormat's own doc comment) selects between that +// default JSON framing and mcpResponseSSE: one "data:" line carrying resp +// as compact JSON, terminated by the blank line the SSE wire format +// requires. The SSE branch reuses sse.go's newSSEWriter/writeData +// verbatim rather than a second, hand-rolled encoder — writeData already +// writes exactly "data: " + b + "\n\n" as one Write call, which both IS +// the single-shot envelope this route needs and already gets the +// terminator right, the exact class of bug (a missing blank-line +// terminator) this negotiation exists to fix elsewhere, not repeat here. +// Its streaming-oriented extras — writeDone's "[DONE]" sentinel, +// per-event flushing meant for incremental delivery — are simply unused: +// this is one event, not a stream, so nothing here ever calls writeDone, +// and a Flush that is a no-op under Yaegi (newSSEWriter's own doc +// comment) costs this single-shot response nothing, unlike a real +// streamed reply that depends on it for time-to-first-byte. +func writeJSONRPCEnvelope(w http.ResponseWriter, format mcpResponseFormat, resp jsonrpcResponse) { resp.JSONRPC = jsonrpcVersion + if format == mcpResponseSSE { + sw := newSSEWriter(w) // commits SSE headers unconditionally, before the encode below — same header-then-body order as the JSON branch + b, err := json.Marshal(resp) + if err != nil { + return // headers already committed; nothing useful to do on encode failure + } + _ = sw.writeData(b) // headers already committed; nothing useful to do on write failure + return + } w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(resp) // headers already committed; nothing useful to do on encode failure } // writeJSONRPCResult marshals result and writes it as a successful -// JSON-RPC response under id. A marshal failure (result is not -// JSON-marshalable — never true for any value this file actually passes, -// but defensively handled rather than panicking) reports itself as a -// JSON-RPC internal error instead. -func writeJSONRPCResult(w http.ResponseWriter, id json.RawMessage, result any) { +// JSON-RPC response under id, framed per format (mcpNegotiateFormat). A +// marshal failure (result is not JSON-marshalable — never true for any +// value this file actually passes, but defensively handled rather than +// panicking) reports itself as a JSON-RPC internal error instead, in the +// same negotiated format. +func writeJSONRPCResult(w http.ResponseWriter, format mcpResponseFormat, id json.RawMessage, result any) { resultBytes, err := json.Marshal(result) if err != nil { - writeJSONRPCErrorResponse(w, id, jsonrpcInternalError, "failed to encode result") + writeJSONRPCErrorResponse(w, format, id, jsonrpcInternalError, "failed to encode result") return } - writeJSONRPCEnvelope(w, jsonrpcResponse{ID: id, Result: resultBytes}) + writeJSONRPCEnvelope(w, format, jsonrpcResponse{ID: id, Result: resultBytes}) } -// writeJSONRPCErrorResponse writes a JSON-RPC error response under id. -func writeJSONRPCErrorResponse(w http.ResponseWriter, id json.RawMessage, code int, message string) { - writeJSONRPCEnvelope(w, jsonrpcResponse{ID: id, Error: &jsonrpcError{Code: code, Message: message}}) +// writeJSONRPCErrorResponse writes a JSON-RPC error response under id, +// framed per format (mcpNegotiateFormat) — a client that asked for SSE +// gets its error in SSE framing too, never a JSON body it cannot parse. +func writeJSONRPCErrorResponse(w http.ResponseWriter, format mcpResponseFormat, id json.RawMessage, code int, message string) { + writeJSONRPCEnvelope(w, format, jsonrpcResponse{ID: id, Error: &jsonrpcError{Code: code, Message: message}}) } // isLegacySSETransportURL reports whether rawURL's path ends in "/sse" — @@ -326,6 +428,15 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u * return } + // Negotiated once, from the request alone, before the body is even + // read: every JSON-RPC-shaped response below — including the parse + // error a malformed body itself produces — must honor it + // (mcpNegotiateFormat's own doc comment). The transport-level error + // paths in this function (writeOAIError, writeLimitViolation) are + // deliberately NOT part of this — see writeJSONRPCEnvelope's own doc + // comment for why those stay out of scope. + format := mcpNegotiateFormat(r) + body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes)) if err != nil { writeOAIError(w, http.StatusBadRequest, "invalid_request_error", "cannot read request body") @@ -337,7 +448,7 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u * // A malformed envelope carries no reliable id at all — JSON-RPC // 2.0 §5.1 requires "id":null for exactly this case, not an // omitted key (jsonrpcResponse.ID's own doc comment). - writeJSONRPCErrorResponse(w, json.RawMessage("null"), jsonrpcParseError, "parse error") + writeJSONRPCErrorResponse(w, format, json.RawMessage("null"), jsonrpcParseError, "parse error") return } @@ -349,20 +460,23 @@ func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u * switch { case req.Method == "initialize": - g.mcpFederatedInitialize(w, req) + g.mcpFederatedInitialize(w, format, req) case req.Method == "ping": - writeJSONRPCResult(w, req.ID, map[string]any{}) + writeJSONRPCResult(w, format, req.ID, map[string]any{}) case req.Method == "tools/list": - g.mcpFederatedToolsList(w, r, req, allowedMCPServerNames(g.cfg, grp)) + g.mcpFederatedToolsList(w, format, r, req, allowedMCPServerNames(g.cfg, grp)) case req.Method == "tools/call": - g.mcpFederatedToolsCall(w, r, req, grp) + g.mcpFederatedToolsCall(w, format, r, req, grp) case strings.HasPrefix(req.Method, "notifications/"): // A JSON-RPC notification carries no id and gets no response body // by definition; MCP's Streamable HTTP transport answers a - // notification POST with 202 Accepted and nothing else. + // notification POST with 202 Accepted and nothing else — format + // plays no part here, deliberately: there is no body to frame + // either way, and wrapping an empty body in an SSE "data:" line + // would fabricate content this response was never meant to carry. w.WriteHeader(http.StatusAccepted) default: - writeJSONRPCErrorResponse(w, req.ID, jsonrpcMethodNotFound, "method not found: "+req.Method) + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcMethodNotFound, "method not found: "+req.Method) } } @@ -414,7 +528,7 @@ var knownMCPProtocolVersions = map[string]bool{ // 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) { +func (g *Gateway) mcpFederatedInitialize(w http.ResponseWriter, format mcpResponseFormat, req jsonrpcRequest) { var params mcpInitializeParams _ = json.Unmarshal(req.Params, ¶ms) // best-effort; empty/malformed params falls back to defaultMCPProtocolVersion below @@ -423,7 +537,7 @@ func (g *Gateway) mcpFederatedInitialize(w http.ResponseWriter, req jsonrpcReque protocolVersion = params.ProtocolVersion } - writeJSONRPCResult(w, req.ID, map[string]any{ + writeJSONRPCResult(w, format, req.ID, map[string]any{ "protocolVersion": protocolVersion, "capabilities": map[string]any{ "tools": map[string]any{}, @@ -795,7 +909,7 @@ const mcpFederatedFanoutConcurrency = 8 // single federated tools/list call can move several targets' own // counters, not 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, names []string) { +func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpResponseFormat, r *http.Request, req jsonrpcRequest, names []string) { var ( mu sync.Mutex wg sync.WaitGroup @@ -901,12 +1015,12 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, // 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") + writeJSONRPCErrorResponse(w, format, 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}) + writeJSONRPCResult(w, format, req.ID, mcpToolsListResult{Tools: merged}) } // mcpToolCallParams is one "tools/call" request's params: name is the @@ -956,16 +1070,16 @@ func resolveFederatedTool(fullName string, allowedNames []string) (serverName, t // 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) { +func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpResponseFormat, r *http.Request, req jsonrpcRequest, grp *group) { var params mcpToolCallParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { - writeJSONRPCErrorResponse(w, req.ID, jsonrpcInvalidParams, "invalid params") + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcInvalidParams, "invalid params") return } serverName, toolName, ok := resolveFederatedTool(params.Name, allowedMCPServerNames(g.cfg, grp)) if !ok { - writeJSONRPCErrorResponse(w, req.ID, jsonrpcInvalidParams, "unknown tool: "+params.Name) + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcInvalidParams, "unknown tool: "+params.Name) return } @@ -986,9 +1100,9 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request, // then-unparsable body would otherwise produce (security review // round 2, 2026-08-22, important finding 4). if errors.Is(err, errMCPResponseTooLarge) { - writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "response too large") + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcInternalError, "response too large") } else { - writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "upstream error") + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcInternalError, "upstream error") } return } @@ -1001,9 +1115,9 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request, // 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") + writeJSONRPCErrorResponse(w, format, req.ID, jsonrpcInternalError, "invalid upstream response") return } - writeJSONRPCEnvelope(w, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error}) + writeJSONRPCEnvelope(w, format, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error}) } diff --git a/mcp_federation_sse_test.go b/mcp_federation_sse_test.go new file mode 100644 index 0000000..f8b8e54 --- /dev/null +++ b/mcp_federation_sse_test.go @@ -0,0 +1,424 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// --- unit: mcpAcceptsSSE / mcpNegotiateFormat --- + +// TestMcpAcceptsSSE covers the tolerant-parsing requirements directly, +// without spinning up a full Gateway: q-values, surrounding whitespace, +// other media types sharing the list, case variation, a near-miss type +// that must NOT match by substring, multiple distinct "Accept:" header +// lines, and the no-header/empty case. +// +// MUTATION VERIFIED: temporarily changing the match to +// strings.Contains(strings.ToLower(header), "text/event-stream") (a +// naive whole-header substring check, the exact anti-pattern this +// function's own doc comment warns against) made the near-miss case +// below ("application/x-text/event-stream-foo" must NOT match) fail, +// since that string DOES contain "text/event-stream" as a substring. +// Reverted before committing. +func TestMcpAcceptsSSE(t *testing.T) { + cases := []struct { + name string + headers []string // one or more "Accept:" header LINES + want bool + }{ + {name: "no Accept header at all", headers: nil, want: false}, + {name: "empty Accept header", headers: []string{""}, want: false}, + {name: "application/json alone", headers: []string{"application/json"}, want: false}, + {name: "exact text/event-stream alone", headers: []string{"text/event-stream"}, want: true}, + { + name: "production reproduction: application/json, text/event-stream", + headers: []string{"application/json, text/event-stream"}, + want: true, + }, + {name: "different case", headers: []string{"TEXT/EVENT-STREAM"}, want: true}, + {name: "mixed case", headers: []string{"Text/Event-Stream"}, want: true}, + {name: "q-value parameter", headers: []string{"text/event-stream;q=0.9"}, want: true}, + {name: "q-value with space before it", headers: []string{"text/event-stream ;q=0.9"}, want: true}, + { + name: "q-values on both types, SSE not preferred but still acceptable", + headers: []string{"application/json;q=0.1, text/event-stream;q=0.9"}, + want: true, + }, + {name: "extra whitespace around the comma-separated element", headers: []string{"application/json , text/event-stream "}, want: true}, + {name: "wildcard alone does not count as an explicit SSE request", headers: []string{"*/*"}, want: false}, + { + name: "near-miss type must not substring-match", + headers: []string{"application/x-text/event-stream-foo"}, + want: false, + }, + { + name: "near-miss type alongside a real other type", + headers: []string{"application/json, application/x-text/event-stream-foo"}, + want: false, + }, + { + name: "two separate Accept header lines, second carries the real token", + headers: []string{"application/json", "text/event-stream"}, + want: true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, federatedMCPPath, nil) + for _, h := range tc.headers { + req.Header.Add("Accept", h) + } + if got := mcpAcceptsSSE(req); got != tc.want { + t.Errorf("mcpAcceptsSSE(Accept=%v) = %v, want %v", tc.headers, got, tc.want) + } + }) + } +} + +// TestMcpNegotiateFormat_StarStar_YieldsJSON pins the deliberate "*/*" +// decision at the negotiation layer (not just the lower-level parser +// above): a caller with no specific preference gets today's default, +// never a framing change it did not explicitly ask for. +// +// MUTATION VERIFIED: temporarily making mcpNegotiateFormat treat "*/*" +// as accepting SSE (return mcpResponseSSE whenever the header is +// non-empty) made this test fail, want mcpResponseJSON got +// mcpResponseSSE. Reverted before committing. +func TestMcpNegotiateFormat_StarStar_YieldsJSON(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, federatedMCPPath, nil) + req.Header.Set("Accept", "*/*") + if got := mcpNegotiateFormat(req); got != mcpResponseJSON { + t.Errorf("mcpNegotiateFormat(Accept: */*) = %v, want mcpResponseJSON — \"*/*\" states no specific preference and must not flip this route's default framing", got) + } +} + +// --- baseline: default framing is byte-identical to today's plain JSON --- +// +// These literal bytes were captured by running the CURRENT (pre-SSE- +// negotiation) handleMCPFederated against these exact requests before +// this file's own SSE-negotiation source changes existed, then pinned +// here verbatim ("assert against the actual current bytes", not merely +// against a re-parsed structural equivalent). Any change to the JSON +// path's byte output — reordering the envelope's fields, adding or +// dropping the trailing newline json.NewEncoder writes, a different +// Content-Type, or accidentally routing the JSON branch through the SSE +// writer — fails this test. +// +// MUTATION VERIFIED: temporarily forcing writeJSONRPCEnvelope's format +// check to `if true` (always taking the SSE branch, even for a plain +// mcpResponseJSON request) made every case below fail, since +// Content-Type became "text/event-stream" and the body gained a +// "data: " prefix and blank-line terminator. Reverted before committing. +func TestHandleMCPFederated_DefaultFraming_ByteIdenticalToBaseline(t *testing.T) { + cases := []struct { + name string + buildTest func(t *testing.T) (http.Handler, *http.Request) + accept string // "" means no Accept header at all + wantBody string + }{ + { + name: "ping, no Accept header", + wantBody: `{"jsonrpc":"2.0","result":{},"id":1}` + "\n", + buildTest: func(t *testing.T) (http.Handler, *http.Request) { + 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) + } + return h, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("1")}) + }, + }, + { + name: "malformed body parse error, Accept: application/json", + accept: "application/json", + wantBody: `{"error":{"message":"parse error","code":-32700},"jsonrpc":"2.0","id":null}` + "\n", + buildTest: func(t *testing.T) (http.Handler, *http.Request) { + 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) + } + req := httptest.NewRequest(http.MethodPost, federatedMCPPath, strings.NewReader("not json")) + req.Header.Set("Authorization", "Bearer sk-alice") + return h, req + }, + }, + { + name: "initialize default version, no Accept header", + wantBody: `{"jsonrpc":"2.0","result":{"capabilities":{"tools":{}},"protocolVersion":"2025-11-25","serverInfo":{"name":"traefik-llmgateway","version":"0.0.0-dev"}},"id":1}` + "\n", + buildTest: func(t *testing.T) (http.Handler, *http.Request) { + 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) + } + return h, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1")}) + }, + }, + { + name: "tools/list two-server aggregate, Accept: application/json", + accept: "application/json", + wantBody: `{"jsonrpc":"2.0","result":{"tools":[{"name":"alpha_lookup","description":"alpha's tool"},{"name":"beta_search","description":"beta's tool"}]},"id":42}` + "\n", + buildTest: func(t *testing.T) (http.Handler, *http.Request) { + alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup", Description: "alpha's tool"}}) + beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search", Description: "beta's tool"}}) + cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false) + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + return h, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("42")}) + }, + }, + { + name: "tools/call unknown prefix error, no Accept header", + wantBody: `{"error":{"message":"unknown tool: nosuch_tool","code":-32602},"jsonrpc":"2.0","id":9}` + "\n", + buildTest: func(t *testing.T) (http.Handler, *http.Request) { + 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(mcpToolCallParams{Name: "nosuch_tool"}) + return h, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("9"), Params: params}) + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + h, req := tc.buildTest(t) + if tc.accept != "" { + req.Header.Set("Accept", tc.accept) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if ct := rec.Header().Get("Content-Type"); ct != "application/json" { + t.Errorf("Content-Type = %q, want %q (default framing must not change)", ct, "application/json") + } + if rec.Body.String() != tc.wantBody { + t.Errorf("body = %q, want the exact pre-negotiation baseline byte sequence %q", rec.Body.String(), tc.wantBody) + } + }) + } +} + +// --- SSE: explicit request gets SSE framing, with a correct terminator --- + +// TestHandleMCPFederated_ProductionRepro_ToolsListAcceptEventStream is +// the literal reproduction from the incident report: a tools/list call +// carrying the client's own "Accept: application/json, text/event-stream" +// header must come back as SSE framing with a parseable "data:" line — +// not the plain JSON body that made the real client fail with "no data +// line in SSE response". +// +// MUTATION VERIFIED: reverting handleMCPFederated to call +// writeJSONRPCResult with a hardcoded mcpResponseJSON (ignoring the +// negotiated format entirely — the bug this whole change fixes) made +// this test fail: Content-Type stayed "application/json" and the body +// had no "data:" line. Reverted before committing. +func TestHandleMCPFederated_ProductionRepro_ToolsListAcceptEventStream(t *testing.T) { + alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup", Description: "alpha's tool"}}) + beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search", Description: "beta's tool"}}) + cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false) + 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("42")}) + req.Header.Set("Accept", "application/json, text/event-stream") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Fatalf("Content-Type = %q, want %q", ct, "text/event-stream") + } + if cc := rec.Header().Get("Cache-Control"); cc != "no-cache" { + t.Errorf("Cache-Control = %q, want %q", cc, "no-cache") + } + + body := rec.Body.String() + dataLine, ok := extractSingleSSEDataLine(t, body) + if !ok { + t.Fatalf("body has no parseable \"data:\" line: %q", body) + } + var got jsonrpcResponse + if err := json.Unmarshal([]byte(dataLine), &got); err != nil { + t.Fatalf("data line is not valid JSON-RPC: %v (line=%q)", err, dataLine) + } + if got.Error != nil { + t.Fatalf("error = %+v, want nil", got.Error) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + if len(result.Tools) != 2 { + t.Errorf("got %d tools, want 2 (alpha_lookup, beta_search)", len(result.Tools)) + } +} + +// TestHandleMCPFederated_SSEFraming_TerminatesWithBlankLine isolates the +// exact defect class this whole change exists to fix — a missing SSE +// event terminator — as its own assertion, independent of payload +// content: the body must end with the blank-line terminator ("\n\n" +// after the data line), and must contain EXACTLY one such terminator +// (not, say, a terminator buried mid-body with trailing junk after it, +// which would equally break a client's per-event framing). +// +// MUTATION VERIFIED: temporarily changing sseWriter.writeData (sse.go) to +// append a single '\n' instead of two ("buf = append(buf, '\n')") made +// this test fail — body no longer ended with "\n\n" — reproducing +// exactly the missing-terminator bug class the brief warns against. +// Reverted before committing. +func TestHandleMCPFederated_SSEFraming_TerminatesWithBlankLine(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) + } + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("1")}) + req.Header.Set("Accept", "text/event-stream") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + body := rec.Body.String() + if !strings.HasSuffix(body, "\n\n") { + t.Fatalf("body does not end with the SSE blank-line terminator: %q", body) + } + if !strings.HasPrefix(body, "data: ") { + t.Fatalf("body does not start with the \"data: \" field prefix: %q", body) + } + // Exactly one event: the terminator ("\n\n") must appear only once, + // at the very end — a stray earlier one would split this into two + // (or more) malformed events instead of the single one intended. + if n := strings.Count(body, "\n\n"); n != 1 { + t.Errorf("body contains %d blank-line terminators, want exactly 1: %q", n, body) + } +} + +// TestHandleMCPFederated_ErrorResponses_HonourNegotiatedSSEFraming +// covers the brief's explicit requirement: "A client that asked for SSE +// must not get a JSON error it cannot parse." Both a locally-produced +// error (parse error, before any backend is ever contacted) and one +// produced deeper in tools/call's own resolution logic get the same SSE +// treatment as a success response. +// +// MUTATION VERIFIED: temporarily hardcoding format = mcpResponseJSON at +// the top of writeJSONRPCErrorResponse (ignoring its own parameter) made +// both cases below fail — Content-Type stayed "application/json". +// Reverted before committing. +func TestHandleMCPFederated_ErrorResponses_HonourNegotiatedSSEFraming(t *testing.T) { + t.Run("parse error", 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) + } + req := httptest.NewRequest(http.MethodPost, federatedMCPPath, strings.NewReader("not json")) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Accept", "text/event-stream") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Fatalf("Content-Type = %q, want %q", ct, "text/event-stream") + } + dataLine, ok := extractSingleSSEDataLine(t, rec.Body.String()) + if !ok { + t.Fatalf("body has no parseable \"data:\" line: %q", rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal([]byte(dataLine), &got); err != nil { + t.Fatalf("data line is not valid JSON-RPC: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcParseError { + t.Errorf("error = %+v, want code %d", got.Error, jsonrpcParseError) + } + }) + + t.Run("tools/call unknown prefix", 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(mcpToolCallParams{Name: "nosuch_tool"}) + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("9"), Params: params}) + req.Header.Set("Accept", "application/json, text/event-stream") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if ct := rec.Header().Get("Content-Type"); ct != "text/event-stream" { + t.Fatalf("Content-Type = %q, want %q", ct, "text/event-stream") + } + dataLine, ok := extractSingleSSEDataLine(t, rec.Body.String()) + if !ok { + t.Fatalf("body has no parseable \"data:\" line: %q", rec.Body.String()) + } + var got jsonrpcResponse + if err := json.Unmarshal([]byte(dataLine), &got); err != nil { + t.Fatalf("data line is not valid JSON-RPC: %v", err) + } + if got.Error == nil || got.Error.Code != jsonrpcInvalidParams { + t.Errorf("error = %+v, want code %d", got.Error, jsonrpcInvalidParams) + } + }) +} + +// TestHandleMCPFederated_Notification_Returns202NoBody_EvenWithAcceptSSE +// pins the brief's other explicit warning: a notification's 202-with- +// no-body response must never get wrapped in an SSE frame, even when the +// caller's Accept header asks for one. +// +// MUTATION VERIFIED: temporarily replacing the notifications/* branch's +// `w.WriteHeader(http.StatusAccepted)` with +// `writeJSONRPCResult(w, format, nil, map[string]any{})` (wrapping it +// like every other response) made this test fail on both the status +// code changing to 200 and the body no longer being empty. Reverted +// before committing. +func TestHandleMCPFederated_Notification_Returns202NoBody_EvenWithAcceptSSE(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) + } + req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "notifications/initialized"}) + req.Header.Set("Accept", "text/event-stream") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusAccepted { + t.Fatalf("status = %d, want 202, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.Len() != 0 { + t.Errorf("body = %q, want empty — a notification must never be wrapped in an SSE frame", rec.Body.String()) + } + if ct := rec.Header().Get("Content-Type"); ct == "text/event-stream" { + t.Errorf("Content-Type = %q, must not be set to text/event-stream for a bodyless 202", ct) + } +} + +// extractSingleSSEDataLine finds the (single) "data:" line in an SSE +// body and returns its payload, reusing sse.go's own lastSSEDataLine so +// this test exercises the same parsing logic a well-behaved reader +// would, rather than a hand-rolled second parser that could disagree +// with it. +func extractSingleSSEDataLine(t *testing.T, body string) (string, bool) { + t.Helper() + data := lastSSEDataLine([]byte(body)) + if len(data) == 0 || string(data) == body { + return "", false + } + return string(data), true +} diff --git a/tools/yaegi-check/main.go b/tools/yaegi-check/main.go index 09d88c4..a9cf29f 100644 --- a/tools/yaegi-check/main.go +++ b/tools/yaegi-check/main.go @@ -1318,6 +1318,10 @@ func exerciseHandler(handler http.Handler, builtinContextTokens int, failoverBHi return err } + if err := exerciseFederatedSSENegotiation(handler); err != nil { + return err + } + if err := exerciseStampedVersion(handler); err != nil { return err } @@ -1625,6 +1629,82 @@ func exerciseFederatedTooLarge(handler http.Handler) error { return nil } +// exerciseFederatedSSENegotiation drives two real POST /mcp "ping" calls +// against the interpreted handler — one carrying no Accept header +// (today's plain JSON framing) and one carrying the reported client's +// exact "Accept: application/json, text/event-stream" header — proving +// content negotiation (mcpNegotiateFormat, mcpAcceptsSSE) and the SSE +// branch's reuse of sse.go's newSSEWriter/writeData both run correctly +// under the REAL interpreter. mcp_federation_sse_test.go already covers +// this compiled; this harness exists because a compiled pass on this +// codebase has repeatedly said nothing about the interpreted shape. +// +// "ping" is deliberately chosen over tools/list or tools/call: it is +// answered entirely locally (handleMCPFederated's own doc comment), so +// this probe needs no additional upstream server and cannot be confused +// with mcpProbeUpstream's own, deliberately oversized, tools/call-only +// response above. +// +// This is also the first time under this harness that sse.go's +// newSSEWriter/writeData actually executes interpreted at all. The +// streaming (chat completions) path is documented as losing its +// incremental Flush under Yaegi (newSSEWriter's own doc comment, +// confirmed by Task 15's integration suite) — harmless here, since this +// is a single-shot envelope that writes once and returns rather than a +// response that depends on Flush for time-to-first-byte, but it means no +// existing gate has ever run this writer interpreted before now. +func exerciseFederatedSSENegotiation(handler http.Handler) error { + pingBody := `{"jsonrpc":"2.0","id":1,"method":"ping"}` + + jsonReq := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(pingBody)) + jsonReq.Header.Set("Authorization", "Bearer "+testDataUserAPIKey) + jsonReq.Header.Set("Content-Type", "application/json") + jsonRec := httptest.NewRecorder() + handler.ServeHTTP(jsonRec, jsonReq) + if jsonRec.Code != http.StatusOK { + return fmt.Errorf("POST /mcp ping (no Accept header): status = %d, want 200, body=%s", jsonRec.Code, jsonRec.Body.String()) + } + if ct := jsonRec.Header().Get("Content-Type"); ct != "application/json" { + return fmt.Errorf("POST /mcp ping (no Accept header): Content-Type = %q, want %q — default framing must not change interpreted", ct, "application/json") + } + + sseReq := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(pingBody)) + sseReq.Header.Set("Authorization", "Bearer "+testDataUserAPIKey) + sseReq.Header.Set("Content-Type", "application/json") + sseReq.Header.Set("Accept", "application/json, text/event-stream") + sseRec := httptest.NewRecorder() + handler.ServeHTTP(sseRec, sseReq) + if sseRec.Code != http.StatusOK { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): status = %d, want 200, body=%s", sseRec.Code, sseRec.Body.String()) + } + if ct := sseRec.Header().Get("Content-Type"); ct != "text/event-stream" { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): Content-Type = %q, want %q — SSE negotiation did not run interpreted", ct, "text/event-stream") + } + + body := sseRec.Body.String() + if !strings.HasPrefix(body, "data: ") { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): body does not start with \"data: \": %q", body) + } + if !strings.HasSuffix(body, "\n\n") { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): body does not end with the SSE blank-line terminator — the exact missing-terminator defect class this negotiation was added to fix: %q", body) + } + + dataLine := strings.TrimSuffix(strings.TrimPrefix(body, "data: "), "\n\n") + var resp struct { + Result map[string]any `json:"result"` + ID json.RawMessage `json:"id"` + } + if err := json.Unmarshal([]byte(dataLine), &resp); err != nil { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): data line is not valid JSON-RPC: %w (line=%q)", err, dataLine) + } + if string(resp.ID) != "1" { + return fmt.Errorf("POST /mcp ping (Accept: text/event-stream): id = %s, want the client's own id 1 echoed back", resp.ID) + } + + fmt.Println("yaegi-check: federated MCP content negotiation served both JSON and SSE framing correctly, interpreted") + return nil +} + // exerciseFailover drives one real POST /v1/chat/completions against the // interpreted handler for failoverModelID — a bare model id both // "failover-a" (always 500s) and "failover-b" (always 200s) are