mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-16 05:37:28 +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/http"
|
||||||
"net/url"
|
"net/url"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -211,12 +212,14 @@ func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *u
|
|||||||
}
|
}
|
||||||
g.limiter.countTargetRequest(targetScopeKind(kind), name)
|
g.limiter.countTargetRequest(targetScopeKind(kind), name)
|
||||||
|
|
||||||
upstreamURL := targetURL
|
upstreamURL, err := buildUpstreamTargetURL(targetURL, rest, r.URL.RawQuery)
|
||||||
if rest != "" {
|
if err != nil {
|
||||||
upstreamURL += "/" + rest
|
// targetURL is operator-configured and already validated at
|
||||||
}
|
// construction (validateTargetURLs) — unreachable in practice,
|
||||||
if r.URL.RawQuery != "" {
|
// kept as an explicit guard rather than trusting that invariant
|
||||||
upstreamURL += "?" + r.URL.RawQuery
|
// silently.
|
||||||
|
writeOAIError(w, http.StatusInternalServerError, "invalid_request_error", "invalid target configuration")
|
||||||
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// feat/target-health: passive recording (target_health.go) — latency
|
// 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))
|
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
|
// validateTargetURLs checks every configured MCP-server and agent entry at
|
||||||
// construction: the map value must be non-nil, the name must pass
|
// construction: the map value must be non-nil, the name must pass
|
||||||
// validateConfigName (the same character-set, dot-only, and reserved-name
|
// 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 {
|
if err := validateTargetURL(ac.URL); err != nil {
|
||||||
return fmt.Errorf("llmgateway: agents %q: %w", name, err)
|
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
|
return nil
|
||||||
}
|
}
|
||||||
|
|||||||
+127
@@ -5,6 +5,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
|
"net/url"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"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 ---
|
// --- dead upstream maps to 502, same as native passthrough ---
|
||||||
|
|
||||||
func TestHandleTargetProxy_DeadUpstream_Returns502(t *testing.T) {
|
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"}},
|
mcpServers: map[string]*TargetConfig{"has a space": {URL: "http://mcp-alpha.internal:8080"}},
|
||||||
wantErr: true,
|
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 {
|
for _, tt := range tests {
|
||||||
t.Run(tt.name, func(t *testing.T) {
|
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 in practice, and "don't exclude" is the safe default for an
|
||||||
// unreachable branch — the same convention sanitizeBaseURL applies to its
|
// unreachable branch — the same convention sanitizeBaseURL applies to its
|
||||||
// own unparsable-URL case (admin.go).
|
// 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 {
|
func isLegacySSETransportURL(rawURL string) bool {
|
||||||
u, err := url.Parse(rawURL)
|
u, err := url.Parse(rawURL)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false
|
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
|
// 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
|
// mcpTool is a convenience literal type this repository's own tests use
|
||||||
// back from a backend server's own response or as re-emitted (name
|
// to build a canned backend tools/list response (newMockJSONRPCServer,
|
||||||
// prefixed) in this gateway's federated aggregate.
|
// 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 {
|
type mcpTool struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Description string `json:"description,omitempty"`
|
Description string `json:"description,omitempty"`
|
||||||
@@ -559,9 +572,118 @@ type mcpTool struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// mcpToolsListResult is the "result" object of an MCP "tools/list"
|
// 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 {
|
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
|
// 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)
|
resp, err := g.targetClient.Do(req)
|
||||||
if err != nil {
|
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
|
return
|
||||||
}
|
}
|
||||||
_ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure
|
_ = 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 —
|
// mcpBackendCall issues a JSON-RPC 2.0 request for method to targetURL —
|
||||||
// bare first, matching the per-request "no persisted CLIENT session"
|
// bare first, matching the per-request "no persisted CLIENT session"
|
||||||
// design (handleMCPFederated's own doc comment) — and falls back to a
|
// design (handleMCPFederated's own doc comment) — and falls back to a
|
||||||
// ONE-SHOT initialize/retry/close handshake ONLY when the bare attempt
|
// ONE-SHOT initialize/retry/close handshake ONLY when the bare attempt's
|
||||||
// signals the server actually requires a session: a JSON-RPC-level error
|
// own outcome narrowly signals a missing session (mcpSessionRequiredSignal,
|
||||||
// response, or an HTTP 4xx status.
|
// 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
|
// 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,
|
// 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
|
return resp, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
retryEligible := (err == nil && resp.Error != nil) || (status >= 400 && status < 500)
|
if !mcpSessionRequiredSignal(method, status, resp, err) {
|
||||||
if !retryEligible {
|
// Not a session-required signal: relay the bare attempt's own
|
||||||
return nil, err
|
// 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)
|
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL, maxBytes)
|
||||||
@@ -772,6 +998,16 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
|
|||||||
}
|
}
|
||||||
return resp, nil
|
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 != "" {
|
if sessionID != "" {
|
||||||
defer func() {
|
defer func() {
|
||||||
closeCtx, cancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout)
|
closeCtx, cancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout)
|
||||||
@@ -821,18 +1057,29 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err
|
|||||||
return &out, nil
|
return &out, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// lastSSEDataLine returns the payload of the last "data:" line in an SSE
|
// lastSSEDataLine returns the payload of the last COMPLETE SSE event's
|
||||||
// body — body itself, unchanged, if it contains no such line at all (so
|
// "data:" field in an SSE body — body itself, unchanged, if it contains
|
||||||
// the caller's own json.Unmarshal fails informatively on whatever body
|
// no "data:" line at all (so the caller's own json.Unmarshal fails
|
||||||
// actually was, rather than this function silently swallowing a
|
// informatively on whatever body actually was, rather than this function
|
||||||
// malformed stream into an empty byte slice).
|
// 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
|
// 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
|
// Streamable-HTTP transport layer would (matching a specific response to
|
||||||
// the request that triggered it by id, ignoring unrelated
|
// the request that triggered it by id, ignoring unrelated
|
||||||
// server-initiated notifications interleaved on the same stream) — it
|
// 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-
|
// documented limitation, acceptable for mcpBackendCall's own one-request-
|
||||||
// in-flight-at-a-time usage of this stream (it never has two outstanding
|
// 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
|
// 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
|
// same stream would be misread by this function. No production server
|
||||||
// probed for this round exhibited that behavior.
|
// probed for this round exhibited that behavior.
|
||||||
func lastSSEDataLine(body []byte) []byte {
|
func lastSSEDataLine(body []byte) []byte {
|
||||||
var last []byte
|
normalized := bytes.ReplaceAll(body, []byte("\r\n"), []byte("\n"))
|
||||||
for _, line := range bytes.Split(body, []byte("\n")) {
|
events := bytes.Split(normalized, []byte("\n\n"))
|
||||||
line = bytes.TrimRight(line, "\r")
|
for i := len(events) - 1; i >= 0; i-- {
|
||||||
if data, ok := bytes.CutPrefix(line, []byte("data: ")); ok {
|
var dataLines [][]byte
|
||||||
last = data
|
for _, line := range bytes.Split(events[i], []byte("\n")) {
|
||||||
} else if data, ok := bytes.CutPrefix(line, []byte("data:")); ok {
|
line = bytes.TrimRight(line, "\r")
|
||||||
last = bytes.TrimSpace(data)
|
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 body
|
|
||||||
}
|
|
||||||
return last
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// mcpFederatedFanoutConcurrency bounds how many of names' backend servers
|
// mcpFederatedFanoutConcurrency bounds how many of names' backend servers
|
||||||
@@ -883,38 +1134,72 @@ func lastSSEDataLine(body []byte) []byte {
|
|||||||
const mcpFederatedFanoutConcurrency = 8
|
const mcpFederatedFanoutConcurrency = 8
|
||||||
|
|
||||||
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
// mcpFederatedToolsList implements the federated "tools/list": it fans
|
||||||
// out one mcpBackendCall per server in names (handleMCPFederated's own
|
// out one per-server fetch (below) per server in names (handleMCPFederated's
|
||||||
// allowedMCPServerNames(grp) call, made once there so the exact set
|
// own allowedMCPServerNames(grp) call, made once there so the exact set
|
||||||
// contacted is always the one the client actually asked to reach), at
|
// contacted is always the one the client actually asked to reach), at
|
||||||
// most mcpFederatedFanoutConcurrency at a time — every backend sharing
|
// most mcpFederatedFanoutConcurrency at a time — every backend sharing
|
||||||
// ONE overall toolsListBackendTimeout deadline (fanoutCtx, below) — then
|
// ONE overall toolsListBackendTimeout deadline (fanoutCtx, below) — then
|
||||||
// merges every reachable server's own tools into one result,
|
// merges every reachable server's own tools into one result,
|
||||||
// each tool's name prefixed "<serverName>_<toolName>", the exact
|
// each tool's name prefixed "<serverName>_<toolName>" (prefixMCPTool; F2's
|
||||||
// underscore-separator convention the old agentgateway this plugin
|
// own doc comment covers why every OTHER field of the tool object passes
|
||||||
// replaces used, and the one pugbot's mcpclient and agentkit have already
|
// through unchanged), the exact underscore-separator convention the old
|
||||||
// persisted into their own tool-id databases (e.g.
|
// agentgateway this plugin replaces used, and the one pugbot's mcpclient
|
||||||
// "brave-search_brave_web_search" — verified live). A server that errors,
|
// and agentkit have already persisted into their own tool-id databases
|
||||||
// times out, panics, or returns an unparsable/invalid result is logged
|
// (e.g. "brave-search_brave_web_search" — verified live). Per server, a
|
||||||
// and skipped, not surfaced as a whole-call failure — UNLESS every single
|
// paginated backend is followed via nextCursor up to
|
||||||
// attempted server failed, in which case an empty tools list would
|
// mcpToolsListMaxPagesPerServer pages (F3, review round 3, 2026-09) — all
|
||||||
// misleadingly look like "this caller's group has no MCP access" — see
|
// still inside that one server's share of fanoutCtx's shared budget, never
|
||||||
// the loud-failure branch below.
|
// 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
|
// countTargetRequests (limits.go) attributes one request to every server
|
||||||
// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not
|
// actually DIALED here, in ONE batched call after wg.Wait() — "dialed",
|
||||||
// "reached" or "succeeded": a server this gateway dialed and got a
|
// not "attempted" in the looser sense the prior revision used (F10,
|
||||||
// response (or a timeout, a connection refusal, or a panic) from still had
|
// review round 3, 2026-09): a goroutine that was still waiting on the
|
||||||
// a real request sent to it and a real slot of this gateway's outbound
|
// fan-out's own semaphore when fanoutCtx expired never sent a byte to
|
||||||
// capacity spent on it, which is what these counters exist to track. A
|
// its server and must not move that server's own per-target counter —
|
||||||
// single federated tools/list call can move several targets' own
|
// only a server this gateway actually opened a connection to (and got a
|
||||||
// counters, not just one — unlike handleTargetProxy's always-exactly-
|
// response, a timeout, a connection refusal, or a panic from) spent a
|
||||||
// one-target shape (mcp_a2a.go).
|
// 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) {
|
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpResponseFormat, r *http.Request, req jsonrpcRequest, names []string) {
|
||||||
var (
|
var (
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
wg sync.WaitGroup
|
wg sync.WaitGroup
|
||||||
merged = make([]mcpTool, 0, len(names))
|
merged = make([]mcpMergedTool, 0, len(names))
|
||||||
failed []string
|
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
|
// fanoutCtx bounds the WHOLE fan-out to toolsListBackendTimeout from
|
||||||
// when THIS call started — not each individual backend's own start
|
// 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 }()
|
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
|
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()
|
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
|
// G4: this goroutine WON the select above, so it was actually
|
||||||
// dialed — only a client hang-up (context.Canceled,
|
// dialed — only a client hang-up (context.Canceled,
|
||||||
// propagating from r.Context() into fanoutCtx) says nothing
|
// 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
|
// shared toolsListBackendTimeout fired while THIS dialed
|
||||||
// server was still in flight — it did not answer within the
|
// server was still in flight — it did not answer within the
|
||||||
// fan-out's own budget, which is a real failure worth
|
// fan-out's own budget, which is a real failure worth
|
||||||
// recording, not a client artifact (F2 over-suppressed this
|
// recording, not a client artifact. Measured across every
|
||||||
// case; a genuinely hung server used to read "unknown"
|
// page this server's own loop actually ran, not just the
|
||||||
// forever under federation-only traffic).
|
// first: a server that answers its first few pages fine and
|
||||||
if err == nil || !errors.Is(err, context.Canceled) {
|
// then times out on a later one is genuinely unhealthy, not
|
||||||
g.targetHealth.record(targetKindMCP, name, targetURL, err == nil, err, time.Since(probeStart), targetHealthSourceTraffic)
|
// 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()
|
mu.Lock()
|
||||||
defer mu.Unlock()
|
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)
|
failed = append(failed, name)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
var result mcpToolsListResult
|
if transportErr != nil || appFailed {
|
||||||
if err := json.Unmarshal(resp.Result, &result); err != nil {
|
failErr := appFailedErr
|
||||||
failed = append(failed, name)
|
if transportErr != nil {
|
||||||
return
|
// 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 {
|
for _, tool := range pageTools {
|
||||||
merged = append(merged, mcpTool{Name: name + "_" + tool.Name, Description: tool.Description, InputSchema: tool.InputSchema})
|
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)
|
}(name)
|
||||||
}
|
}
|
||||||
wg.Wait()
|
wg.Wait()
|
||||||
|
|
||||||
if len(names) > 0 {
|
if len(dialed) > 0 {
|
||||||
attempted := make([]limitScope, len(names))
|
attempted := make([]limitScope, len(dialed))
|
||||||
for i, name := range names {
|
for i, name := range dialed {
|
||||||
attempted[i] = limitScope{kind: targetKindMCP, id: name}
|
attempted[i] = limitScope{kind: targetKindMCP, id: name}
|
||||||
}
|
}
|
||||||
g.limiter.countTargetRequests(attempted)
|
g.limiter.countTargetRequests(attempted)
|
||||||
@@ -1063,16 +1435,77 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, format mcpRespons
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name })
|
// F4: sort by (name, server) so a collision's outcome is deterministic
|
||||||
writeJSONRPCResult(w, format, req.ID, mcpToolsListResult{Tools: merged})
|
// — 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
|
// mcpToolCallParams is one "tools/call" request's params: name is the
|
||||||
// caller-facing, server-prefixed tool id federation itself minted in
|
// 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 {
|
type mcpToolCallParams struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Arguments json.RawMessage `json:"arguments,omitempty"`
|
Arguments json.RawMessage `json:"arguments,omitempty"`
|
||||||
|
Meta json.RawMessage `json:"_meta,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// resolveFederatedTool finds which of allowedNames is fullName's own
|
// 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
|
// out above: err != nil is a failure, a JSON-RPC-level resp.Error is
|
||||||
// not.
|
// not.
|
||||||
probeStart := time.Now()
|
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,
|
// G4: same dialed-server rule as mcpFederatedToolsList's own fan-out,
|
||||||
// above — ctx here derives from r.Context() too, so a client cancel
|
// above — ctx here derives from r.Context() too, so a client cancel
|
||||||
// (context.Canceled) must record nothing. Unlike the fan-out,
|
// (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)
|
g.limiter.countTargetRequest(targetKindMCP, serverName)
|
||||||
if err != nil {
|
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
|
// A response that hit mcpBackendCallResponseMaxBytes gets its own
|
||||||
// specific message, not the generic "upstream error" a truncated-
|
// specific message, not the generic "upstream error" a truncated-
|
||||||
// then-unparsable body would otherwise produce (security review
|
// 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",
|
name: "tools/list two-server aggregate, Accept: application/json",
|
||||||
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) {
|
buildTest: func(t *testing.T) (http.Handler, *http.Request) {
|
||||||
alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup", Description: "alpha's tool"}})
|
alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup", Description: "alpha's tool"}})
|
||||||
beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search", Description: "beta'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",
|
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}`,
|
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",
|
name: "malformed JSON",
|
||||||
contentType: "application/json",
|
contentType: "application/json",
|
||||||
@@ -147,6 +161,84 @@ func TestParseBackendJSONRPC(t *testing.T) {
|
|||||||
|
|
||||||
// --- end-to-end: POST /mcp ---
|
// --- 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:
|
// mockJSONRPCServer is a minimal MCP-shaped backend for federation tests:
|
||||||
// it answers "tools/list" with a fixed tool set and "tools/call" by
|
// it answers "tools/list" with a fixed tool set and "tools/call" by
|
||||||
// echoing back the resolved (prefix-stripped) name/arguments it was
|
// 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")
|
w.Header().Set("Content-Type", "application/json")
|
||||||
switch req.Method {
|
switch req.Method {
|
||||||
case "tools/list":
|
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})
|
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||||
case "tools/call":
|
case "tools/call":
|
||||||
var params mcpToolCallParams
|
var params mcpToolCallParams
|
||||||
@@ -244,7 +336,7 @@ func newSessionRequiredMockServer(t *testing.T, tools []mcpTool) *sessionRequire
|
|||||||
}
|
}
|
||||||
switch req.Method {
|
switch req.Method {
|
||||||
case "tools/list":
|
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})
|
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||||
case "tools/call":
|
case "tools/call":
|
||||||
var params mcpToolCallParams
|
var params mcpToolCallParams
|
||||||
@@ -523,7 +615,7 @@ func TestHandleMCPFederated_ToolsList_AggregatesPrefixesAndRespectsGroupAccess(t
|
|||||||
|
|
||||||
gotNames := make([]string, len(result.Tools))
|
gotNames := make([]string, len(result.Tools))
|
||||||
for i, tool := range result.Tools {
|
for i, tool := range result.Tools {
|
||||||
gotNames[i] = tool.Name
|
gotNames[i] = mustToolName(t, tool)
|
||||||
}
|
}
|
||||||
if len(gotNames) != len(tc.wantToolNames) {
|
if len(gotNames) != len(tc.wantToolNames) {
|
||||||
t.Fatalf("tools = %v, want %v", gotNames, 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 {
|
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||||
t.Fatalf("decode result: %v", err)
|
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)
|
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
|
// TestHandleMCPFederated_ToolsCall_RoutesToCorrectServer_StripsPrefix
|
||||||
// proves tools/call resolves the server prefix, strips it before
|
// proves tools/call resolves the server prefix, strips it before
|
||||||
// forwarding, contacts ONLY the resolved server (beta must never be
|
// 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
|
// TestHandleMCPFederated_ToolsCall_BackendError_RelayedToClient proves a
|
||||||
// backend's own JSON-RPC error (e.g. a tool execution failure) is relayed
|
// 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
|
// 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
|
// shapes and the one non-trigger shape (a network failure is covered
|
||||||
// separately below, since it needs a different mock shape entirely: no
|
// separately below, since it needs a different mock shape entirely: no
|
||||||
// server to answer at all).
|
// 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) {
|
func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
||||||
cases := []struct {
|
cases := []struct {
|
||||||
bareHandler http.HandlerFunc
|
bareHandler http.HandlerFunc
|
||||||
name string
|
name string
|
||||||
|
method string
|
||||||
wantHandshake bool
|
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) {
|
bareHandler: func(w http.ResponseWriter, r *http.Request) {
|
||||||
var req jsonrpcRequest
|
var req jsonrpcRequest
|
||||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||||
@@ -995,11 +1376,25 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
|||||||
wantHandshake: true,
|
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) {
|
bareHandler: func(w http.ResponseWriter, r *http.Request) {
|
||||||
w.WriteHeader(http.StatusForbidden)
|
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",
|
name: "HTTP 500 never triggers the handshake",
|
||||||
@@ -1008,9 +1403,43 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
|||||||
},
|
},
|
||||||
wantHandshake: false,
|
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 {
|
for _, tc := range cases {
|
||||||
t.Run(tc.name, func(t *testing.T) {
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
method := tc.method
|
||||||
|
if method == "" {
|
||||||
|
method = "tools/list"
|
||||||
|
}
|
||||||
var callCount atomic.Int32
|
var callCount atomic.Int32
|
||||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
if callCount.Add(1) == 1 {
|
if callCount.Add(1) == 1 {
|
||||||
@@ -1040,14 +1469,14 @@ func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
|||||||
t.Fatal("handler is not *Gateway")
|
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()
|
gotCalls := callCount.Load()
|
||||||
if tc.wantHandshake && gotCalls < 2 {
|
if tc.wantHandshake && gotCalls < 2 {
|
||||||
t.Errorf("calls = %d, want >=2 (bare attempt + initialize handshake)", gotCalls)
|
t.Errorf("calls = %d, want >=2 (bare attempt + initialize handshake)", gotCalls)
|
||||||
}
|
}
|
||||||
if !tc.wantHandshake && gotCalls != 1 {
|
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 {
|
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||||
t.Fatalf("decode result: %v", err)
|
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)
|
t.Fatalf("tools = %+v, want exactly alpha_search", result.Tools)
|
||||||
}
|
}
|
||||||
|
|
||||||
seq := strict.callSequence()
|
seq := strict.callSequence()
|
||||||
if len(seq) != 4 {
|
// F5, review round 3, 2026-09: notifications/initialized is now sent
|
||||||
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)
|
// 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=" {
|
if seq[0] != "tools/list:session=" {
|
||||||
t.Errorf("call[0] = %q, want a bare tools/list attempt with no session header", seq[0])
|
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=" {
|
if seq[1] != "initialize:session=" {
|
||||||
t.Errorf("call[1] = %q, want an initialize handshake with no session header", seq[1])
|
t.Errorf("call[1] = %q, want an initialize handshake with no session header", seq[1])
|
||||||
}
|
}
|
||||||
if seq[2] != "tools/list:session=sess-1" {
|
if seq[2] != "notifications/initialized:session=sess-1" {
|
||||||
t.Errorf("call[2] = %q, want the retried tools/list carrying the session header the handshake returned", seq[2])
|
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" {
|
if seq[3] != "tools/list:session=sess-1" {
|
||||||
t.Errorf("call[3] = %q, want the best-effort session close carrying the same session header", seq[3])
|
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 {
|
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||||
t.Fatalf("decode result: %v", err)
|
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)
|
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", true},
|
||||||
{"http://playwright.internal:8080/sse?foo=bar", 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", false},
|
||||||
{"http://mcp.internal:8080/mcp", false},
|
{"http://mcp.internal:8080/mcp", false},
|
||||||
|
{"http://mcp.internal:8080/mcp/", false},
|
||||||
{"http://mcp.internal:8080/sse/extra", false},
|
{"http://mcp.internal:8080/sse/extra", false},
|
||||||
{"://not a valid url", false}, // malformed -> safe default: not excluded
|
{"://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 {
|
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||||
t.Fatalf("decode result: %v", err)
|
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)
|
t.Fatalf("tools = %+v, want exactly alpha_lookup (playwright excluded)", result.Tools)
|
||||||
}
|
}
|
||||||
if legacy.called.Load() != 0 {
|
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) == "":
|
case req.Method == "tools/list" && r.Header.Get(mcpSessionHeader) == "":
|
||||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}})
|
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}})
|
||||||
default:
|
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})
|
_ = 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 {
|
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||||
t.Fatalf("decode result: %v", err)
|
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)
|
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 ---
|
// --- feat/target-health: passive recording from federation traffic ---
|
||||||
|
|
||||||
// TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer proves
|
// TestHandleMCPFederated_ToolsList_TargetHealth_RecordsPerServer proves
|
||||||
@@ -2187,6 +2738,31 @@ func TestHandleMCPFederated_ToolsList_NeverDialedSemaphoreStarved_RecordsNothing
|
|||||||
if neverDialed != 1 {
|
if neverDialed != 1 {
|
||||||
t.Errorf("neverDialed = %d, want exactly 1", neverDialed)
|
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
|
// TestHandleMCPFederated_ToolsCall_ClientCanceled_RecordsNothing mirrors
|
||||||
|
|||||||
Reference in New Issue
Block a user