diff --git a/README.md b/README.md
index d8ccaf6..95a3c43 100644
--- a/README.md
+++ b/README.md
@@ -42,7 +42,7 @@ Catalog — see [Development](#development) for what that needs.
```mermaid
flowchart LR
Req["incoming request"] --> Route{"route
ServeHTTP: path+method dispatch"}
- Route -->|"a recognized route:
/v1/models, /v1/chat/completions,
/v1/embeddings, /v1/mcp/servers,
/v1/agents, /mcp/name/...,
/a2a/name/..., or /providerName/...
for a known provider"| Auth["auth (per branch)
identify(): key digest lookup"]
+ Route -->|"a recognized route:
/v1/models, /v1/chat/completions,
/v1/embeddings, /v1/mcp/servers,
/v1/agents, /mcp (federated),
/mcp/name/...,
/a2a/name/..., or /providerName/...
for a known provider"| Auth["auth (per branch)
identify(): key digest lookup"]
Route -->|"no route matches"| Fallback["passthroughUnknown ? next : 404
never calls identify()"]
Auth -->|"no match"| E401["401 authentication_error"]
Auth -->|"user, group"| Branch{"branch"}
@@ -50,21 +50,26 @@ flowchart LR
Branch --> Unified["POST /v1/chat/completions
POST /v1/embeddings"]
Branch --> Registry["GET /v1/mcp/servers, /v1/agents"]
Branch --> MCP["MCP/A2A proxy"]
+ Branch --> Federated["POST /mcp (federated JSON-RPC)"]
Branch --> Passthrough["native passthrough"]
Unified --> Resolve["resolve model
registry.resolve"]
Resolve -->|"unknown/denied"| E404["404 / 403"]
Resolve --> Limit1["limiter.checkAndCount"]
MCP --> Limit2["limiter.checkAndCount"]
+ Federated --> Limit4["limiter.checkAndCount"]
Passthrough --> Limit3["limiter.checkAndCount"]
Limit1 -->|"over budget"| E429["429 / 503"]
Limit2 -->|"over budget"| E429
Limit3 -->|"over budget"| E429
+ Limit4 -->|"over budget"| E429
Limit1 --> Exec["execute: translate + call adapter"]
Limit2 --> ExecProxy["execute: reverse proxy"]
Limit3 --> ExecProxy
+ Limit4 --> ExecFed["execute: answer locally, or
JSON-RPC call(s) to backend server(s)"]
Exec --> Account["account: usage, cost"]
Account --> Resp["respond to client"]
ExecProxy --> Resp
+ ExecFed --> Resp
Models --> Resp
Registry --> Resp
```
@@ -500,8 +505,9 @@ http:
pushed the scope over.
- **A synthetic total scope** (`kind: "total"`, `id: "all"`) is counted
alongside every request's own user/group scopes on every metered route
- (unified chat/embeddings, the media endpoints, native passthrough, and
- the MCP/A2A proxy) — the sum of all LLM traffic combined. It carries no
+ (unified chat/embeddings, the media endpoints, native passthrough, the
+ per-server MCP/A2A proxy, and the federated `/mcp` endpoint) — the sum
+ of all LLM traffic combined. It carries no
limit and is never evaluated: no configuration can throttle it. The
admin usage API's `total` row and every `scope=total` usage-history
query read this scope.
@@ -824,13 +830,27 @@ or in CI.
`{"scope","metric","window","points":[{"bucket":"2026082114","value":123},...]}`.
The Charts view fetches this once per selection change, plus a 30s
auto-refresh of the current selection.
+- **`GET /admin/api/targets`** returns every configured MCP server and
+ agent for the dashboard's "MCP & Agents" tab:
+ `{"mcpServers":[...],"agents":[...]}`, each entry
+ `{"name","url","access","counters"}`. `url` has any userinfo/query
+ string stripped, same as a provider's `baseUrl` in `overview`. `access`
+ is the list of group names actually allowed to reach that target,
+ computed via the identical glob match `mcpServers`/`agents`
+ authorization itself uses (`GroupConfig.mcpServers`/`agents`), so this
+ view can never disagree with what the proxy enforces; omitted when
+ every configured group can reach it. `counters` is
+ `requestsPerMinute`/`requestsPerDay`/`requestsPerMonth` — requests
+ only, no tokens or cost (see [MCP and A2A](#mcp-and-a2a) for why). Polled
+ every 5 seconds, in the same batch as `overview` and `usage`.
- **What's exposed**: provider names, types, base URLs (with any
userinfo/query string stripped before it's ever echoed), model counts,
- discovery status, group/user names, membership, limits, and live usage
- counters. **Never exposed**: API keys (not even digests), provider
- keys, the Redis password, or users-file path contents — every
- secret-bearing field is redacted from every response.
-- **Admin traffic is never counted**: none of the three `/admin/api/*`
+ discovery status, group/user names, membership, limits, MCP/agent
+ target names/URLs/access, and live usage counters. **Never exposed**:
+ API keys (not even digests), provider keys, the Redis password, or
+ users-file path contents — every secret-bearing field is redacted from
+ every response.
+- **Admin traffic is never counted**: none of the four `/admin/api/*`
JSON routes call `checkAndCount` — admin polling never moves any
user's or group's `requestsPerMinute`/`requestsPerDay` counters, and
usage statistics reflect real LLM traffic only (operator directive: an
@@ -838,18 +858,18 @@ or in CI.
never itself distort the numbers it displays). The direct consequence:
an admin's own `requestsPerMinute`/`requestsPerDay` limit, if
configured, is never enforced against admin-route traffic either — an
- admin key holder can poll any of the three routes as fast as they like.
+ admin key holder can poll any of the four routes as fast as they like.
This is an accepted trade-off, not an oversight: these are admin-gated,
cheap reads, and an admin holder polling aggressively is a
self-inflicted, not a shared, resource cost. `GET /admin` and
`GET /admin/assets/*` count nothing either, for the simpler reason that
they are unauthenticated — there is no identified user to count a
request against.
-- **Response headers**: all three JSON routes set
+- **Response headers**: all four JSON routes set
`X-Content-Type-Options: nosniff` and `Cache-Control: no-store`; every
hashed asset sets `X-Content-Type-Options: nosniff` and its own
- immutable `Cache-Control` (above). Every one of the four routes — `GET
- /admin`, every hashed asset, and the three JSON routes — shares one
+ immutable `Cache-Control` (above). Every one of these routes — `GET
+ /admin`, every hashed asset, and the four JSON routes — shares one
`Content-Security-Policy` header (`default-src 'none'; script-src
'self'; style-src 'self'; connect-src 'self'; img-src 'self' data:;
frame-ancestors 'none'; base-uri 'none'; form-action 'none'`) — no
@@ -880,10 +900,50 @@ or in CI.
- Request-rate limits (`requestsPerMinute`/`requestsPerDay`) apply to
target-proxy calls the same as everywhere else; token/cost metrics do
not, since there is no usage to parse from an arbitrary MCP/A2A
- response.
+ response. Every proxied call (per-server proxy or federated `/mcp`) is
+ additionally counted against its own target's request-rate counters —
+ requests only, no configurable limit of its own this round — visible via
+ `GET /admin/api/targets` above.
- Group visibility (`mcpServers`/`agents` glob lists) gates both the
registry listing and the proxy path itself — a denied name is a 403 at
the proxy, not just hidden from the listing.
+- **Federated `POST /mcp`**: a single JSON-RPC 2.0 endpoint aggregating
+ every MCP server the caller's group can reach — for clients built
+ against a single aggregated endpoint (the earlier, pre-plugin gateway
+ this replaces) rather than the per-server `/mcp/{name}/...` proxy above.
+ Unlike that proxy, this is not a byte-level reverse proxy: the gateway
+ decodes the request itself, answers `initialize`/`ping` locally, and for
+ `tools/list`/`tools/call` issues its own outbound JSON-RPC call(s) to
+ the relevant server(s) as part of handling that one request — no
+ session state is cached across separate calls (per-request upstream
+ sessions, matching the per-server proxy's own complete absence of
+ gateway-side session tracking).
+ - `tools/list` fans out to every allowed server concurrently and merges
+ the results, prefixing each tool's name `"_"` —
+ the same underscore-separator convention the earlier agentgateway used
+ and that existing MCP clients (pugbot's `mcpclient`, agentkit) already
+ persist in their own tool-id records (e.g.
+ `brave-search_brave_web_search`). A server that errors or is
+ unreachable is skipped, not surfaced as a whole-call failure — the
+ aggregate degrades to every other server's tools.
+ - `tools/call` resolves the target server by the LONGEST matching
+ `"_"` prefix against the caller's allowed servers (so two
+ configured servers where one name prefixes the other, e.g. `foo` and
+ `foo_bar`, resolve unambiguously), strips the prefix, forwards the
+ call, and relays the result under the caller's own JSON-RPC `id`. An
+ unresolvable prefix — including one that names a real but
+ group-restricted server — is a JSON-RPC `invalid params` error
+ (`-32602`), not an HTTP `404`: the request reached a real route and
+ method, it just named a tool nothing could route.
+ - `notifications/*` gets a `202 Accepted` with no body (no response is
+ ever due for a JSON-RPC notification); any other method is a JSON-RPC
+ `method not found` (`-32601`). Malformed JSON is a JSON-RPC `parse
+ error` (`-32700`) with `id: null`, per spec.
+ - `checkAndCount` runs once per request (the caller's own user/group/
+ total scopes); each backend server actually contacted is additionally
+ attributed its own per-target counters — see `GET /admin/api/targets`
+ above and the request-rate note below. Applies to MCP servers only,
+ not agents: there is no equivalent aggregated `/a2a` endpoint.
## Security notes
@@ -940,7 +1000,13 @@ or in CI.
though, a translate-time 400 (an unsupported parameter combination) or
Anthropic's embeddings 501 still count toward the caller's request-rate
limit — the limiter already ran by then — though neither ever bills a
- token or cost, since usage stays at zero either way.
+ token or cost, since usage stays at zero either way. The federated
+ `/mcp` endpoint follows the same shape: an unreadable body, an
+ unsupported `Upgrade`, or a malformed JSON-RPC envelope are all rejected
+ before the limiter runs, but a JSON-RPC-level error discovered AFTER
+ dispatch — `method not found`, an unresolvable `tools/call` prefix —
+ still counts, since `checkAndCount` already ran by the time the method
+ is even looked at.
- **Gateway-to-client responses on passthrough and MCP/A2A routes are
identity-encoded.** The gateway strips a client's own `Accept-Encoding`
header before forwarding to the upstream, so Go's `http.Transport` adds
diff --git a/llmgateway.go b/llmgateway.go
index d091535..781cc95 100644
--- a/llmgateway.go
+++ b/llmgateway.go
@@ -449,6 +449,23 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return
}
+ // Federated /mcp (Feature C, v0.21) is matched ahead of the per-server
+ // "/mcp/{name}/..." target proxy below: targetRoute's own splitting
+ // already rejects a bare "/mcp" (no name segment to route to), so
+ // there is no ambiguity between the two, but this ordering keeps every
+ // exact-path route grouped together (mirroring "/v1/mcp/servers" and
+ // "/v1/agents" immediately above) ahead of the prefix-matched one.
+ if r.Method == http.MethodPost && r.URL.Path == federatedMCPPath {
+ u, grp, ok := g.auth.identify(r)
+ g.logAuthEvent(ok, authEventUserName(u), r)
+ if !ok {
+ writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key")
+ return
+ }
+ g.handleMCPFederated(sw, r, u, grp)
+ return
+ }
+
if name, rest, ok := targetRoute(r.URL.EscapedPath(), targetKindMCP); ok {
u, grp, authOK := g.auth.identify(r)
g.logAuthEvent(authOK, authEventUserName(u), r)
diff --git a/mcp_federation.go b/mcp_federation.go
new file mode 100644
index 0000000..b49cf41
--- /dev/null
+++ b/mcp_federation.go
@@ -0,0 +1,490 @@
+package traefikllmgateway
+
+import (
+ "bytes"
+ "context"
+ "encoding/json"
+ "io"
+ "net/http"
+ "sort"
+ "strings"
+ "sync"
+)
+
+// federatedMCPPath is the bare "/mcp" route this file implements —
+// distinct from the per-server "/mcp/{name}/..." target proxy
+// (targetRoute/handleTargetProxy, mcp_a2a.go), which rejects a bare "/mcp"
+// or "/mcp/" outright (targetRoute's own doc comment: "carries no target
+// name to route to"). Two real consumers (pugbot's mcpclient, agentkit)
+// were built against the earlier, pre-plugin gateway's single aggregated
+// MCP endpoint and reach every configured MCP server through ONE URL —
+// this route restores that shape on top of the plugin's per-server
+// config, rather than forcing those consumers into a lossy
+// single-server-at-a-time workaround.
+const federatedMCPPath = "/mcp"
+
+// JSON-RPC 2.0 envelope fields and standard error codes (jsonrpc.org's
+// spec §5.1) this file's own responses use. MCP's wire format is JSON-RPC
+// 2.0 verbatim — no plugin-specific envelope on top.
+const (
+ jsonrpcVersion = "2.0"
+ jsonrpcParseError = -32700
+ jsonrpcInvalidParams = -32602
+ jsonrpcMethodNotFound = -32601
+ jsonrpcInternalError = -32603
+)
+
+// jsonrpcRequest is one JSON-RPC 2.0 request/notification, decoded from
+// the client's POST body and re-encoded (with a synthetic id) for this
+// gateway's own outbound calls to a backend MCP server (mcpBackendCall).
+// ID is nil for a notification (no "id" key in the wire form) — the
+// signal notifications/* dispatch checks for a response-shaped request.
+// Only a single JSON-RPC object is supported, never a batch array: MCP's
+// own Streamable HTTP transport is the target wire format here, and
+// nothing in this round's brief calls for JSON-RPC batch support.
+type jsonrpcRequest struct {
+ Params json.RawMessage `json:"params,omitempty"`
+ JSONRPC string `json:"jsonrpc"`
+ Method string `json:"method"`
+ ID json.RawMessage `json:"id,omitempty"`
+}
+
+// jsonrpcError is a JSON-RPC 2.0 error object.
+type jsonrpcError struct {
+ Message string `json:"message"`
+ Code int `json:"code"`
+}
+
+// jsonrpcResponse is one JSON-RPC 2.0 response, used both to decode a
+// backend MCP server's reply (mcpBackendCall) and to encode this
+// gateway's own reply to the client — Result is left as json.RawMessage
+// deliberately (never `any`) so a backend's raw result can be relayed to
+// the client verbatim (mcpFederatedToolsCall) without a decode-then-
+// re-encode round trip, and so building a fresh result (mcpFederatedInitialize,
+// mcpFederatedToolsList) is one explicit json.Marshal away rather than
+// implicit. ID has no "omitempty": JSON-RPC 2.0 requires an "id" key on
+// every non-notification response, explicitly null when the request's own
+// id could not even be determined (a parse error) — omitting the key
+// entirely in that case would be spec-incorrect, not just unusual.
+type jsonrpcResponse struct {
+ Error *jsonrpcError `json:"error,omitempty"`
+ JSONRPC string `json:"jsonrpc"`
+ Result json.RawMessage `json:"result,omitempty"`
+ ID json.RawMessage `json:"id"`
+}
+
+// writeJSONRPCEnvelope writes resp as the response body, matching
+// setAdminJSONHeaders' own "declared Content-Type, nothing else asserted"
+// minimalism — a JSON-RPC error is still carried at HTTP 200 (JSON-RPC
+// errors are a payload-level concept, not a transport-level one; every
+// genuinely transport-level failure this route can hit — an unreadable
+// body, an unauthenticated/forbidden caller, a rate limit — is handled
+// before this function is ever reached, via the gateway's own
+// writeOAIError/writeLimitViolation envelope, matching every other route
+// in this package).
+func writeJSONRPCEnvelope(w http.ResponseWriter, resp jsonrpcResponse) {
+ resp.JSONRPC = jsonrpcVersion
+ w.Header().Set("Content-Type", "application/json")
+ _ = json.NewEncoder(w).Encode(resp) // headers already committed; nothing useful to do on encode failure
+}
+
+// writeJSONRPCResult marshals result and writes it as a successful
+// JSON-RPC response under id. A marshal failure (result is not
+// JSON-marshalable — never true for any value this file actually passes,
+// but defensively handled rather than panicking) reports itself as a
+// JSON-RPC internal error instead.
+func writeJSONRPCResult(w http.ResponseWriter, id json.RawMessage, result any) {
+ resultBytes, err := json.Marshal(result)
+ if err != nil {
+ writeJSONRPCErrorResponse(w, id, jsonrpcInternalError, "failed to encode result")
+ return
+ }
+ writeJSONRPCEnvelope(w, jsonrpcResponse{ID: id, Result: resultBytes})
+}
+
+// writeJSONRPCErrorResponse writes a JSON-RPC error response under id.
+func writeJSONRPCErrorResponse(w http.ResponseWriter, id json.RawMessage, code int, message string) {
+ writeJSONRPCEnvelope(w, jsonrpcResponse{ID: id, Error: &jsonrpcError{Code: code, Message: message}})
+}
+
+// allowedMCPServerNames returns the sorted names of every MCP server
+// grp.allowsMCP permits — the SAME per-server authorization check
+// handleTargetProxy enforces (mcp_a2a.go), applied here across the whole
+// catalog: tools/list aggregation only ever fans out to these names,
+// and tools/call prefix resolution (resolveFederatedTool) only ever
+// searches them — so a group restricted to N servers can see or invoke
+// only those N servers' tools through POST /mcp too, exactly as it can
+// only reach them through POST /mcp/{name} directly.
+func allowedMCPServerNames(cfg *Config, grp *group) []string {
+ names := make([]string, 0, len(cfg.MCPServers))
+ for name := range cfg.MCPServers {
+ if grp.allowsMCP(name) {
+ names = append(names, name)
+ }
+ }
+ sort.Strings(names)
+ return names
+}
+
+// handleMCPFederated implements POST /mcp (federatedMCPPath): the
+// federated JSON-RPC endpoint this file's own doc comment explains the
+// need for.
+//
+// Unlike handleTargetProxy, this is not a byte-level reverse proxy — it
+// cannot be: aggregating N servers' tools/list results into one response,
+// or routing one tools/call to the single right server, both require
+// actually parsing the JSON-RPC envelope, which a blind proxy never does.
+// The gateway terminates the client's JSON-RPC exchange itself: it
+// answers "initialize" and "ping" locally (mcpFederatedInitialize; ping
+// needs no backend at all) and, for "tools/list"/"tools/call", issues its
+// OWN outbound JSON-RPC call(s) to the relevant backend server(s) as part
+// of handling THIS one incoming request — no session state is ever
+// persisted across separate client requests (mcpBackendCall's own doc
+// comment), matching how the per-server proxy today carries no
+// gateway-side session state of its own either (it is a pure byte pipe;
+// whatever MCP session semantics exist are negotiated directly between
+// the client and that one backend, via headers proxyUpstream forwards
+// unmodified — see handleTargetProxy, mcp_a2a.go). A response can never be
+// a raw byte relay here, even for a single-target call (tools/call): every
+// outbound backend call carries its own synthetic id (mcpBackendCall), so
+// the client's own id must always be decoded and re-attached on the way
+// back out — this is the concrete reason "relay the response" cannot mean
+// what it means for handleTargetProxy.
+//
+// checkAndCount runs once per incoming request, before method dispatch —
+// covering the caller's own user/group/total scopes exactly like every
+// other metered route — with per-target attribution (countTargetRequest)
+// added separately, per backend server actually contacted, inside
+// mcpFederatedToolsList/mcpFederatedToolsCall: a single tools/list call
+// may contact several servers, so "one request in, one target scope
+// incremented" does not hold here the way it does for handleTargetProxy.
+func (g *Gateway) handleMCPFederated(w http.ResponseWriter, r *http.Request, u *user, grp *group) {
+ if r.Header.Get("Upgrade") != "" {
+ writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported")
+ return
+ }
+
+ body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes))
+ if err != nil {
+ writeOAIError(w, http.StatusBadRequest, "invalid_request_error", "cannot read request body")
+ return
+ }
+
+ var req jsonrpcRequest
+ if err := json.Unmarshal(body, &req); err != nil {
+ // A malformed envelope carries no reliable id at all — JSON-RPC
+ // 2.0 §5.1 requires "id":null for exactly this case, not an
+ // omitted key (jsonrpcResponse.ID's own doc comment).
+ writeJSONRPCErrorResponse(w, json.RawMessage("null"), jsonrpcParseError, "parse error")
+ return
+ }
+
+ scopes := withTotalScope(buildLimitScopes(u, grp))
+ if violation := g.limiter.checkAndCount(scopes); violation != nil {
+ writeLimitViolation(w, violation)
+ return
+ }
+
+ switch {
+ case req.Method == "initialize":
+ g.mcpFederatedInitialize(w, req)
+ case req.Method == "ping":
+ writeJSONRPCResult(w, req.ID, map[string]any{})
+ case req.Method == "tools/list":
+ g.mcpFederatedToolsList(w, r, req, grp)
+ case req.Method == "tools/call":
+ g.mcpFederatedToolsCall(w, r, req, grp)
+ case strings.HasPrefix(req.Method, "notifications/"):
+ // A JSON-RPC notification carries no id and gets no response body
+ // by definition; MCP's Streamable HTTP transport answers a
+ // notification POST with 202 Accepted and nothing else.
+ w.WriteHeader(http.StatusAccepted)
+ default:
+ writeJSONRPCErrorResponse(w, req.ID, jsonrpcMethodNotFound, "method not found: "+req.Method)
+ }
+}
+
+// mcpInitializeParams is the one field of MCP's "initialize" request
+// params this gateway reads.
+type mcpInitializeParams struct {
+ ProtocolVersion string `json:"protocolVersion"`
+}
+
+// defaultMCPProtocolVersion is reported only when the caller's own
+// "initialize" request omits protocolVersion entirely. This is a
+// fallback default this plugin picks for its own local response, not a
+// verified compatibility claim about any downstream MCP server's
+// supported version — see mcpFederatedInitialize's own doc comment for
+// why the caller's own value is preferred whenever it sends one.
+const defaultMCPProtocolVersion = "2025-06-18"
+
+// mcpFederatedInitialize answers "initialize" locally — no backend server
+// is ever contacted for it (handleMCPFederated's own doc comment).
+// protocolVersion echoes the caller's own requested value when present:
+// version negotiation is the caller's choice to make, and echoing it back
+// avoids this plugin hardcoding a claim about which dated MCP protocol
+// revision it actually tracks. capabilities.tools is present (empty
+// object, no sub-fields) to declare tool-calling support without
+// overclaiming listChanged notifications this gateway never sends.
+func (g *Gateway) mcpFederatedInitialize(w http.ResponseWriter, req jsonrpcRequest) {
+ var params mcpInitializeParams
+ _ = json.Unmarshal(req.Params, ¶ms) // best-effort; empty/malformed params falls back to defaultMCPProtocolVersion below
+
+ protocolVersion := params.ProtocolVersion
+ if protocolVersion == "" {
+ protocolVersion = defaultMCPProtocolVersion
+ }
+
+ writeJSONRPCResult(w, req.ID, map[string]any{
+ "protocolVersion": protocolVersion,
+ "capabilities": map[string]any{
+ "tools": map[string]any{},
+ },
+ "serverInfo": map[string]any{
+ "name": "traefik-llmgateway",
+ "version": pluginVersion,
+ },
+ })
+}
+
+// mcpTool is one tool entry in an MCP "tools/list" result, either as read
+// back from a backend server's own response or as re-emitted (name
+// prefixed) in this gateway's federated aggregate.
+type mcpTool struct {
+ Name string `json:"name"`
+ Description string `json:"description,omitempty"`
+ InputSchema json.RawMessage `json:"inputSchema,omitempty"`
+}
+
+// mcpToolsListResult is the "result" object of an MCP "tools/list"
+// response, both a backend's own and this gateway's federated one.
+type mcpToolsListResult struct {
+ Tools []mcpTool `json:"tools"`
+}
+
+// mcpBackendCall issues one JSON-RPC 2.0 request to targetURL — an MCP
+// server's own configured base URL, the same single endpoint
+// handleTargetProxy reverse-proxies every method to — and returns its
+// parsed response. It carries a synthetic id ("federated"), never the
+// client's own: tools/list may issue this call several times over for one
+// incoming client request (once per allowed server), so reusing the
+// client's id across N independent outbound calls would be meaningless,
+// and the id on the way back is discarded by every caller of this
+// function either way — only Result/Error is read.
+//
+// No outbound "initialize" handshake precedes this call, and none of its
+// own session/cookie state is cached across separate mcpBackendCall
+// invocations — deliberately: per-request, throwaway upstream calls are
+// what "keep upstream sessions per-request" (handleMCPFederated's own doc
+// comment) means in practice here, matching the per-server proxy's own
+// complete absence of gateway-side session tracking. A backend MCP server
+// that hard-requires a prior initialize handshake before honoring
+// tools/list or tools/call will reject this call; that is a real,
+// documented trade-off of this design (see README's federation section),
+// not an oversight.
+//
+// A backend may answer either a bare JSON object (Content-Type:
+// application/json) or an MCP Streamable HTTP SSE stream (Content-Type:
+// text/event-stream) carrying one or more "data:" events — both are
+// accepted (parseBackendJSONRPC).
+func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) {
+ paramsRaw, err := json.Marshal(params)
+ if err != nil {
+ return nil, err
+ }
+ reqBody, err := json.Marshal(jsonrpcRequest{JSONRPC: jsonrpcVersion, ID: json.RawMessage(`"federated"`), Method: method, Params: paramsRaw})
+ if err != nil {
+ return nil, err
+ }
+
+ // gosec G704 (SSRF via taint analysis): targetURL is an
+ // operator-configured MCP server base (Config.MCPServers[name].URL,
+ // validated at construction by validateTargetURLs) — no client-
+ // supplied path segment or query ever reaches this call, unlike
+ // proxyUpstream's own upstreamURL (routes_passthrough.go), which
+ // deliberately does forward client-supplied path material and
+ // documents that same taint-analysis false positive at its own call
+ // site.
+ httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(reqBody)) //nolint:gosec // operator-configured target URL, no client-supplied path/query
+ if err != nil {
+ return nil, err
+ }
+ httpReq.Header.Set("Content-Type", "application/json")
+ httpReq.Header.Set("Accept", "application/json, text/event-stream")
+
+ resp, err := g.targetClient.Do(httpReq)
+ if err != nil {
+ return nil, err
+ }
+ defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure
+
+ respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxRequestBytes))
+ if err != nil {
+ return nil, err
+ }
+ return parseBackendJSONRPC(resp.Header.Get("Content-Type"), respBody)
+}
+
+// parseBackendJSONRPC decodes body as a backend MCP server's JSON-RPC
+// response, accepting either shape the MCP Streamable HTTP transport
+// allows a POST response to take: a bare JSON object (any other
+// Content-Type is treated as this default case), or an SSE stream
+// (Content-Type: text/event-stream) of one or more "data:" lines — the
+// LAST such line is parsed, matching a minimal backend that emits exactly
+// one data event carrying its final response.
+func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, error) {
+ if strings.Contains(strings.ToLower(contentType), "text/event-stream") {
+ body = lastSSEDataLine(body)
+ }
+ var out jsonrpcResponse
+ if err := json.Unmarshal(body, &out); err != nil {
+ return nil, err
+ }
+ return &out, nil
+}
+
+// lastSSEDataLine returns the payload of the last "data:" line in an SSE
+// body — body itself, unchanged, if it contains no such line at all (so
+// the caller's own json.Unmarshal fails informatively on whatever body
+// actually was, rather than this function silently swallowing a
+// malformed stream into an empty byte slice).
+func lastSSEDataLine(body []byte) []byte {
+ var last []byte
+ for _, line := range bytes.Split(body, []byte("\n")) {
+ line = bytes.TrimRight(line, "\r")
+ if data, ok := bytes.CutPrefix(line, []byte("data: ")); ok {
+ last = data
+ } else if data, ok := bytes.CutPrefix(line, []byte("data:")); ok {
+ last = bytes.TrimSpace(data)
+ }
+ }
+ if last == nil {
+ return body
+ }
+ return last
+}
+
+// mcpFederatedToolsList implements the federated "tools/list": it fans
+// out one mcpBackendCall per server allowedMCPServerNames(grp) permits,
+// concurrently, then merges every reachable server's own tools into one
+// result — each tool's name prefixed "_", the exact
+// underscore-separator convention the old agentgateway this plugin
+// replaces used, and the one pugbot's mcpclient and agentkit have already
+// persisted into their own tool-id databases (e.g.
+// "brave-search_brave_web_search" — verified live). A server that errors,
+// times out, or returns an unparsable result is logged and skipped, not
+// surfaced as a whole-call failure: a federated aggregate degrading to
+// "every OTHER server's tools" when one upstream is down is the useful
+// behavior an aggregator should have, matching how a client would want a
+// federated view to behave in practice.
+//
+// countTargetRequest attributes one request to EVERY server actually
+// contacted here (handleMCPFederated's own doc comment) — a single
+// federated tools/list call can move several targets' own counters, not
+// just one.
+func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) {
+ names := allowedMCPServerNames(g.cfg, grp)
+
+ var (
+ mu sync.Mutex
+ wg sync.WaitGroup
+ merged = make([]mcpTool, 0, len(names))
+ failed []string
+ )
+ for _, name := range names {
+ wg.Add(1)
+ go func(name string) {
+ defer wg.Done()
+ targetURL := g.cfg.MCPServers[name].URL
+ resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/list", struct{}{})
+ g.limiter.countTargetRequest(targetKindMCP, name)
+
+ mu.Lock()
+ defer mu.Unlock()
+ if err != nil || resp.Error != nil {
+ failed = append(failed, name)
+ return
+ }
+ var result mcpToolsListResult
+ if err := json.Unmarshal(resp.Result, &result); err != nil {
+ failed = append(failed, name)
+ return
+ }
+ for _, tool := range result.Tools {
+ merged = append(merged, mcpTool{Name: name + "_" + tool.Name, Description: tool.Description, InputSchema: tool.InputSchema})
+ }
+ }(name)
+ }
+ wg.Wait()
+
+ if len(failed) > 0 {
+ sort.Strings(failed)
+ g.logf("federated tools/list: %d of %d allowed MCP server(s) failed and were skipped: %v", len(failed), len(names), failed)
+ }
+ sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name })
+
+ writeJSONRPCResult(w, req.ID, mcpToolsListResult{Tools: merged})
+}
+
+// mcpToolCallParams is one "tools/call" request's params: name is the
+// caller-facing, server-prefixed tool id federation itself minted in
+// tools/list ("_").
+type mcpToolCallParams struct {
+ Name string `json:"name"`
+ Arguments json.RawMessage `json:"arguments,omitempty"`
+}
+
+// resolveFederatedTool finds which of allowedNames is fullName's own
+// server prefix, matched by LONGEST "_" prefix — disambiguates
+// two configured, allowed servers whose names are themselves prefixes of
+// one another (e.g. both "foo" and "foo_bar" configured: a tool id
+// "foo_bar_lookup" must resolve to server "foo_bar"'s tool "lookup", not
+// server "foo"'s tool "bar_lookup"). ok is false when no allowed server's
+// prefix matches at all — including a tool that belongs to a real,
+// configured server the caller's group simply cannot reach, since
+// allowedNames (allowedMCPServerNames) already excludes it; this is what
+// makes tools/call authz-aware without a second, separate access check.
+func resolveFederatedTool(fullName string, allowedNames []string) (serverName, toolName string, ok bool) {
+ best := ""
+ for _, name := range allowedNames {
+ if strings.HasPrefix(fullName, name+"_") && len(name) > len(best) {
+ best = name
+ }
+ }
+ if best == "" {
+ return "", "", false
+ }
+ return best, strings.TrimPrefix(fullName, best+"_"), true
+}
+
+// mcpFederatedToolsCall resolves req's "name" param to one allowed MCP
+// server (resolveFederatedTool), forwards the call there with the server
+// prefix stripped, and relays its result or error back to the client
+// under the CLIENT's own id — never the synthetic one mcpBackendCall's
+// outbound call used. An unresolvable prefix is a JSON-RPC invalid-params
+// error, not an HTTP 404: the request reached a real method on a real
+// route, it just named a tool nothing federation could route (an
+// unconfigured server, or one this caller's group cannot reach).
+func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) {
+ var params mcpToolCallParams
+ if err := json.Unmarshal(req.Params, ¶ms); err != nil {
+ writeJSONRPCErrorResponse(w, req.ID, jsonrpcInvalidParams, "invalid params")
+ return
+ }
+
+ serverName, toolName, ok := resolveFederatedTool(params.Name, allowedMCPServerNames(g.cfg, grp))
+ if !ok {
+ writeJSONRPCErrorResponse(w, req.ID, jsonrpcInvalidParams, "unknown tool: "+params.Name)
+ return
+ }
+
+ targetURL := g.cfg.MCPServers[serverName].URL
+ resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments})
+ g.limiter.countTargetRequest(targetKindMCP, serverName)
+ if err != nil {
+ g.logf("federated tools/call: server %q: %v", serverName, err)
+ writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "upstream error")
+ return
+ }
+ writeJSONRPCEnvelope(w, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error})
+}
diff --git a/mcp_federation_test.go b/mcp_federation_test.go
new file mode 100644
index 0000000..4b938a8
--- /dev/null
+++ b/mcp_federation_test.go
@@ -0,0 +1,705 @@
+package traefikllmgateway
+
+import (
+ "bytes"
+ "context"
+ "encoding/json"
+ "net/http"
+ "net/http/httptest"
+ "sync/atomic"
+ "testing"
+ "time"
+)
+
+// --- unit: resolveFederatedTool ---
+
+// TestResolveFederatedTool covers good/bad/edge cases: an unambiguous
+// match, the longest-prefix disambiguation between two configured server
+// names where one is a prefix of the other, and an unresolvable prefix
+// (unconfigured server, or one simply not in the allowed set — the same
+// thing from this function's point of view).
+func TestResolveFederatedTool(t *testing.T) {
+ cases := []struct {
+ name string
+ fullName string
+ wantServerName string
+ wantToolName string
+ allowedNames []string
+ wantOK bool
+ }{
+ {
+ name: "unambiguous single server match",
+ fullName: "brave-search_brave_web_search",
+ allowedNames: []string{"brave-search"},
+ wantServerName: "brave-search",
+ wantToolName: "brave_web_search",
+ wantOK: true,
+ },
+ {
+ name: "longest prefix wins when one server name prefixes another",
+ fullName: "foo_bar_lookup",
+ allowedNames: []string{"foo", "foo_bar"},
+ wantServerName: "foo_bar",
+ wantToolName: "lookup",
+ wantOK: true,
+ },
+ {
+ name: "unknown prefix (no configured server matches)",
+ fullName: "nosuchserver_tool",
+ allowedNames: []string{"alpha", "beta"},
+ wantOK: false,
+ },
+ {
+ name: "prefix belongs to a real server not in the allowed set (authz-excluded)",
+ fullName: "beta_tool",
+ allowedNames: []string{"alpha"}, // beta exists in config elsewhere but is not allowed here
+ wantOK: false,
+ },
+ {
+ name: "empty allowed set",
+ fullName: "alpha_tool",
+ allowedNames: nil,
+ wantOK: false,
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ serverName, toolName, ok := resolveFederatedTool(tc.fullName, tc.allowedNames)
+ if ok != tc.wantOK {
+ t.Fatalf("ok = %v, want %v", ok, tc.wantOK)
+ }
+ if !tc.wantOK {
+ return
+ }
+ if serverName != tc.wantServerName || toolName != tc.wantToolName {
+ t.Errorf("resolveFederatedTool(%q, %v) = (%q, %q), want (%q, %q)",
+ tc.fullName, tc.allowedNames, serverName, toolName, tc.wantServerName, tc.wantToolName)
+ }
+ })
+ }
+}
+
+// --- unit: parseBackendJSONRPC (JSON and SSE-shaped backend responses) ---
+
+func TestParseBackendJSONRPC(t *testing.T) {
+ cases := []struct {
+ name string
+ contentType string
+ body string
+ wantResult string
+ wantErr bool
+ }{
+ {
+ name: "plain JSON object",
+ contentType: "application/json",
+ body: `{"jsonrpc":"2.0","id":"federated","result":{"ok":true}}`,
+ wantResult: `{"ok":true}`,
+ },
+ {
+ name: "SSE stream, single data event",
+ contentType: "text/event-stream",
+ body: "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":\"federated\",\"result\":{\"ok\":true}}\n\n",
+ wantResult: `{"ok":true}`,
+ },
+ {
+ name: "SSE stream, last of multiple data events wins",
+ contentType: "text/event-stream",
+ 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}`,
+ },
+ {
+ name: "malformed JSON",
+ contentType: "application/json",
+ body: `not json`,
+ wantErr: true,
+ },
+ {
+ name: "SSE stream with no data: line at all falls through to the raw body",
+ contentType: "text/event-stream",
+ body: "event: ping\n\n",
+ wantErr: true, // "event: ping" is not valid JSON on its own
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ got, err := parseBackendJSONRPC(tc.contentType, []byte(tc.body))
+ if tc.wantErr {
+ if err == nil {
+ t.Fatal("want error, got nil")
+ }
+ return
+ }
+ if err != nil {
+ t.Fatalf("unexpected error: %v", err)
+ }
+ if string(got.Result) != tc.wantResult {
+ t.Errorf("Result = %s, want %s", got.Result, tc.wantResult)
+ }
+ })
+ }
+}
+
+// --- end-to-end: POST /mcp ---
+
+// mockJSONRPCServer is a minimal MCP-shaped backend for federation tests:
+// it answers "tools/list" with a fixed tool set and "tools/call" by
+// echoing back the resolved (prefix-stripped) name/arguments it was
+// actually called with, so a test can assert routing stripped the prefix
+// correctly. called counts every request this server received, so a test
+// can assert a server the caller's group must NOT reach was never
+// contacted at all.
+type mockJSONRPCServer struct {
+ srv *httptest.Server
+ tools []mcpTool
+ called atomic.Int32
+}
+
+func newMockJSONRPCServer(t *testing.T, tools []mcpTool) *mockJSONRPCServer {
+ t.Helper()
+ m := &mockJSONRPCServer{tools: tools}
+ m.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ m.called.Add(1)
+ var req jsonrpcRequest
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ t.Fatalf("mock server: decode request: %v", err)
+ }
+ w.Header().Set("Content-Type", "application/json")
+ switch req.Method {
+ case "tools/list":
+ result, _ := json.Marshal(mcpToolsListResult{Tools: m.tools})
+ _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
+ case "tools/call":
+ var params mcpToolCallParams
+ _ = json.Unmarshal(req.Params, ¶ms)
+ echo, _ := json.Marshal(map[string]any{"calledName": params.Name, "calledArguments": json.RawMessage(params.Arguments)})
+ _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: echo})
+ default:
+ _ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: jsonrpcMethodNotFound, Message: "method not found"}})
+ }
+ }))
+ t.Cleanup(m.srv.Close)
+ return m
+}
+
+// newFederationTestConfig builds a Config with two MCP servers ("alpha",
+// "beta"), a group whose mcpServers glob restricts it to alphaAllowed
+// only when alphaOnly is true (otherwise unrestricted), and one inline
+// user "alice" in that group.
+func newFederationTestConfig(alphaURL, betaURL string, alphaOnly bool) *Config {
+ cfg := CreateConfig()
+ cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
+ cfg.MCPServers = map[string]*TargetConfig{
+ "alpha": {URL: alphaURL},
+ "beta": {URL: betaURL},
+ }
+ grpCfg := &GroupConfig{}
+ if alphaOnly {
+ grpCfg.MCPServers = []string{"alpha"}
+ }
+ cfg.Groups = map[string]*GroupConfig{"default": grpCfg}
+ cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
+ return cfg
+}
+
+func newFederatedRequest(t *testing.T, apiKey string, req jsonrpcRequest) *http.Request {
+ t.Helper()
+ body, err := json.Marshal(req)
+ if err != nil {
+ t.Fatalf("marshal request: %v", err)
+ }
+ httpReq := httptest.NewRequest(http.MethodPost, federatedMCPPath, bytes.NewReader(body))
+ if apiKey != "" {
+ httpReq.Header.Set("Authorization", "Bearer "+apiKey)
+ }
+ return httpReq
+}
+
+func TestHandleMCPFederated_Unauthenticated_Returns401(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("1")}))
+ if rec.Code != http.StatusUnauthorized {
+ t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String())
+ }
+}
+
+func TestHandleMCPFederated_MalformedBody_ReturnsJSONRPCParseError(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ req := httptest.NewRequest(http.MethodPost, federatedMCPPath, bytes.NewReader([]byte("not json")))
+ 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 (JSON-RPC errors are payload-level, not transport-level), body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error == nil || got.Error.Code != jsonrpcParseError {
+ t.Errorf("error = %+v, want code %d (parse error)", got.Error, jsonrpcParseError)
+ }
+ if string(got.ID) != "null" {
+ t.Errorf("id = %s, want literal null (JSON-RPC 2.0 §5.1: id must be null, not omitted, when it cannot be determined)", got.ID)
+ }
+}
+
+func TestHandleMCPFederated_UnknownMethod_ReturnsMethodNotFound(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "resources/list", ID: json.RawMessage("7")}))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error == nil || got.Error.Code != jsonrpcMethodNotFound {
+ t.Errorf("error = %+v, want code %d (method not found)", got.Error, jsonrpcMethodNotFound)
+ }
+ if string(got.ID) != "7" {
+ t.Errorf("id = %s, want the client's own id 7 echoed back", got.ID)
+ }
+}
+
+func TestHandleMCPFederated_Notification_Returns202NoBody(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "notifications/initialized"}))
+ if rec.Code != http.StatusAccepted {
+ t.Fatalf("status = %d, want 202, body=%s", rec.Code, rec.Body.String())
+ }
+ if rec.Body.Len() != 0 {
+ t.Errorf("body = %q, want empty (a notification gets no response body)", rec.Body.String())
+ }
+}
+
+func TestHandleMCPFederated_Initialize_EchoesProtocolVersionLocally(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ params, _ := json.Marshal(map[string]any{"protocolVersion": "2024-11-05"})
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
+ JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1"), Params: params,
+ }))
+ 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 struct {
+ Capabilities struct {
+ Tools map[string]any `json:"tools"`
+ } `json:"capabilities"`
+ ProtocolVersion string `json:"protocolVersion"`
+ }
+ if err := json.Unmarshal(got.Result, &result); err != nil {
+ t.Fatalf("decode result: %v", err)
+ }
+ if result.ProtocolVersion != "2024-11-05" {
+ t.Errorf("protocolVersion = %q, want the caller's own requested version echoed back", result.ProtocolVersion)
+ }
+ if result.Capabilities.Tools == nil {
+ t.Error("capabilities.tools must be present (tool-calling support declared)")
+ }
+}
+
+func TestHandleMCPFederated_Initialize_DefaultsWhenCallerOmitsVersion(t *testing.T) {
+ cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1")}))
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ var result struct {
+ ProtocolVersion string `json:"protocolVersion"`
+ }
+ if err := json.Unmarshal(got.Result, &result); err != nil {
+ t.Fatalf("decode result: %v", err)
+ }
+ if result.ProtocolVersion != defaultMCPProtocolVersion {
+ t.Errorf("protocolVersion = %q, want the fallback default %q", result.ProtocolVersion, defaultMCPProtocolVersion)
+ }
+}
+
+func TestHandleMCPFederated_Ping_AnsweredLocally(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, nil)
+ cfg := newFederationTestConfig(alpha.srv.URL, "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)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("1")}))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+ if alpha.called.Load() != 0 {
+ t.Error("ping must never contact any backend server")
+ }
+}
+
+// TestHandleMCPFederated_ToolsList_AggregatesPrefixesAndRespectsGroupAccess
+// is the core aggregation test: two mock MCP servers ("alpha", "beta"),
+// each with their own distinctly-named tool. A group unrestricted to
+// mcpServers sees both, prefixed and merged, sorted by name; a group
+// restricted to "alpha" only sees alpha's own tool, beta is never even
+// contacted, and per-target counters land only on the server(s) actually
+// reached.
+func TestHandleMCPFederated_ToolsList_AggregatesPrefixesAndRespectsGroupAccess(t *testing.T) {
+ cases := []struct {
+ name string
+ wantToolNames []string
+ wantBetaCalls int32
+ alphaOnly bool
+ }{
+ {
+ name: "unrestricted group sees both servers' tools, merged and sorted",
+ alphaOnly: false,
+ wantToolNames: []string{"alpha_lookup", "beta_search"},
+ wantBetaCalls: 1,
+ },
+ {
+ name: "group restricted to alpha never contacts beta at all",
+ alphaOnly: true,
+ wantToolNames: []string{"alpha_lookup"},
+ wantBetaCalls: 0,
+ },
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup", Description: "alpha's tool"}})
+ beta := newMockJSONRPCServer(t, []mcpTool{{Name: "search", Description: "beta's tool"}})
+
+ cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, tc.alphaOnly)
+ 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")
+ }
+
+ 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] = tool.Name
+ }
+ if len(gotNames) != len(tc.wantToolNames) {
+ t.Fatalf("tools = %v, want %v", gotNames, tc.wantToolNames)
+ }
+ for i, want := range tc.wantToolNames {
+ if gotNames[i] != want {
+ t.Errorf("tools[%d] = %q, want %q", i, gotNames[i], want)
+ }
+ }
+
+ if beta.called.Load() != tc.wantBetaCalls {
+ t.Errorf("beta server called %d times, want %d", beta.called.Load(), tc.wantBetaCalls)
+ }
+
+ // Per-target counter attribution (Feature B machinery reused
+ // by Feature C): alpha is always contacted in both cases.
+ now := time.Now()
+ alphaReq, ok := gw.limiter.getCounter(targetKindMCP, "alpha", metricReq, windowDay, now)
+ if !ok || alphaReq != 1 {
+ t.Errorf("mcp/alpha req:day counter = %d (ok=%v), want 1", alphaReq, ok)
+ }
+ betaReq, ok := gw.limiter.getCounter(targetKindMCP, "beta", metricReq, windowDay, now)
+ if !ok || betaReq != int64(tc.wantBetaCalls) {
+ t.Errorf("mcp/beta req:day counter = %d (ok=%v), want %d", betaReq, ok, tc.wantBetaCalls)
+ }
+ })
+ }
+}
+
+// TestHandleMCPFederated_ToolsList_UnreachableServerDegradesNotFails proves
+// one failing/unreachable allowed server never fails the whole aggregate —
+// the other server's tools still come back.
+func TestHandleMCPFederated_ToolsList_UnreachableServerDegradesNotFails(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
+ // "beta" points at a URL nothing listens on.
+ cfg := newFederationTestConfig(alpha.srv.URL, "http://127.0.0.1:1", false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error != nil {
+ t.Fatalf("error = %+v, want nil (a down server degrades the aggregate, it does not fail the call)", got.Error)
+ }
+ var result mcpToolsListResult
+ if err := json.Unmarshal(got.Result, &result); err != nil {
+ t.Fatalf("decode result: %v", err)
+ }
+ if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_lookup" {
+ t.Errorf("tools = %+v, want exactly alpha_lookup (beta unreachable, skipped)", result.Tools)
+ }
+}
+
+// TestHandleMCPFederated_ToolsCall_RoutesToCorrectServer_StripsPrefix
+// proves tools/call resolves the server prefix, strips it before
+// forwarding, contacts ONLY the resolved server (beta must never be
+// called for an alpha-prefixed tool), relays the backend's result, and
+// re-attaches the CLIENT's own id (never mcpBackendCall's synthetic one).
+func TestHandleMCPFederated_ToolsCall_RoutesToCorrectServer_StripsPrefix(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, nil)
+ beta := newMockJSONRPCServer(t, nil)
+ cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, false)
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+ gw, ok := h.(*Gateway)
+ if !ok {
+ t.Fatal("handler is not *Gateway")
+ }
+
+ callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup", Arguments: json.RawMessage(`{"q":"traefik"}`)})
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
+ JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage(`"client-id-42"`), Params: callParams,
+ }))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error != nil {
+ t.Fatalf("error = %+v, want nil", got.Error)
+ }
+ if string(got.ID) != `"client-id-42"` {
+ t.Errorf("id = %s, want the client's own id echoed back, not mcpBackendCall's synthetic one", got.ID)
+ }
+
+ var echoed struct {
+ CalledName string `json:"calledName"`
+ CalledArguments json.RawMessage `json:"calledArguments"`
+ }
+ if err := json.Unmarshal(got.Result, &echoed); err != nil {
+ t.Fatalf("decode result: %v", err)
+ }
+ if echoed.CalledName != "lookup" {
+ t.Errorf("backend saw tool name %q, want the prefix stripped down to %q", echoed.CalledName, "lookup")
+ }
+ if string(echoed.CalledArguments) != `{"q":"traefik"}` {
+ t.Errorf("backend saw arguments %s, want the client's own arguments forwarded verbatim", echoed.CalledArguments)
+ }
+
+ if alpha.called.Load() != 1 {
+ t.Errorf("alpha called %d times, want exactly 1", alpha.called.Load())
+ }
+ if beta.called.Load() != 0 {
+ t.Error("beta must never be contacted for an alpha-prefixed tool call")
+ }
+
+ betaReq, ok := gw.limiter.getCounter(targetKindMCP, "beta", metricReq, windowDay, time.Now())
+ if !ok || betaReq != 0 {
+ t.Errorf("mcp/beta req:day counter = %d (ok=%v), want 0 (never contacted)", betaReq, ok)
+ }
+ alphaReq, ok := gw.limiter.getCounter(targetKindMCP, "alpha", metricReq, windowDay, time.Now())
+ if !ok || alphaReq != 1 {
+ t.Errorf("mcp/alpha req:day counter = %d (ok=%v), want 1 (the resolved target)", alphaReq, ok)
+ }
+}
+
+// TestHandleMCPFederated_ToolsCall_BackendError_RelayedToClient proves a
+// backend's own JSON-RPC error (e.g. a tool execution failure) is relayed
+// to the client verbatim, under the client's own id — not swallowed or
+// replaced with a generic gateway error.
+func TestHandleMCPFederated_ToolsCall_BackendError_RelayedToClient(t *testing.T) {
+ alpha := httptest.NewServer(http.HandlerFunc(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: "tool execution failed"}})
+ }))
+ defer alpha.Close()
+
+ cfg := newFederationTestConfig(alpha.URL, "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)
+ }
+
+ callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"})
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
+ JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("5"), Params: callParams,
+ }))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error == nil || got.Error.Code != -32000 || got.Error.Message != "tool execution failed" {
+ t.Errorf("error = %+v, want the backend's own error relayed verbatim", got.Error)
+ }
+ if string(got.ID) != "5" {
+ t.Errorf("id = %s, want the client's own id 5", got.ID)
+ }
+}
+
+// TestHandleMCPFederated_ToolsCall_UnreachableServer_ReturnsInternalError
+// proves a resolved-but-unreachable backend server is a JSON-RPC internal
+// error, not a raw connection failure surfaced to the client.
+func TestHandleMCPFederated_ToolsCall_UnreachableServer_ReturnsInternalError(t *testing.T) {
+ cfg := newFederationTestConfig("http://127.0.0.1:1", "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)
+ }
+
+ callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"})
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
+ JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("6"), Params: callParams,
+ }))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error == nil || got.Error.Code != jsonrpcInternalError {
+ t.Errorf("error = %+v, want code %d (internal error)", got.Error, jsonrpcInternalError)
+ }
+}
+
+// TestHandleMCPFederated_ToolsCall_UnknownPrefix_ReturnsJSONRPCError proves
+// an unresolvable tool name is a JSON-RPC invalid-params error (still HTTP
+// 200), not an HTTP 404 — the request reached a real route and a real
+// method, it just named a tool nothing could route.
+func TestHandleMCPFederated_ToolsCall_UnknownPrefix_ReturnsJSONRPCError(t *testing.T) {
+ cases := []struct {
+ name string
+ toolName string
+ restricted bool
+ }{
+ {name: "no configured server matches the prefix at all", toolName: "nosuchserver_tool"},
+ {name: "prefix belongs to a real but group-restricted-away server", toolName: "beta_tool", restricted: true},
+ }
+ for _, tc := range cases {
+ t.Run(tc.name, func(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, nil)
+ beta := newMockJSONRPCServer(t, nil)
+ cfg := newFederationTestConfig(alpha.srv.URL, beta.srv.URL, tc.restricted)
+ 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: tc.toolName})
+ rec := httptest.NewRecorder()
+ h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{
+ JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("9"), Params: callParams,
+ }))
+ if rec.Code != http.StatusOK {
+ t.Fatalf("status = %d, want 200 (JSON-RPC error, not an HTTP error), body=%s", rec.Code, rec.Body.String())
+ }
+ var got jsonrpcResponse
+ if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
+ t.Fatalf("decode: %v", err)
+ }
+ if got.Error == nil || got.Error.Code != jsonrpcInvalidParams {
+ t.Errorf("error = %+v, want code %d (invalid params)", got.Error, jsonrpcInvalidParams)
+ }
+ if alpha.called.Load() != 0 || beta.called.Load() != 0 {
+ t.Error("no backend server must ever be contacted for an unresolvable tool name")
+ }
+ })
+ }
+}
+
+func TestHandleMCPFederated_RateLimited_Returns429(t *testing.T) {
+ alpha := newMockJSONRPCServer(t, nil)
+ cfg := newFederationTestConfig(alpha.srv.URL, "http://beta.invalid", false)
+ cfg.Groups["default"].Limits = &LimitsConfig{RequestsPerMinute: 1}
+ h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
+ if err != nil {
+ t.Fatalf("New: %v", err)
+ }
+
+ rec1 := httptest.NewRecorder()
+ h.ServeHTTP(rec1, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("1")}))
+ if rec1.Code != http.StatusOK {
+ t.Fatalf("first request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
+ }
+
+ rec2 := httptest.NewRecorder()
+ h.ServeHTTP(rec2, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "ping", ID: json.RawMessage("2")}))
+ if rec2.Code != http.StatusTooManyRequests {
+ t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String())
+ }
+}