mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
fix(mcp): stop duplicate tool calls and keep full tool metadata
Federated tools/call re-initializes a session only when the backend signals a missing session: HTTP 400 or 404, or a JSON-RPC error that names a session. A failing tool is no longer called twice. tools/list now: - passes each tool object through and rewrites only its name - follows nextCursor up to a page cap and keeps the earlier pages if a later page fails - skips tools with no name - resolves name collisions to the server that tools/call routes to The handshake sends notifications/initialized. tools/call forwards _meta. Multi-line SSE data is joined. A2A target URLs are joined with url.JoinPath and keep their query string. agents[].card must start with "/". Logged target URLs and errors are scrubbed of credentials. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
+84
-6
@@ -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
|
||||
}
|
||||
|
||||
+127
@@ -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) {
|
||||
|
||||
+523
-83
@@ -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
|
||||
// "<serverName>_<original name>" — 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 "<serverName>_" 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 "<serverName>_<toolName>", 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 "<serverName>_<toolName>" (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 "<serverName>_<toolName>" 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 ("<serverName>_<toolName>").
|
||||
// tools/list ("<serverName>_<toolName>"). 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
|
||||
|
||||
@@ -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"}})
|
||||
|
||||
+596
-20
@@ -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 "<server>_" 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 "<server>_<tool>" 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
|
||||
|
||||
Reference in New Issue
Block a user