mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat(mcp): federated /mcp aggregation endpoint
New POST /mcp: a single JSON-RPC 2.0 endpoint aggregating every MCP
server the caller's group can reach, for clients built against the
earlier pre-plugin gateway's one aggregated endpoint (pugbot's
mcpclient, agentkit) rather than this plugin's per-server
/mcp/{name}/... proxy.
Unlike the per-server proxy, this cannot be a byte-level reverse
proxy: aggregating N servers' tools/list results, or routing one
tools/call to the single right server, both require actually parsing
the JSON-RPC envelope. The gateway answers initialize/ping locally
and, for tools/list/tools/call, issues its own outbound JSON-RPC
call(s) to the relevant backend server(s) per incoming request — no
session state persisted across separate client requests, matching the
per-server proxy's own complete absence of gateway-side session
tracking. A backend may answer bare JSON or an SSE stream; both are
accepted (parseBackendJSONRPC).
- tools/list fans out concurrently to every allowed server and merges
results, prefixing each tool "<serverName>_<toolName>" — the
underscore convention the earlier agentgateway used and that
pugbot/agentkit already persist in their own tool-id records (e.g.
brave-search_brave_web_search, verified live). A failing/unreachable
server is skipped, not a whole-call failure.
- tools/call resolves the target by LONGEST matching
"<serverName>_" prefix against the caller's allowed servers
(disambiguates a server name that prefixes another), strips it,
forwards, and relays the result/error under the CLIENT's own id. An
unresolvable prefix (unconfigured, or authz-excluded) is a JSON-RPC
invalid-params error, not a 404.
- Reuses Feature B's countTargetRequest for per-target attribution —
every backend server actually contacted increments its own scope.
- notifications/* -> 202 no body; unknown method -> JSON-RPC
method-not-found; malformed body -> JSON-RPC parse error with
id:null (spec-required, not omitted).
No new limits config surface. Tests cover prefix resolution
(table-driven, incl. the ambiguous-prefix case), SSE/JSON backend
parsing, group-restricted aggregation (a denied server is never even
contacted), call routing/prefix-stripping, backend error relay,
unreachable-server handling, unknown-prefix errors, and per-target
counter attribution. README's MCP and A2A section documents the
federation contract and naming convention; the architecture diagram
and the total-scope/rejected-before-quota notes are updated to
include the new route.
This commit is contained in:
@@ -42,7 +42,7 @@ Catalog — see [Development](#development) for what that needs.
|
||||
```mermaid
|
||||
flowchart LR
|
||||
Req["incoming request"] --> Route{"route<br/>ServeHTTP: path+method dispatch"}
|
||||
Route -->|"a recognized route:<br/>/v1/models, /v1/chat/completions,<br/>/v1/embeddings, /v1/mcp/servers,<br/>/v1/agents, /mcp/name/...,<br/>/a2a/name/..., or /providerName/...<br/>for a known provider"| Auth["auth (per branch)<br/>identify(): key digest lookup"]
|
||||
Route -->|"a recognized route:<br/>/v1/models, /v1/chat/completions,<br/>/v1/embeddings, /v1/mcp/servers,<br/>/v1/agents, /mcp (federated),<br/>/mcp/name/...,<br/>/a2a/name/..., or /providerName/...<br/>for a known provider"| Auth["auth (per branch)<br/>identify(): key digest lookup"]
|
||||
Route -->|"no route matches"| Fallback["passthroughUnknown ? next : 404<br/>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<br/>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<br/>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<br/>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 `"<serverName>_<toolName>"` —
|
||||
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
|
||||
`"<serverName>_"` 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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 "<serverName>_<toolName>", the exact
|
||||
// underscore-separator convention the old agentgateway this plugin
|
||||
// replaces used, and the one pugbot's mcpclient and agentkit have already
|
||||
// persisted into their own tool-id databases (e.g.
|
||||
// "brave-search_brave_web_search" — verified live). A server that errors,
|
||||
// times out, 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 ("<serverName>_<toolName>").
|
||||
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 "<serverName>_" 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})
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user