diff --git a/mcp_a2a.go b/mcp_a2a.go index 32df7e9..805adf7 100644 --- a/mcp_a2a.go +++ b/mcp_a2a.go @@ -6,6 +6,7 @@ import ( "net/http" "net/url" "sort" + "strings" "time" ) @@ -211,12 +212,14 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u } g.limiter.countTargetRequest(targetScopeKind(kind), name) - upstreamURL := targetURL - if rest != "" { - upstreamURL += "/" + rest - } - if r.URL.RawQuery != "" { - upstreamURL += "?" + r.URL.RawQuery + upstreamURL, err := buildUpstreamTargetURL(targetURL, rest, r.URL.RawQuery) + if err != nil { + // targetURL is operator-configured and already validated at + // construction (validateTargetURLs) — unreachable in practice, + // kept as an explicit guard rather than trusting that invariant + // silently. + writeOAIError(w, http.StatusInternalServerError, "invalid_request_error", "invalid target configuration") + return } // feat/target-health: passive recording (target_health.go) — latency @@ -229,6 +232,70 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u g.recordTargetProxyHealth(kind, name, targetURL, result, ok, time.Since(start)) } +// buildUpstreamTargetURL joins targetURL (an MCP server or A2A agent's own +// configured base, already validated at construction by +// validateTargetURL) with rest (handleTargetProxy's own path remainder, +// already checked for a traversal segment by its caller) using +// url.URL.JoinPath rather than plain string concatenation (F7, review +// round 3, 2026-09): the prior "+ \"/\" + rest" produced a doubled slash +// whenever targetURL's own path already ended in "/" (e.g. configured as +// "http://x/mcp/"), and, worse, a targetURL carrying its own query string +// — the only way TargetConfig/AgentConfig can pass a credential to an +// upstream today, since neither carries a header field — got that query +// glued onto the wrong side of "?" + the request's own RawQuery, +// producing something like "...?token=abc/foo" or two "?" characters and +// corrupting the token federation was relying on. The target's own query +// string is preserved (JoinPath never touches RawQuery) and merged with +// the incoming request's own query (mergeQueryStrings, below) — neither +// silently overwrites the other. +func buildUpstreamTargetURL(targetURL, rest, requestRawQuery string) (string, error) { + base, err := url.Parse(targetURL) + if err != nil { + return "", err + } + joined := base + if rest != "" { + joined = base.JoinPath(rest) + } + joined.RawQuery = mergeQueryStrings(joined.RawQuery, requestRawQuery) + return joined.String(), nil +} + +// mergeQueryStrings combines targetQuery (the configured target URL's own +// query string, returned verbatim when it stands alone) with +// requestQuery (the incoming client request's own), so neither silently +// discards the other (F7). A request query that fails to parse is +// dropped, not treated as a whole-request failure: the request's own +// method/body may still be perfectly valid, and the target's own query — +// a credential, per this file's own TargetConfig/AgentConfig doc +// comments — is the one half of this that must never be lost silently. +func mergeQueryStrings(targetQuery, requestQuery string) string { + if targetQuery == "" { + return requestQuery + } + if requestQuery == "" { + return targetQuery + } + merged, err := url.ParseQuery(targetQuery) + if err != nil { + // targetQuery came from an already-parsed url.URL's own RawQuery, + // so this is unreachable in practice — falls back to the + // request's own query alone rather than panicking or silently + // dropping both. + return requestQuery + } + reqValues, err := url.ParseQuery(requestQuery) + if err != nil { + return targetQuery + } + for k, vs := range reqValues { + for _, v := range vs { + merged.Add(k, v) + } + } + return merged.Encode() +} + // validateTargetURLs checks every configured MCP-server and agent entry at // construction: the map value must be non-nil, the name must pass // validateConfigName (the same character-set, dot-only, and reserved-name @@ -261,6 +328,17 @@ func validateTargetURLs(cfg *Config) error { if err := validateTargetURL(ac.URL); err != nil { return fmt.Errorf("llmgateway: agents %q: %w", name, err) } + // F8, review round 3, 2026-09: handleAgents builds the card URL as + // "/a2a/" + name + cardPath (its own doc comment) with no + // separator of its own — a Card that does not start with "/" + // (e.g. "agent.json") concatenates straight onto the agent name + // instead of starting a new path segment, producing something + // like "/a2a/myagentagent.json": a 404, or a completely different + // route, never the card the operator configured. Empty (the + // default, defaultAgentCardPath applies) is always valid. + if ac.Card != "" && !strings.HasPrefix(ac.Card, "/") { + return fmt.Errorf("llmgateway: agents %q: card %q must start with \"/\"", name, ac.Card) + } } return nil } diff --git a/mcp_a2a_test.go b/mcp_a2a_test.go index ab2f1c2..ee3f53f 100644 --- a/mcp_a2a_test.go +++ b/mcp_a2a_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "net/url" "strings" "testing" "time" @@ -544,6 +545,112 @@ func TestHandleTargetProxy_EmptyRest_ProxiesToRoot(t *testing.T) { } } +// --- F7 (review round 3, 2026-09): URL joining, not string concatenation --- + +// TestBuildUpstreamTargetURL covers buildUpstreamTargetURL directly: +// good (ordinary join), bad (a target whose own path already ends in +// "/", the doubled-slash bug), and edge (a target carrying its own query +// string, which must survive alongside the request's own). +func TestBuildUpstreamTargetURL(t *testing.T) { + cases := []struct { + name string + targetURL string + rest string + requestRawQuery string + want string + }{ + { + name: "ordinary join", + targetURL: "http://x.internal/mcp", + rest: "foo", + want: "http://x.internal/mcp/foo", + }, + { + name: "target path already ends in / no longer doubles the slash", + targetURL: "http://x.internal/mcp/", + rest: "foo", + want: "http://x.internal/mcp/foo", + }, + { + name: "empty rest proxies to the target root unchanged", + targetURL: "http://x.internal/mcp", + rest: "", + want: "http://x.internal/mcp", + }, + { + name: "target's own query survives alongside the request's own", + targetURL: "http://x.internal/mcp?token=abc", + rest: "foo", + requestRawQuery: "q=1", + want: "http://x.internal/mcp/foo?q=1&token=abc", + }, + { + name: "target's own query survives with no request query at all", + targetURL: "http://x.internal/mcp?token=abc", + rest: "foo", + want: "http://x.internal/mcp/foo?token=abc", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := buildUpstreamTargetURL(tc.targetURL, tc.rest, tc.requestRawQuery) + if err != nil { + t.Fatalf("buildUpstreamTargetURL: %v", err) + } + if got != tc.want { + t.Errorf("buildUpstreamTargetURL(%q, %q, %q) = %q, want %q", tc.targetURL, tc.rest, tc.requestRawQuery, got, tc.want) + } + }) + } +} + +// TestHandleTargetProxy_TargetQueryString_MergedWithRequestQuery is the +// F7 end-to-end regression test: a target URL configured with its own +// query string (the only way TargetConfig can carry a credential today) +// must reach the upstream intact, merged with — not overwritten or +// corrupted by — the caller's own request query. +func TestHandleTargetProxy_TargetQueryString_MergedWithRequestQuery(t *testing.T) { + var gotURL string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotURL = r.URL.String() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL + "/mcp?token=secret"}} + 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) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools?q=1", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + 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()) + } + u, err := url.Parse(gotURL) + if err != nil { + t.Fatalf("parse upstream-observed URL %q: %v", gotURL, err) + } + if u.Path != "/mcp/tools" { + t.Errorf("upstream path = %q, want %q", u.Path, "/mcp/tools") + } + got := u.Query() + if got.Get("token") != "secret" { + t.Errorf("upstream query token = %q, want %q (the target's own configured query must survive)", got.Get("token"), "secret") + } + if got.Get("q") != "1" { + t.Errorf("upstream query q = %q, want %q (the request's own query must survive too)", got.Get("q"), "1") + } +} + // --- dead upstream maps to 502, same as native passthrough --- func TestHandleTargetProxy_DeadUpstream_Returns502(t *testing.T) { @@ -815,6 +922,26 @@ func TestNewGateway_TargetURLValidation(t *testing.T) { mcpServers: map[string]*TargetConfig{"has a space": {URL: "http://mcp-alpha.internal:8080"}}, wantErr: true, }, + { + // F8, review round 3, 2026-09: handleAgents builds the card + // URL as "/a2a/" + name + cardPath with no separator of its + // own — a Card that does not start with "/" concatenates + // straight onto the agent name instead of starting a new + // path segment. + name: "agent card not starting with / is a constructor error", + agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: "agent.json"}}, + wantErr: true, + }, + { + name: "agent card starting with / is valid", + agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: "/custom-card.json"}}, + wantErr: false, + }, + { + name: "agent card empty (default applies) is valid", + agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal", Card: ""}}, + wantErr: false, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/mcp_federation.go b/mcp_federation.go index 17fe2f6..4de827f 100644 --- a/mcp_federation.go +++ b/mcp_federation.go @@ -335,12 +335,21 @@ func writeJSONRPCErrorResponse(w http.ResponseWriter, format mcpResponseFormat, // 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). +// +// The trailing slash is trimmed before the suffix check (F9, review round +// 3, 2026-09): a server configured as "http://host/sse/" is the identical +// legacy-SSE endpoint as "http://host/sse" — trailing slashes are common +// in operator-typed URLs and carry no transport meaning here — but the +// bare HasSuffix check missed it, leaving such a server IN the fan-out +// where every tools/list call wasted a semaphore slot, failed, and marked +// the target unhealthy for a transport this file was never built to +// speak at all. func isLegacySSETransportURL(rawURL string) bool { u, err := url.Parse(rawURL) if err != nil { return false } - return strings.HasSuffix(u.Path, "/sse") + return strings.HasSuffix(strings.TrimSuffix(u.Path, "/"), "/sse") } // allowedMCPServerNames returns the sorted names of every MCP server BOTH @@ -549,9 +558,13 @@ func (g *Gateway) mcpFederatedInitialize(w http.ResponseWriter, format mcpRespon }) } -// mcpTool is one tool entry in an MCP "tools/list" result, either as read -// back from a backend server's own response or as re-emitted (name -// prefixed) in this gateway's federated aggregate. +// mcpTool is a convenience literal type this repository's own tests use +// to build a canned backend tools/list response (newMockJSONRPCServer, +// sessionRequiredMockServer) — name/description/inputSchema are the three +// fields every such test fixture actually needs to set. Production code +// no longer round-trips a real tool through this fixed-field shape (F2, +// review round 3, 2026-09) — see mcpToolsListResult's own doc comment for +// why. type mcpTool struct { Name string `json:"name"` Description string `json:"description,omitempty"` @@ -559,9 +572,118 @@ type mcpTool struct { } // mcpToolsListResult is the "result" object of an MCP "tools/list" -// response, both a backend's own and this gateway's federated one. +// response, both a backend's own (decoded, in mcpFederatedToolsList's own +// per-server fan-out) and this gateway's federated one (re-encoded from +// the merged tools, mcpFederatedToolsList's own final writeJSONRPCResult +// call). +// +// Tools decodes as []map[string]json.RawMessage, not the fixed mcpTool +// struct above (F2, review round 3, 2026-09): the earlier fixed-field +// struct kept only name/description/inputSchema and silently dropped +// every OTHER field a real tool object carries on the round trip through +// this gateway — title, annotations (destructiveHint, readOnlyHint, ...), +// outputSchema, _meta. A federated client that relies on annotations to +// decide whether a tool needs user confirmation before a call, or on +// outputSchema to validate structuredContent, silently lost that signal. +// Decoding into a raw map per tool and rewriting only "name" +// (prefixMCPTool, below) passes every other field through unchanged, +// byte-for-byte, exactly as the backend sent it. One side effect: a Go +// map's keys marshal back out in SORTED order (encoding/json's own +// documented behavior), not the backend's original field order — a +// cosmetic reordering, never a content change. +// +// NextCursor is the MCP pagination cursor a backend sets when it has more +// tools than fit in one response (F3, review round 3, 2026-09) — +// mcpFederatedToolsList's own per-server loop follows it, bounded by +// mcpToolsListMaxPagesPerServer, until a page comes back with none. type mcpToolsListResult struct { - Tools []mcpTool `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` + Tools []map[string]json.RawMessage `json:"tools"` +} + +// mcpToolsListParams is one "tools/list" request's own params. Cursor is +// omitted (its zero value, "") for a server's first page; a non-empty +// value echoes back exactly the nextCursor that same server's own +// previous page reported (F3) — never a cursor minted by this gateway or +// borrowed from another server, since MCP's cursor is an opaque, +// per-server token. +type mcpToolsListParams struct { + Cursor string `json:"cursor,omitempty"` +} + +// mcpToolsListMaxPagesPerServer bounds how many pages +// mcpFederatedToolsList will follow via nextCursor for any ONE server +// (F3, review round 3, 2026-09) — independent of, and in addition to, +// the fan-out's own shared toolsListBackendTimeout budget (fanoutCtx in +// mcpFederatedToolsList), which already bounds total wall-clock +// regardless of page count. This cap exists so a misbehaving backend +// that always sets nextCursor (an unbounded pagination loop) cannot hold +// this goroutine's semaphore slot for the fan-out's entire timeout +// budget one page at a time. 20 is chosen generously above any real MCP +// server's tool catalog this codebase has seen (the largest production +// tools/list this gateway federates today is a single, unpaginated +// page), so it is never expected to actually trim a legitimate server's +// results — reaching it degrades to a partial tool list for that one +// server plus a logged note, the same "partial, not a hard failure" +// shape the fan-out's own timeout budget already uses. +const mcpToolsListMaxPagesPerServer = 20 + +// prefixMCPTool returns a shallow copy of tool (F2, review round 3, +// 2026-09) with only its "name" field rewritten to +// "_" — the merge's own established +// server-prefix convention (mcpFederatedToolsList's own doc comment) — +// and every other field passed through unchanged. fullName is that same +// rewritten name, returned separately so a caller building an +// mcpMergedTool for the final sort/dedupe (F4) never has to decode it back +// out of raw. ok is false when tool carries no "name" field at all, "name" +// is not a JSON string, or "name" is null or the empty string (P3, review +// round 4: json.Unmarshal("null", &toolName) succeeds and leaves toolName +// as "", indistinguishable from an explicit "name":"" without this check) +// — none a valid MCP tool object, skipped by the caller like any other +// malformed entry from that server, without failing the whole server's +// contribution. +func prefixMCPTool(serverName string, tool map[string]json.RawMessage) (out map[string]json.RawMessage, fullName string, ok bool) { + rawName, present := tool["name"] + if !present { + return nil, "", false + } + var toolName string + if err := json.Unmarshal(rawName, &toolName); err != nil { + return nil, "", false + } + // P3 (review round 4): "name":null unmarshals into toolName's zero + // value ("") with no error, same as an explicit "name":"" — either + // way there is no real tool name to prefix, so this would otherwise + // mint and list a bare "_" tool nothing can ever call by + // a meaningful name. Treated exactly like the missing-field and + // non-string cases above: skipped, not emitted. + if toolName == "" { + return nil, "", false + } + fullName = serverName + "_" + toolName + prefixedRaw, err := json.Marshal(fullName) + if err != nil { + return nil, "", false + } + out = make(map[string]json.RawMessage, len(tool)) + for k, v := range tool { + out[k] = v + } + out["name"] = json.RawMessage(prefixedRaw) + return out, fullName, true +} + +// mcpMergedTool pairs one server's already-prefixed tool object (raw map, +// F2) with the two fields mcpFederatedToolsList's own final sort and +// collision check (F4, review round 3, 2026-09) need without re-decoding +// "name" out of raw on every comparison: its own prefixed name, and the +// server it came from — the deterministic tiebreak the collision check +// below uses, so the outcome never depends on which goroutine's fan-out +// happened to finish first. +type mcpMergedTool struct { + raw map[string]json.RawMessage + name string + server string } // doBackendJSONRPC issues one JSON-RPC 2.0 POST to targetURL — an MCP @@ -708,18 +830,115 @@ func (g *Gateway) mcpBackendCloseSession(ctx context.Context, targetURL, session resp, err := g.targetClient.Do(req) if err != nil { - g.logf("federated mcp: best-effort session close failed for %q: %v", targetURL, err) + // P4 (review round 4): targetURL is this gateway's own credential + // channel (buildUpstreamTargetURL's doc comment, mcp_a2a.go), and + // err here is commonly a *url.Error wrapping it verbatim — logged + // through the same sanitizeBaseURL/sanitizeProviderErr pair every + // other target-URL-bearing log line in this codebase uses (e.g. + // adminTargetView.URL, admin.go). + g.logf("federated mcp: best-effort session close failed for %q: %v", sanitizeBaseURL(targetURL), sanitizeProviderErr(err.Error(), targetURL)) return } _ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure } +// mcpBackendSendInitializedNotification best-effort POSTs the MCP +// lifecycle's own "notifications/initialized" to targetURL, carrying +// sessionID (mcpBackendHandshake's own Mcp-Session-Id, when the backend +// set one) as the session header — the client-side half of the +// initialize handshake the MCP spec requires before any OTHER request on +// that session (F5, review round 3, 2026-09). mcpBackendCall's own +// handshake fallback previously sent "initialize" and went straight to +// the retried real call, skipping this notification entirely; a server +// that also gates its OTHER methods on having received it (not just on +// having answered "initialize") would still reject the retry. +// +// A notification carries no "id" (jsonrpcRequest's own zero value already +// omits it, ID being omitempty) and gets no JSON-RPC response body to +// interpret — the MCP Streamable HTTP transport answers a notification +// POST with 202 and nothing else — so there is nothing here to act on +// besides logging a transport-level failure; never surfaced to the +// caller, matching mcpBackendCloseSession's own best-effort convention +// right above. ctx is the same context the handshake and retry share: +// this call must not outlive the operation it is a step of, unlike the +// deliberately-separate mcpSessionCloseTimeout the best-effort DELETE +// runs under afterward (that one runs only once the real response has +// already been read and returned, so it alone needs its own budget). +func (g *Gateway) mcpBackendSendInitializedNotification(ctx context.Context, targetURL, sessionID string) { + body, err := json.Marshal(jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "notifications/initialized"}) + if err != nil { + return + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) //nolint:gosec // same operator-configured target URL as doBackendJSONRPC + if err != nil { + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + req.Header.Set(mcpSessionHeader, sessionID) + } + + resp, err := g.targetClient.Do(req) + if err != nil { + // P4 (review round 4): same targetURL-is-a-credential-channel + // scrub as mcpBackendCloseSession's identical log line above. + g.logf("federated mcp: best-effort notifications/initialized failed for %q: %v", sanitizeBaseURL(targetURL), sanitizeProviderErr(err.Error(), targetURL)) + return + } + _ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure +} + +// mcpSessionRequiredSignal decides whether the bare (session-less) +// attempt's own outcome — status, the parsed response when there is one, +// and err — specifically signals "this backend requires a session", the +// ONLY case mcpBackendCall's handshake-fallback retry may fire for (F1, +// review round 3, 2026-09). The PRIOR rule ("any JSON-RPC error OR any +// HTTP 4xx") mistook two different things for a session problem: a +// genuine tool-level failure a backend answers as an ordinary JSON-RPC +// error — retrying THAT silently double-executes whatever side effect a +// non-idempotent tools/call already ran, since the backend already did +// the work before reporting its own error — and a 429 (rate-limited, +// the opposite of what warrants a retry: it adds more load to a backend +// that already asked to be slowed down). +// +// method distinguishes the two call shapes mcpBackendCall is ever used +// for. tools/list's own params are the fixed, always-valid struct{}{} +// (mcpFederatedToolsList) and it has no side effect to double-run, so ANY +// JSON-RPC error there is still treated as session-missing — this is the +// exact behavior verified live against readitall ("method tools/list is +// invalid during...") and fetch (-32602 with or without params), neither +// of which literally names "session" in its own wording (mcpBackendCall's +// own doc comment covers the full live probe). tools/call gets the +// narrow rule instead: a JSON-RPC error must itself mention "session" +// (case-insensitive) to retry — any OTHER JSON-RPC error there is +// presumed to be the tool's own real failure, not a protocol-state +// problem. +// +// HTTP 400 and 404 are the only statuses ever eligible, for either +// method — narrowed from the prior "any 400-499", which is what let a +// 429 through; 401/403/409/422/etc. are real authorization/validation +// outcomes a retry cannot fix and must not paper over. +func mcpSessionRequiredSignal(method string, status int, resp *jsonrpcResponse, err error) bool { + if status == http.StatusBadRequest || status == http.StatusNotFound { + return true + } + if err != nil || resp == nil || resp.Error == nil { + return false + } + if method == "tools/list" { + return true + } + return strings.Contains(strings.ToLower(resp.Error.Message), "session") +} + // 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. +// ONE-SHOT initialize/retry/close handshake ONLY when the bare attempt's +// own outcome narrowly signals a missing session (mcpSessionRequiredSignal, +// its own doc comment covers exactly what qualifies and why it was +// narrowed, F1). // // 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, @@ -754,9 +973,16 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, return resp, nil } - retryEligible := (err == nil && resp.Error != nil) || (status >= 400 && status < 500) - if !retryEligible { - return nil, err + if !mcpSessionRequiredSignal(method, status, resp, err) { + // Not a session-required signal: relay the bare attempt's own + // outcome as-is — a JSON-RPC error response when the backend gave + // one (err is nil in that case; resp carries it), otherwise the + // transport-level err — and never retry. Retrying here is exactly + // what F1 fixes: a tools/call retry against an arbitrary tool-level + // error could silently double-execute a non-idempotent tool, and a + // 429 retry only adds load to a backend that is already + // rate-limiting. + return resp, err } sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL, maxBytes) @@ -772,6 +998,16 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, } return resp, nil } + // F5, review round 3, 2026-09: the MCP lifecycle requires the client + // to send notifications/initialized once it has the initialize + // result, before any other request — a step this handshake fallback + // previously skipped entirely, going straight from "initialize" to + // the retried real call. Best-effort, matching mcpBackendCloseSession's + // own convention right below: a notification gets no JSON-RPC + // response to interpret (a 202 and nothing else), so a failure here + // is logged, never surfaced — the retry that follows is what actually + // decides whether the handshake as a whole succeeded. + g.mcpBackendSendInitializedNotification(ctx, targetURL, sessionID) if sessionID != "" { defer func() { closeCtx, cancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout) @@ -821,18 +1057,29 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err return &out, nil } -// lastSSEDataLine returns the payload of the last "data:" line in an SSE -// body — body itself, unchanged, if it contains no such line at all (so -// 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). +// lastSSEDataLine returns the payload of the last COMPLETE SSE event's +// "data:" field in an SSE body — body itself, unchanged, if it contains +// no "data:" line at all (so 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). +// +// An SSE event's own multi-line "data:" field is joined with "\n" per the +// wire format (F6, review round 3, 2026-09; this gateway's own +// sseWriter.writeData, sse.go, emits exactly this shape for a multi-line +// payload) — every "data:" line up to the event's +// terminating blank line ("\n\n") is one logical field, not a separate +// one. A PRIOR revision of this function took only the single, textually +// LAST "data:" line in the whole body, which misread a backend's own +// pretty-printed or line-split JSON payload (each physical line +// individually invalid JSON on its own) as a parse failure, misdiagnosing +// a perfectly well-formed backend as failed. // // This does NOT implement per-request id correlation across multiple -// concurrent SSE-delivered messages the way a real MCP client's +// concurrent SSE-delivered EVENTS 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, +// takes the LAST complete event, 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 @@ -840,19 +1087,23 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err // 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")) { - line = bytes.TrimRight(line, "\r") - if data, ok := bytes.CutPrefix(line, []byte("data: ")); ok { - last = data - } else if data, ok := bytes.CutPrefix(line, []byte("data:")); ok { - last = bytes.TrimSpace(data) + normalized := bytes.ReplaceAll(body, []byte("\r\n"), []byte("\n")) + events := bytes.Split(normalized, []byte("\n\n")) + for i := len(events) - 1; i >= 0; i-- { + var dataLines [][]byte + for _, line := range bytes.Split(events[i], []byte("\n")) { + line = bytes.TrimRight(line, "\r") + if data, ok := bytes.CutPrefix(line, []byte("data: ")); ok { + dataLines = append(dataLines, data) + } else if data, ok := bytes.CutPrefix(line, []byte("data:")); ok { + dataLines = append(dataLines, bytes.TrimSpace(data)) + } + } + if dataLines != nil { + return bytes.Join(dataLines, []byte("\n")) } } - if last == nil { - return body - } - return last + return body } // mcpFederatedFanoutConcurrency bounds how many of names' backend servers @@ -883,38 +1134,72 @@ func lastSSEDataLine(body []byte) []byte { const mcpFederatedFanoutConcurrency = 8 // mcpFederatedToolsList implements the federated "tools/list": it fans -// out one mcpBackendCall per server in names (handleMCPFederated's own -// allowedMCPServerNames(grp) call, made once there so the exact set +// out one per-server fetch (below) per server in names (handleMCPFederated's +// own allowedMCPServerNames(grp) call, made once there so the exact set // contacted is always the one the client actually asked to reach), at // most mcpFederatedFanoutConcurrency at a time — every backend sharing // ONE overall toolsListBackendTimeout deadline (fanoutCtx, below) — 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, panics, 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. +// each tool's name prefixed "_" (prefixMCPTool; F2's +// own doc comment covers why every OTHER field of the tool object passes +// through unchanged), 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). Per server, a +// paginated backend is followed via nextCursor up to +// mcpToolsListMaxPagesPerServer pages (F3, review round 3, 2026-09) — all +// still inside that one server's share of fanoutCtx's shared budget, never +// a fresh timeout per page. A server that errors, times out, panics, 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. // -// 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, a connection refusal, or a panic) 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 — unlike handleTargetProxy's always-exactly- -// one-target shape (mcp_a2a.go). +// countTargetRequests (limits.go) attributes one request to every server +// actually DIALED here, in ONE batched call after wg.Wait() — "dialed", +// not "attempted" in the looser sense the prior revision used (F10, +// review round 3, 2026-09): a goroutine that was still waiting on the +// fan-out's own semaphore when fanoutCtx expired never sent a byte to +// its server and must not move that server's own per-target counter — +// only a server this gateway actually opened a connection to (and got a +// response, a timeout, a connection refusal, or a panic from) spent a +// real slot of this gateway's outbound capacity, which is what these +// counters exist to track. A 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). +// +// Two servers can legitimately mint the identical final prefixed name — +// e.g. server "a" has tool "b_c" and server "a_b" has tool "c", both +// producing "a_b_c" (F4, review round 3, 2026-09) — resolveFederatedTool's +// own longest-prefix rule can only ever route such a name to ONE of them, +// so shipping both in tools/list would advertise a tool the client can +// never actually reach as advertised, and some client-side LLM tool APIs +// reject a duplicate name outright. The final merge below detects this by +// sorting on (name, server) — for a stable, deterministic grouping that +// never depends on which goroutine's fan-out happened to finish first — +// then, for every colliding name, keeps the entry whose server is the one +// resolveFederatedTool would ACTUALLY route a tools/call for that name to +// (P2, review round 4: the prior revision instead kept whichever server +// sorted first alphabetically, which is a different server whenever the +// alphabetically-first name is not also the longest-prefix match — +// advertising a schema/description tools/list never lets the caller +// reach, because every call for that name is routed elsewhere). Logs +// exactly which server's tool was dropped and which server tools/call +// actually routes this name to. func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpResponseFormat, r *http.Request, req jsonrpcRequest, names []string) { var ( mu sync.Mutex wg sync.WaitGroup - merged = make([]mcpTool, 0, len(names)) + merged = make([]mcpMergedTool, 0, len(names)) failed []string + // dialed is the F10 fix's own accounting list — every name a + // goroutine actually won the semaphore for (and so is about to + // call mcpBackendCall for), kept separate from failed: a server + // starved out by fanoutCtx before ever winning the semaphore + // belongs in failed (the client-visible degradation log) but NOT + // in dialed (the per-target request counters below). + dialed []string ) // fanoutCtx bounds the WHOLE fan-out to toolsListBackendTimeout from // when THIS call started — not each individual backend's own start @@ -992,18 +1277,81 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons } defer func() { <-sem }() + // F10: this goroutine won the semaphore, so it is about to + // actually dial the server — recorded now, unconditionally, + // regardless of what the paginated fetch below returns, so + // countTargetRequests below reflects real outbound requests + // sent, never a server that only ever waited on the + // semaphore. + mu.Lock() + dialed = append(dialed, name) + mu.Unlock() + targetURL := g.cfg.MCPServers[name].URL - // feat/target-health: passive recording (target_health.go) — - // err != nil is a failure; a JSON-RPC-level resp.Error is NOT - // (the server answered, it just reported its own error). The - // record call right below already covers both: it reads ok - // from err == nil alone, so a non-nil resp.Error alongside a - // nil err still records success here — the failed-list - // append further down (used only for the merged tools/list - // result and the loud-failure short-circuit) is a separate - // concern and never touches targetHealth itself. probeStart := time.Now() - resp, err := g.mcpBackendCall(fanoutCtx, targetURL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes) + + // F3: follow nextCursor up to mcpToolsListMaxPagesPerServer + // pages, all sharing fanoutCtx's one already-ticking budget — + // never a fresh per-page timeout. transportErr and appFailed + // are kept separate deliberately: transportErr is a real + // mcpBackendCall failure (network, non-2xx, the size cap) and + // is what targetHealth below treats as unhealthy; appFailed + // is "the backend answered but this fan-out cannot use what + // it said" (a JSON-RPC error, or a page that fails to + // unmarshal) — the server is still reachable and healthy, it + // just has nothing usable to contribute this call, matching + // the pre-F3 rule that a JSON-RPC-level resp.Error never + // marks a server unhealthy. + var ( + pageTools []map[string]json.RawMessage + cursor string + transportErr error + appFailed bool + // appFailedErr names WHY appFailed was set (a JSON-RPC + // error object carries no Go error to reuse), so the F3 + // partial-keep log below can say something more useful + // than "it failed". No wrapped url — a JSON-RPC error + // message and a json.Unmarshal error never carry + // targetURL, unlike transportErr, so this never needs the + // sanitizeProviderErr treatment transportErr gets below. + appFailedErr error + ) + page := 0 + for ; page < mcpToolsListMaxPagesPerServer; page++ { + resp, err := g.mcpBackendCall(fanoutCtx, targetURL, "tools/list", mcpToolsListParams{Cursor: cursor}, mcpBackendResponseMaxBytes) + if err != nil { + transportErr = err + break + } + if resp.Error != nil { + appFailed = true + appFailedErr = fmt.Errorf("%s (code %d)", resp.Error.Message, resp.Error.Code) + break + } + var result mcpToolsListResult + if err := json.Unmarshal(resp.Result, &result); err != nil { + appFailed = true + appFailedErr = err + break + } + pageTools = append(pageTools, result.Tools...) + if result.NextCursor == "" { + cursor = "" + break + } + cursor = result.NextCursor + } + // F3 (review round 4): the loop above has no other way to + // tell "stopped because the server said no more pages" + // (cursor == "" here) from "stopped because + // mcpToolsListMaxPagesPerServer ran out while the server's + // last fetched page still set a NextCursor" — the latter + // silently truncates a legitimately larger catalog unless + // logged. + if page == mcpToolsListMaxPagesPerServer && cursor != "" { + g.logf("federated tools/list: server %q hit the %d-page cap; results truncated", name, mcpToolsListMaxPagesPerServer) + } + // G4: this goroutine WON the select above, so it was actually // dialed — only a client hang-up (context.Canceled, // propagating from r.Context() into fanoutCtx) says nothing @@ -1013,34 +1361,58 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons // shared toolsListBackendTimeout fired while THIS dialed // server was still in flight — it did not answer within the // fan-out's own budget, which is a real failure worth - // recording, not a client artifact (F2 over-suppressed this - // case; a genuinely hung server used to read "unknown" - // forever under federation-only traffic). - if err == nil || !errors.Is(err, context.Canceled) { - g.targetHealth.record(targetKindMCP, name, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic) + // recording, not a client artifact. Measured across every + // page this server's own loop actually ran, not just the + // first: a server that answers its first few pages fine and + // then times out on a later one is genuinely unhealthy, not + // a client artifact. + if transportErr == nil || !errors.Is(transportErr, context.Canceled) { + g.targetHealth.record(targetKindMCP, name, targetURL, transportErr == nil, transportErr, time.Since(probeStart), targetHealthSourceTraffic) } mu.Lock() defer mu.Unlock() - if err != nil || resp.Error != nil { + // F3 (review round 4): a failure on the very FIRST page (page + // == 0) leaves nothing to salvage, so the server is still + // reported failed exactly as before this fix. A failure on a + // LATER page — this server answered one or more pages fine + // and only broke partway through pagination — now keeps the + // tools already fetched instead of discarding the whole + // server's contribution; a warning names what was lost so the + // degradation is still visible, matching this fan-out's own + // "partial, not silent" convention for the page-cap case + // above. + if page == 0 && (transportErr != nil || appFailed) { failed = append(failed, name) return } - var result mcpToolsListResult - if err := json.Unmarshal(resp.Result, &result); err != nil { - failed = append(failed, name) - return + if transportErr != nil || appFailed { + failErr := appFailedErr + if transportErr != nil { + // P4: transportErr can be a *url.Error wrapping + // targetURL verbatim (an operator-configured target + // URL's query string is this gateway's own credential + // channel, buildUpstreamTargetURL's doc comment, + // mcp_a2a.go) — scrubbed the same way every other + // target-URL-bearing log line in this file is. + failErr = errors.New(sanitizeProviderErr(transportErr.Error(), targetURL)) + } + g.logf("federated tools/list: server %q: page %d failed: %v; keeping %d tool(s) already fetched from earlier pages", name, page, failErr, len(pageTools)) } - for _, tool := range result.Tools { - merged = append(merged, mcpTool{Name: name + "_" + tool.Name, Description: tool.Description, InputSchema: tool.InputSchema}) + for _, tool := range pageTools { + prefixed, fullName, ok := prefixMCPTool(name, tool) + if !ok { + continue // malformed tool entry from this server; skip just that one, not the whole server + } + merged = append(merged, mcpMergedTool{name: fullName, server: name, raw: prefixed}) } }(name) } wg.Wait() - if len(names) > 0 { - attempted := make([]limitScope, len(names)) - for i, name := range names { + if len(dialed) > 0 { + attempted := make([]limitScope, len(dialed)) + for i, name := range dialed { attempted[i] = limitScope{kind: targetKindMCP, id: name} } g.limiter.countTargetRequests(attempted) @@ -1063,16 +1435,77 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons return } - sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name }) - writeJSONRPCResult(w, format, req.ID, mcpToolsListResult{Tools: merged}) + // F4: sort by (name, server) so a collision's outcome is deterministic + // — never "whichever goroutine happened to append first" — then group + // every entry sharing a name and pick the survivor below. + sort.Slice(merged, func(i, j int) bool { + if merged[i].name != merged[j].name { + return merged[i].name < merged[j].name + } + return merged[i].server < merged[j].server + }) + deduped := make([]map[string]json.RawMessage, 0, len(merged)) + for i := 0; i < len(merged); { + j := i + 1 + for j < len(merged) && merged[j].name == merged[i].name { + j++ + } + group := merged[i:j] + survivor := group[0] + if len(group) > 1 { + // P2 (review round 4): the survivor must be the server + // resolveFederatedTool would actually route a tools/call for + // this name to (longest matching server-name prefix) — the + // SAME rule, called against the SAME names this fan-out was + // given, so tools/list can never advertise a schema/ + // description the caller then cannot reach through + // tools/call. Every candidate server here is by construction + // one of names (mcpMergedTool.server, set from the fan-out's + // own per-server loop) and a valid prefix of group[i].name + // (prefixMCPTool's own "_" convention), + // so resolveFederatedTool always resolves to one of them; the + // !ok branch below is unreachable in practice and only keeps + // group[0] as the same safe default the rest of this file + // uses for its own unreachable branches. + if routedServer, _, ok := resolveFederatedTool(group[0].name, names); ok { + for _, m := range group { + if m.server == routedServer { + survivor = m + break + } + } + } + for _, m := range group { + if m.server == survivor.server { + continue + } + g.logf("federated tools/list: tool %q from server %q dropped: tools/call routes this name to server %q instead", m.name, m.server, survivor.server) + } + } + deduped = append(deduped, survivor.raw) + i = j + } + + writeJSONRPCResult(w, format, req.ID, mcpToolsListResult{Tools: deduped}) } // mcpToolCallParams is one "tools/call" request's params: name is the // caller-facing, server-prefixed tool id federation itself minted in -// tools/list ("_"). +// tools/list ("_"). Meta carries the request's own +// "_meta" object verbatim — MCP's out-of-band per-request extension +// point, whose only field this gateway's own callers have actually used +// is progressToken (F11, review round 3, 2026-09) — forwarded to the +// resolved backend unchanged (mcpFederatedToolsCall) so a caller that +// expects progress correlated by that token at least reaches a backend +// that might honor it. This gateway itself never relays a progress +// notification back to the client (mcpFederatedToolsCall's own doc +// comment covers what it does and does not relay), so a caller depending +// on the notification arriving still needs a direct connection to the +// backend for that half. type mcpToolCallParams struct { Name string `json:"name"` Arguments json.RawMessage `json:"arguments,omitempty"` + Meta json.RawMessage `json:"_meta,omitempty"` } // resolveFederatedTool finds which of allowedNames is fullName's own @@ -1140,7 +1573,9 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons // out above: err != nil is a failure, a JSON-RPC-level resp.Error is // not. probeStart := time.Now() - resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments}, mcpBackendCallResponseMaxBytes) + // Meta: params.Meta (F11) — the client's own "_meta" (progressToken + // included, when set) forwarded verbatim to the resolved backend. + resp, err := g.mcpBackendCall(ctx, targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments, Meta: params.Meta}, mcpBackendCallResponseMaxBytes) // G4: same dialed-server rule as mcpFederatedToolsList's own fan-out, // above — ctx here derives from r.Context() too, so a client cancel // (context.Canceled) must record nothing. Unlike the fan-out, @@ -1154,7 +1589,12 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, format mcpRespons } g.limiter.countTargetRequest(targetKindMCP, serverName) if err != nil { - g.logf("federated tools/call: server %q: %v", serverName, err) + // P4 (review round 4): err here is commonly a *url.Error wrapping + // targetURL verbatim, including its query string (this gateway's + // own credential channel for a target — mcp_a2a.go) — scrubbed + // the same way as every other target-URL-bearing log line in + // this file. + g.logf("federated tools/call: server %q: %v", serverName, sanitizeProviderErr(err.Error(), targetURL)) // A response that hit mcpBackendCallResponseMaxBytes gets its own // specific message, not the generic "upstream error" a truncated- // then-unparsable body would otherwise produce (security review diff --git a/mcp_federation_sse_test.go b/mcp_federation_sse_test.go index f8b8e54..3432144 100644 --- a/mcp_federation_sse_test.go +++ b/mcp_federation_sse_test.go @@ -160,9 +160,16 @@ func TestHandleMCPFederated_DefaultFraming_ByteIdenticalToBaseline(t *testing.T) }, }, { + // F2, review round 3, 2026-09: each tool now round-trips + // through map[string]json.RawMessage (mcpToolsListResult's + // own doc comment) rather than the earlier fixed-field + // struct, so its own keys marshal back out in SORTED order — + // encoding/json's own documented behavior for a map — not + // the struct-declared "name" then "description" order the + // prior baseline pinned. Cosmetic only: content is identical. 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", + wantBody: `{"jsonrpc":"2.0","result":{"tools":[{"description":"alpha's tool","name":"alpha_lookup"},{"description":"beta's tool","name":"beta_search"}]},"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"}}) diff --git a/mcp_federation_test.go b/mcp_federation_test.go index 9fc1f41..f83be4a 100644 --- a/mcp_federation_test.go +++ b/mcp_federation_test.go @@ -113,6 +113,20 @@ func TestParseBackendJSONRPC(t *testing.T) { body: "data: {\"jsonrpc\":\"2.0\",\"id\":\"federated\",\"result\":{\"stage\":1}}\n\ndata: {\"jsonrpc\":\"2.0\",\"id\":\"federated\",\"result\":{\"stage\":2}}\n\n", wantResult: `{"stage":2}`, }, + { + // F6, review round 3, 2026-09: a single event's own "data:" + // field can legitimately span several physical lines (this + // gateway's own sseWriter.writeData emits exactly this shape + // for a multi-line payload) — the SSE wire format joins them + // with "\n" into one logical payload, not two separate + // events. Split right after the "result" key's colon, where + // JSON permits whitespace (including a newline) before the + // value. + name: "SSE event with one data field split across two data: lines, joined with newline", + contentType: "text/event-stream", + body: "data: {\"jsonrpc\":\"2.0\",\"id\":\"federated\",\"result\":\ndata: {\"ok\":true}}\n\n", + wantResult: `{"ok":true}`, + }, { name: "malformed JSON", contentType: "application/json", @@ -147,6 +161,84 @@ func TestParseBackendJSONRPC(t *testing.T) { // --- end-to-end: POST /mcp --- +// mcpToolsToRaw converts the test-fixture []mcpTool shape into the +// []map[string]json.RawMessage wire shape mcpToolsListResult.Tools +// actually decodes/encodes as (F2, review round 3, 2026-09) — every mock +// server in this file builds its own canned tools/list response through +// this, so a test still writes the compact mcpTool{Name: ...} literal +// while the bytes on the wire (and so what mcpFederatedToolsList's own +// F2 field-passthrough sees) match a real backend's shape. +func mcpToolsToRaw(t *testing.T, tools []mcpTool) []map[string]json.RawMessage { + t.Helper() + if tools == nil { + return nil + } + out := make([]map[string]json.RawMessage, len(tools)) + for i, tool := range tools { + b, err := json.Marshal(tool) + if err != nil { + t.Fatalf("marshal mcpTool: %v", err) + } + var m map[string]json.RawMessage + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("unmarshal mcpTool: %v", err) + } + out[i] = m + } + return out +} + +// mustToolName decodes tool's own "name" field, failing the test if it is +// missing or not a JSON string — used everywhere a test needs to compare +// mcpFederatedToolsList's own merged output (now []map[string]json.RawMessage, +// F2) against an expected name string. +func mustToolName(t *testing.T, tool map[string]json.RawMessage) string { + t.Helper() + var name string + if err := json.Unmarshal(tool["name"], &name); err != nil { + t.Fatalf("tool has no valid \"name\" field: %+v (%v)", tool, err) + } + return name +} + +// TestPrefixMCPTool_SkipsInvalidName is the P3 regression test (review +// round 4): json.Unmarshal("null", &toolName) succeeds and leaves +// toolName as its zero value (""), identical to an explicit "name":"" — +// without the fix, either one mints and lists a bare "_" tool +// with no real name behind it, rather than being skipped like the +// already-handled missing-field and non-string-name cases. +func TestPrefixMCPTool_SkipsInvalidName(t *testing.T) { + cases := []struct { + tool map[string]json.RawMessage + name string + }{ + {name: "missing name field", tool: map[string]json.RawMessage{"x": json.RawMessage(`1`)}}, + {name: "null name", tool: map[string]json.RawMessage{"name": json.RawMessage(`null`)}}, + {name: "empty string name", tool: map[string]json.RawMessage{"name": json.RawMessage(`""`)}}, + {name: "non-string name", tool: map[string]json.RawMessage{"name": json.RawMessage(`42`)}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + out, fullName, ok := prefixMCPTool("srv", tc.tool) + if ok || out != nil || fullName != "" { + t.Errorf("prefixMCPTool(%+v) = (%+v, %q, %v), want (nil, \"\", false)", tc.tool, out, fullName, ok) + } + }) + } + // Good case, alongside the bad ones above (table style): a real name + // still prefixes normally. + t.Run("valid name", func(t *testing.T) { + out, fullName, ok := prefixMCPTool("srv", map[string]json.RawMessage{"name": json.RawMessage(`"search"`)}) + if !ok || fullName != "srv_search" { + t.Fatalf("prefixMCPTool = (%+v, %q, %v), want ok=true, fullName=\"srv_search\"", out, fullName, ok) + } + var gotName string + if err := json.Unmarshal(out["name"], &gotName); err != nil || gotName != "srv_search" { + t.Errorf("out[\"name\"] = %s, want \"srv_search\"", out["name"]) + } + }) +} + // mockJSONRPCServer is a minimal MCP-shaped backend for federation tests: // it answers "tools/list" with a fixed tool set and "tools/call" by // echoing back the resolved (prefix-stripped) name/arguments it was @@ -172,7 +264,7 @@ func newMockJSONRPCServer(t *testing.T, tools []mcpTool) *mockJSONRPCServer { w.Header().Set("Content-Type", "application/json") switch req.Method { case "tools/list": - result, _ := json.Marshal(mcpToolsListResult{Tools: m.tools}) + result, _ := json.Marshal(mcpToolsListResult{Tools: mcpToolsToRaw(t, m.tools)}) _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) case "tools/call": var params mcpToolCallParams @@ -244,7 +336,7 @@ func newSessionRequiredMockServer(t *testing.T, tools []mcpTool) *sessionRequire } switch req.Method { case "tools/list": - result, _ := json.Marshal(mcpToolsListResult{Tools: m.tools}) + result, _ := json.Marshal(mcpToolsListResult{Tools: mcpToolsToRaw(t, m.tools)}) _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) case "tools/call": var params mcpToolCallParams @@ -523,7 +615,7 @@ func TestHandleMCPFederated_ToolsList_AggregatesPrefixesAndRespectsGroupAccess(t gotNames := make([]string, len(result.Tools)) for i, tool := range result.Tools { - gotNames[i] = tool.Name + gotNames[i] = mustToolName(t, tool) } if len(gotNames) != len(tc.wantToolNames) { t.Fatalf("tools = %v, want %v", gotNames, tc.wantToolNames) @@ -581,11 +673,249 @@ func TestHandleMCPFederated_ToolsList_UnreachableServerDegradesNotFails(t *testi 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" { + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "alpha_lookup" { t.Errorf("tools = %+v, want exactly alpha_lookup (beta unreachable, skipped)", result.Tools) } } +// TestHandleMCPFederated_ToolsList_FollowsNextCursor is the F3 regression +// test (review round 3, 2026-09): a paginated backend's tools all arrive +// in the federated aggregate, not just the first page, and the cursor +// each successive request carries is exactly the previous page's own +// nextCursor. +func TestHandleMCPFederated_ToolsList_FollowsNextCursor(t *testing.T) { + var calls []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + var params mcpToolsListParams + _ = json.Unmarshal(req.Params, ¶ms) + calls = append(calls, "cursor="+params.Cursor) + + var page mcpToolsListResult + switch params.Cursor { + case "": + page = mcpToolsListResult{Tools: mcpToolsToRaw(t, []mcpTool{{Name: "one"}}), NextCursor: "page2"} + case "page2": + page = mcpToolsListResult{Tools: mcpToolsToRaw(t, []mcpTool{{Name: "two"}})} // no NextCursor: last page + default: + t.Fatalf("unexpected cursor %q", params.Cursor) + } + result, _ := json.Marshal(page) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + })) + defer srv.Close() + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) // alphaOnly + 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", got.Error) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + gotNames := make([]string, len(result.Tools)) + for i, tool := range result.Tools { + gotNames[i] = mustToolName(t, tool) + } + wantNames := []string{"alpha_one", "alpha_two"} + if len(gotNames) != len(wantNames) || gotNames[0] != wantNames[0] || gotNames[1] != wantNames[1] { + t.Fatalf("tools = %v, want %v (both pages merged)", gotNames, wantNames) + } + if len(calls) != 2 || calls[0] != "cursor=" || calls[1] != "cursor=page2" { + t.Errorf("calls = %v, want [\"cursor=\" \"cursor=page2\"] (first page bare, second page echoing the first page's own nextCursor)", calls) + } +} + +// TestHandleMCPFederated_ToolsList_LaterPageFailure_KeepsEarlierPages is +// the F3 partial-keep regression test (review round 4): page 1 succeeds +// and sets a NextCursor, page 2 fails with a JSON-RPC error. Before this +// fix, ANY page failure discarded the whole server's contribution, +// including the tools from pages already fetched successfully; now only +// page 0 (the very first page) failing does that — a later page failing +// keeps what came before and logs a warning naming what was lost. +func TestHandleMCPFederated_ToolsList_LaterPageFailure_KeepsEarlierPages(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + var params mcpToolsListParams + _ = json.Unmarshal(req.Params, ¶ms) + w.Header().Set("Content-Type", "application/json") + switch params.Cursor { + case "": + result, _ := json.Marshal(mcpToolsListResult{Tools: mcpToolsToRaw(t, []mcpTool{{Name: "one"}}), NextCursor: "page2"}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + case "page2": + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "backend broke mid-pagination"}}) + default: + t.Fatalf("unexpected cursor %q", params.Cursor) + } + })) + defer srv.Close() + + cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) // alphaOnly + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + var rec *httptest.ResponseRecorder + logOutput := captureStderr(t, func() { + 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 — page 1's tool must still be served, not turned into a total 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 || mustToolName(t, result.Tools[0]) != "alpha_one" { + t.Fatalf("tools = %+v, want exactly [alpha_one] (page 1's tool kept despite page 2 failing)", result.Tools) + } + if !strings.Contains(logOutput, "page 1 failed") || !strings.Contains(logOutput, "keeping 1 tool") { + t.Errorf("log output = %q, want a warning naming the failed page and how many tools were kept", logOutput) + } +} + +// TestHandleMCPFederated_ToolsList_NextCursorLoop_BoundedByPageCap proves +// a backend that ALWAYS sets nextCursor (an unbounded pagination loop, by +// bug or by design) cannot hold the fan-out hostage forever — F3's own +// mcpToolsListMaxPagesPerServer cap stops it, degrading to a partial tool +// list rather than exhausting the fan-out's shared timeout budget one +// page at a time. Also covers F3's page-cap warning (review round 4): a +// truncation this silent otherwise gives an operator no signal a +// legitimately larger catalog got cut off. +func TestHandleMCPFederated_ToolsList_NextCursorLoop_BoundedByPageCap(t *testing.T) { + var pageCount atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := pageCount.Add(1) + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + page := mcpToolsListResult{ + Tools: mcpToolsToRaw(t, []mcpTool{{Name: fmt.Sprintf("tool%d", n)}}), + NextCursor: fmt.Sprintf("next%d", n), // never empty: an infinite pagination loop + } + result, _ := json.Marshal(page) + w.Header().Set("Content-Type", "application/json") + _ = 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) + } + + var rec *httptest.ResponseRecorder + logOutput := captureStderr(t, func() { + 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()) + } + if !strings.Contains(logOutput, "hit the") || !strings.Contains(logOutput, "page cap") { + t.Errorf("log output = %q, want a warning naming the page-cap truncation", logOutput) + } + + if got := pageCount.Load(); got != int32(mcpToolsListMaxPagesPerServer) { + t.Errorf("pages fetched = %d, want exactly mcpToolsListMaxPagesPerServer (%d)", got, mcpToolsListMaxPagesPerServer) + } +} + +// TestHandleMCPFederated_ToolsList_CollisionDropsLaterDeterministically is +// the F4 regression test (review round 3, 2026-09), extended for P2 +// (review round 4): two configured, allowed servers whose names collide +// under the "_" prefix convention (server "a" has tool +// "b_c"; server "a_b" has tool "c"; both mint "a_b_c") must not both +// appear in the merged list — one wins, deterministically, and the drop +// is logged. The survivor must be server "a_b", not "a": resolveFederatedTool +// resolves "a_b_c" against allowed names ["a","a_b"] by LONGEST matching +// prefix, which is "a_b" — the OLD alphabetical tiebreak instead kept "a" +// (it sorts first), which would advertise a tool description tools/call +// could never actually reach, since every call for "a_b_c" is routed to +// "a_b" regardless of which one tools/list lists. Each server's tool +// carries its own description so the assertion below proves the LISTED +// description is the ROUTED server's own, not just that a name survived. +func TestHandleMCPFederated_ToolsList_CollisionDropsLaterDeterministically(t *testing.T) { + a := newMockJSONRPCServer(t, []mcpTool{{Name: "b_c", Description: "FROM_A"}}) + ab := newMockJSONRPCServer(t, []mcpTool{{Name: "c", Description: "FROM_A_B"}}) + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{ + "a": {URL: a.srv.URL}, + "a_b": {URL: ab.srv.URL}, + } + 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")})) + 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) + } + var result mcpToolsListResult + if err := json.Unmarshal(got.Result, &result); err != nil { + t.Fatalf("decode result: %v", err) + } + // Exactly one "a_b_c" survives — never two entries sharing one name, + // which some client-side LLM tool APIs reject outright, and never + // zero (the collision must resolve to a deterministic winner, not + // drop both). + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "a_b_c" { + t.Fatalf("tools = %+v, want exactly one entry named a_b_c", result.Tools) + } + // P2: the listed description must belong to server "a_b" (FROM_A_B) — + // the one resolveFederatedTool actually routes "a_b_c" calls to — + // never server "a"'s (FROM_A), which the old alphabetical tiebreak + // would have kept despite tools/call never reaching it under that + // name. + var desc string + if err := json.Unmarshal(result.Tools[0]["description"], &desc); err != nil { + t.Fatalf("tool has no valid \"description\" field: %+v (%v)", result.Tools[0], err) + } + if desc != "FROM_A_B" { + t.Errorf("description = %q, want %q (the routed server a_b's own, not a's)", desc, "FROM_A_B") + } + if routedServer, _, ok := resolveFederatedTool("a_b_c", []string{"a", "a_b"}); !ok || routedServer != "a_b" { + t.Fatalf("sanity check failed: resolveFederatedTool(%q) = (%q, ok=%v), want (\"a_b\", true) — the test's own premise is wrong if this fails", "a_b_c", routedServer, ok) + } +} + // TestHandleMCPFederated_ToolsCall_RoutesToCorrectServer_StripsPrefix // proves tools/call resolves the server prefix, strips it before // forwarding, contacts ONLY the resolved server (beta must never be @@ -655,6 +985,46 @@ func TestHandleMCPFederated_ToolsCall_RoutesToCorrectServer_StripsPrefix(t *test } } +// TestHandleMCPFederated_ToolsCall_ForwardsMeta is the F11 regression test +// (review round 3, 2026-09): a client's own "_meta" object on tools/call — +// progressToken included — reaches the resolved backend verbatim, +// alongside the usual prefix-stripped name and forwarded arguments. +func TestHandleMCPFederated_ToolsCall_ForwardsMeta(t *testing.T) { + var gotMeta json.RawMessage + alpha := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req jsonrpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + var params mcpToolCallParams + _ = json.Unmarshal(req.Params, ¶ms) + gotMeta = params.Meta + w.Header().Set("Content-Type", "application/json") + result, _ := json.Marshal(map[string]any{"ok": true}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + })) + defer alpha.Close() + + cfg := newFederationTestConfig(alpha.URL, "http://beta.invalid", true) // alphaOnly + 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", + Meta: json.RawMessage(`{"progressToken":"tok-1"}`), + }) + 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()) + } + if string(gotMeta) != `{"progressToken":"tok-1"}` { + t.Errorf("backend saw _meta = %s, want the client's own _meta forwarded verbatim", gotMeta) + } +} + // TestHandleMCPFederated_ToolsCall_BackendError_RelayedToClient proves a // backend's own JSON-RPC error (e.g. a tool execution failure) is relayed // to the client verbatim, under the client's own id — not swallowed or @@ -978,14 +1348,25 @@ func TestHandleMCPFederated_ToolsCall_BackendEmptyEnvelope_ConvertsToInternalErr // 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). +// TestMcpBackendCall_HandshakeFallbackTrigger covers mcpSessionRequiredSignal's +// narrowed rule (F1, review round 3, 2026-09) end to end through +// mcpBackendCall: only a narrow session-missing signal may trigger the +// handshake retry, never an arbitrary tools/call error and never a 429 — +// the exact over-eager-retry bug F1 fixes, where a non-idempotent +// tools/call could be silently double-executed. method defaults to +// "tools/list" when empty, matching the broader "any JSON-RPC error is a +// session signal" rule that method gets (mcpSessionRequiredSignal's own +// doc comment); the "tool called once" cases below explicitly exercise +// "tools/call" to prove the narrower rule that method gets instead. func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) { cases := []struct { bareHandler http.HandlerFunc name string + method string wantHandshake bool }{ { - name: "bare JSON-RPC error triggers the handshake", + name: "bare JSON-RPC error triggers the handshake (tools/list: any error is a session signal)", bareHandler: func(w http.ResponseWriter, r *http.Request) { var req jsonrpcRequest _ = json.NewDecoder(r.Body).Decode(&req) @@ -995,11 +1376,25 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) { wantHandshake: true, }, { - name: "HTTP 403 with no JSON-RPC body triggers the handshake", + name: "HTTP 404 with no JSON-RPC body triggers the handshake", + bareHandler: func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + }, + wantHandshake: true, + }, + { + name: "HTTP 403 with no JSON-RPC body does NOT trigger the handshake (F1: narrowed from any 4xx to just 400/404)", bareHandler: func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusForbidden) }, - wantHandshake: true, + wantHandshake: false, + }, + { + name: "HTTP 429 never triggers the handshake (F1: rate-limited, retrying only adds load)", + bareHandler: func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + }, + wantHandshake: false, }, { name: "HTTP 500 never triggers the handshake", @@ -1008,9 +1403,43 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) { }, wantHandshake: false, }, + { + // The F1 scenario itself: a tools/call answered with an + // ordinary, unrelated JSON-RPC error (a real tool-level + // failure, e.g. bad arguments) must be called exactly ONCE — + // retrying it could silently double-execute a non-idempotent + // tool's side effect. + name: "tools/call: arbitrary JSON-RPC error does not trigger the handshake (tool called once)", + method: "tools/call", + 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: -32602, Message: "invalid arguments: missing required field \"query\""}}) + }, + wantHandshake: false, + }, + { + // The narrow allow-case for tools/call: the SAME JSON-RPC + // error shape, but its own message clearly names a session + // problem, still gets the handshake. + name: "tools/call: JSON-RPC error naming a session problem still triggers the handshake", + method: "tools/call", + 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: call initialize first"}}) + }, + wantHandshake: true, + }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { + method := tc.method + if method == "" { + method = "tools/list" + } var callCount atomic.Int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if callCount.Add(1) == 1 { @@ -1040,14 +1469,14 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) { t.Fatal("handler is not *Gateway") } - _, _ = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes) + _, _ = gw.mcpBackendCall(context.Background(), srv.URL, method, struct{}{}, mcpBackendResponseMaxBytes) 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) + t.Errorf("calls = %d, want exactly 1 (bare attempt only, no handshake retry)", gotCalls) } }) } @@ -1227,13 +1656,15 @@ func TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetr 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" { + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "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) + // F5, review round 3, 2026-09: notifications/initialized is now sent + // between the handshake and the retry, one more call than before. + if len(seq) != 5 { + t.Fatalf("call sequence = %v, want exactly 5 calls (bare attempt, initialize, notifications/initialized, 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]) @@ -1241,11 +1672,98 @@ func TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetr 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[2] != "notifications/initialized:session=sess-1" { + t.Errorf("call[2] = %q, want the lifecycle's own notifications/initialized, 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]) + if seq[3] != "tools/list:session=sess-1" { + t.Errorf("call[3] = %q, want the retried tools/list carrying the session header the handshake returned", seq[3]) + } + if seq[4] != "DELETE:session=sess-1" { + t.Errorf("call[4] = %q, want the best-effort session close carrying the same session header", seq[4]) + } +} + +// TestMcpBackendCall_SendsNotificationsInitialized_AfterHandshake is the +// F5 regression test (review round 3, 2026-09), pinning the notification's +// own exact wire shape directly: no "id" field at all (it is a JSON-RPC +// notification, not a request awaiting a reply), method +// "notifications/initialized", and the session header carrying the +// handshake's own Mcp-Session-Id. A backend that answers it with an +// error (as this one deliberately does, mirroring a real server that +// does not implement this endpoint) must never surface as a failure of +// the call this handshake serves — fire-and-forget, as documented. +func TestMcpBackendCall_SendsNotificationsInitialized_AfterHandshake(t *testing.T) { + var ( + mu sync.Mutex + sawNotification bool + notificationHasID bool + notificationSessID string + ) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodDelete { + w.WriteHeader(http.StatusNoContent) + return + } + body, _ := io.ReadAll(r.Body) + var raw map[string]json.RawMessage + _ = json.Unmarshal(body, &raw) + var req jsonrpcRequest + _ = json.Unmarshal(body, &req) + w.Header().Set("Content-Type", "application/json") + + switch req.Method { + case "initialize": + w.Header().Set(mcpSessionHeader, "sess-f5") + result, _ := json.Marshal(map[string]any{}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + case "notifications/initialized": + mu.Lock() + sawNotification = true + _, notificationHasID = raw["id"] + notificationSessID = r.Header.Get(mcpSessionHeader) + mu.Unlock() + // Deliberately answer with an error, to prove this response + // is never inspected by the caller. + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: jsonrpcMethodNotFound, Message: "not implemented"}}) + default: + if r.Header.Get(mcpSessionHeader) == "sess-f5" { + // The retried real call, now carrying the session the + // handshake returned: succeeds. + result, _ := json.Marshal(mcpToolsListResult{}) + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) + return + } + // The bare attempt: reject with a session-required signal so + // the handshake fires. + _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}}) + } + })) + 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") + } + + if _, err := gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{}, mcpBackendResponseMaxBytes); err != nil { + t.Fatalf("mcpBackendCall: %v, want the retry to succeed despite the notification's own error response", err) + } + + mu.Lock() + defer mu.Unlock() + if !sawNotification { + t.Fatal("backend never received notifications/initialized") + } + if notificationHasID { + t.Error("notifications/initialized carried an \"id\" field; want none (it is a notification, not a request)") + } + if notificationSessID != "sess-f5" { + t.Errorf("notifications/initialized session header = %q, want %q", notificationSessID, "sess-f5") } } @@ -1506,7 +2024,7 @@ func TestHandleMCPFederated_ToolsList_PanickingBackendDegradesNotCrashes(t *test 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" { + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "alpha_lookup" { t.Errorf("tools = %+v, want exactly alpha_lookup (beta panicked, skipped)", result.Tools) } } @@ -1624,8 +2142,13 @@ func TestIsLegacySSETransportURL(t *testing.T) { }{ {"http://playwright.internal:8080/sse", true}, {"http://playwright.internal:8080/sse?foo=bar", true}, + // F9, review round 3, 2026-09: a trailing slash is the identical + // legacy-SSE endpoint, not a different path shape. + {"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/mcp/", false}, {"http://mcp.internal:8080/sse/extra", false}, {"://not a valid url", false}, // malformed -> safe default: not excluded } @@ -1683,7 +2206,7 @@ func TestHandleMCPFederated_ToolsList_ExcludesLegacySSEServer_NeverContacted(t * 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" { + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "alpha_lookup" { t.Fatalf("tools = %+v, want exactly alpha_lookup (playwright excluded)", result.Tools) } if legacy.called.Load() != 0 { @@ -1849,7 +2372,7 @@ func TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller(t *testing.T) 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"}}}) + result, _ := json.Marshal(mcpToolsListResult{Tools: mcpToolsToRaw(t, []mcpTool{{Name: "search"}})}) _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result}) } })) @@ -1877,11 +2400,39 @@ func TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller(t *testing.T) 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" { + if len(result.Tools) != 1 || mustToolName(t, result.Tools[0]) != "alpha_search" { t.Fatalf("tools = %+v, want alpha_search despite the session-close failure", result.Tools) } } +// TestMcpBackendCloseSession_LogScrubsCredentialFromTargetURL is the P4 +// regression test (review round 4): a configured target URL's query +// string is this gateway's own credential channel for an MCP target +// (buildUpstreamTargetURL's doc comment, mcp_a2a.go), and a failed Do() +// against it commonly returns a *url.Error wrapping that URL verbatim. +// Before this fix, mcpBackendCloseSession's best-effort failure log +// printed both the raw targetURL and the raw err, leaking the credential +// straight to stderr; both must now come out scrubbed. +func TestMcpBackendCloseSession_LogScrubsCredentialFromTargetURL(t *testing.T) { + gw := newTestGatewayForLogger(t) + // 127.0.0.1:1 is a reserved, always-refused port — Do() fails fast + // with a real *url.Error wrapping this exact URL, no network flake + // risk and no server to stand up for a pure "the log scrubs its own + // error text" assertion. + const targetURL = "http://127.0.0.1:1/mcp?api-key=SUPERSECRET123" + + logOutput := captureStderr(t, func() { + gw.mcpBackendCloseSession(context.Background(), targetURL, "sess-1") + }) + + if strings.Contains(logOutput, "SUPERSECRET123") { + t.Errorf("log output = %q, want the api-key query value scrubbed", logOutput) + } + if !strings.Contains(logOutput, "127.0.0.1") { + t.Errorf("log output = %q, want the target host still present (only credentials must be stripped)", logOutput) + } +} + // --- feat/target-health: passive recording from federation traffic --- // TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer proves @@ -2187,6 +2738,31 @@ func TestHandleMCPFederated_ToolsList_NeverDialedSemaphoreStarved_RecordsNothing if neverDialed != 1 { t.Errorf("neverDialed = %d, want exactly 1", neverDialed) } + + // F10, review round 3, 2026-09: countTargetRequests must attribute a + // request to every DIALED server only — the one server target_health + // above confirms was never dialed (snap.observed == false) must show + // ZERO on its per-target request counter, even though it was one of + // the "allowed" servers this call named. Before the fix, EVERY name + // in the allowed set moved this counter regardless of whether a + // goroutine ever won the semaphore for it. getCounter's own "ok" + // return is not useful here — the in-process fallback store this test + // runs against (no Redis configured) reports ok=true unconditionally, + // even for a key that was never incremented (limiter.storeGet's own + // doc, limits.go) — so the count itself is the only signal that + // matters. + now := time.Now() + for _, name := range names { + snap := gw.targetHealth.snapshot(targetKindMCP, name) + got, _ := gw.limiter.getCounter(targetKindMCP, name, metricReq, windowDay, now) + want := int64(0) + if snap.observed { + want = 1 + } + if got != want { + t.Errorf("mcp/%s req:day counter = %d, want %d (dialed=%v)", name, got, want, snap.observed) + } + } } // TestHandleMCPFederated_ToolsCall_ClientCanceled_RecordsNothing mirrors