mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
fix(mcp): backend timeouts, response validity, session handshake fallback
Review round 2 fix batch for the federated /mcp endpoint, driven by a
live probe against all 11 production MCP servers behind this gateway
through the byte-pipe proxy (session-less tools/list): 8 answer bare,
readitall and fetch reject it and need a session handshake, playwright
is a legacy HTTP+SSE-transport server (400 Missing sessionId) that no
handshake fixes.
MF1: every backend call now runs under its own context.WithTimeout
derived from r.Context() — toolsListBackendTimeout (20s) per fan-out
call, toolsCallBackendTimeout (120s) for tools/call, named constants.
MF2: doBackendJSONRPC rejects any non-2xx upstream status as an error
(previously unchecked — a non-JSON-RPC error body only failed by
accident of not parsing). parseBackendJSONRPC additionally rejects a
validly-JSON body that carries neither "result" nor "error" — the
exact silent-failure shape that would otherwise decode into an
all-zero jsonrpcResponse and hand the client a bare
{"jsonrpc":"2.0","id":1}. mcpFederatedToolsCall carries a
defensive final guard against the same shape before relaying.
MF3: mcpBackendCall is now bare-first-with-fallback — attempt the
bare call; on a JSON-RPC error OR HTTP 4xx (never a network failure,
timeout, or 5xx), retry once via initialize -> capture Mcp-Session-Id
-> retry with the header -> best-effort DELETE close. A legacy
HTTP+SSE-transport server (URL path ending "/sse") is excluded from
federation entirely before ever reaching this — isLegacySSETransportURL,
folded into allowedMCPServerNames — since no handshake fixes a
different transport; it stays reachable via the per-server proxy. When
every attempted server fails, tools/list returns a JSON-RPC internal
error ("no MCP server reachable") instead of a silently-empty list.
SF6: defaultMCPProtocolVersion -> "2025-11-25"; a caller's own
protocolVersion is echoed back only from a small explicit allowlist
(knownMCPProtocolVersions), falling back to the default otherwise —
never echoes an arbitrary caller string as this gateway's own claimed
capability.
SF8: GET/DELETE/etc. on exactly /mcp is now 405 with Allow: POST,
checked before auth (llmgateway.go).
Folded minors: countTargetRequests batches per-target counter writes
for tools/list's whole fan-out into ONE storeIncrMulti call after
wg.Wait() (was one call per server); every "actually contacted" doc
comment now says "attempted" (attribution is attempt-, not
success-, accounting); lastSSEDataLine documents its no-cross-message-
id-correlation limitation; routes_unified.go's stale buildLimitScopes/
withTotalScope comments (which claimed handleAdminAPI calls
buildLimitScopes directly — it never has) are corrected; jsonrpcRequest
documents why its "jsonrpc" field is decoded but never validated.
Tests: ~30 new/updated cases covering per-backend timeout derivation
(via the shorter-of-two-context-deadlines technique, no real 20s/120s
waits), non-2xx rejection, the empty-envelope guard end-to-end,
handshake-trigger classification (table-driven: JSON-RPC error and
HTTP 4xx trigger, HTTP 500 and network failure do not), the full
handshake sequence (bare -> initialize -> retry -> close, asserted
call-by-call), handshake-gives-up-too, all-servers-failed, legacy-SSE
exclusion (never contacted, in both tools/list and tools/call),
protocol-version allowlist (all 4 versions + one rejected), 405/Allow,
and best-effort session-close failure never surfacing to the caller.
This commit is contained in:
@@ -1011,14 +1011,16 @@ func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) {
|
||||
l.storeIncrMulti(entries) // no error return by contract (doc comment above); ok is intentionally discarded
|
||||
}
|
||||
|
||||
// countTargetRequest increments one MCP-server's or A2A agent's per-target
|
||||
// request counters (Feature B, v0.21) at min, hour, day, AND month — the
|
||||
// only scope kind this package tracks a month window of REQUESTS for.
|
||||
// kind is targetScopeKind's own output ("mcp" or scopeKindAgent —
|
||||
// mcp_a2a.go), id is the configured target name; both come embedded in
|
||||
// windowKey's own key string, so a target scope can never collide with a
|
||||
// user/group/total scope of the same id even by coincidence — kind is
|
||||
// part of the key, not just a struct field.
|
||||
// countTargetRequests increments MULTIPLE target scopes' per-target request
|
||||
// counters (Feature B, v0.21) — each at min, hour, day, AND month, the only
|
||||
// scope kind this package tracks a month window of REQUESTS for — in ONE
|
||||
// storeIncrMulti call. kind is targetScopeKind's own output ("mcp" or
|
||||
// scopeKindAgent — mcp_a2a.go), id is the configured target name; both come
|
||||
// embedded in windowKey's own key string, so a target scope can never
|
||||
// collide with a user/group/total scope of the same id even by coincidence
|
||||
// — kind is part of the key, not just a struct field. A scope's own
|
||||
// `limits` field is ignored entirely (never read); callers pass plain
|
||||
// {kind, id} pairs.
|
||||
//
|
||||
// This is deliberately NOT folded into checkAndCount: a target scope
|
||||
// carries no limits by design (this round adds no config surface for
|
||||
@@ -1033,20 +1035,44 @@ func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) {
|
||||
// rejected as scope creep touching user/group accounting nothing else in
|
||||
// this round asked to change.
|
||||
//
|
||||
// Callers: handleTargetProxy (mcp_a2a.go, once per proxied request, after
|
||||
// its own user/group/total admission check passes) and the federated /mcp
|
||||
// endpoint (mcp_federation.go, once per actual backend server contacted —
|
||||
// tools/list may contact several in one incoming client request, tools/call
|
||||
// exactly one), so a call attributed to a specific target is counted the
|
||||
// same way regardless of which route reached it.
|
||||
func (l *limiter) countTargetRequest(kind, id string) {
|
||||
// Batching (review round 2, 2026-08-21): the federated /mcp endpoint's own
|
||||
// tools/list fan-out can attempt several backend servers for ONE incoming
|
||||
// client request — the earlier per-server countTargetRequest call inside
|
||||
// each fan-out goroutine paid one round trip per server; this collects
|
||||
// every scope from the whole fan-out into a single call after the fan-out's
|
||||
// own wg.Wait(), the same "one request in, one round trip out" discipline
|
||||
// checkAndCount/account already apply per scope-slice.
|
||||
//
|
||||
// Callers: handleTargetProxy (mcp_a2a.go, a 1-element slice, once per
|
||||
// proxied request, after its own user/group/total admission check passes)
|
||||
// and the federated /mcp endpoint (mcp_federation.go: a 1-element slice for
|
||||
// tools/call's single resolved target, or one element per server ATTEMPTED
|
||||
// — not necessarily reached or succeeded — for tools/list's fan-out; see
|
||||
// mcpFederatedToolsList's own doc comment for why "attempted" is the
|
||||
// correct word here).
|
||||
func (l *limiter) countTargetRequests(scopes []limitScope) {
|
||||
if len(scopes) == 0 {
|
||||
return
|
||||
}
|
||||
now := l.now()
|
||||
l.storeIncrMulti([]counterIncr{
|
||||
newCounterIncr(kind, id, metricReq, windowMin, now, 1, minWindowTTL),
|
||||
newCounterIncr(kind, id, metricReq, windowHour, now, 1, hourWindowTTL),
|
||||
newCounterIncr(kind, id, metricReq, windowDay, now, 1, dayWindowTTL),
|
||||
newCounterIncr(kind, id, metricReq, windowMonth, now, 1, monthWindowTTL),
|
||||
})
|
||||
entries := make([]counterIncr, 0, len(scopes)*4)
|
||||
for _, sc := range scopes {
|
||||
entries = append(entries,
|
||||
newCounterIncr(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL),
|
||||
newCounterIncr(sc.kind, sc.id, metricReq, windowHour, now, 1, hourWindowTTL),
|
||||
newCounterIncr(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL),
|
||||
newCounterIncr(sc.kind, sc.id, metricReq, windowMonth, now, 1, monthWindowTTL),
|
||||
)
|
||||
}
|
||||
l.storeIncrMulti(entries)
|
||||
}
|
||||
|
||||
// countTargetRequest is countTargetRequests for exactly one target scope —
|
||||
// handleTargetProxy's own single-target shape (mcp_a2a.go), where a batch
|
||||
// of one buys nothing but keeps the call site a plain two-argument call
|
||||
// instead of a one-element slice literal.
|
||||
func (l *limiter) countTargetRequest(kind, id string) {
|
||||
l.countTargetRequests([]limitScope{{kind: kind, id: id}})
|
||||
}
|
||||
|
||||
// targetCounterKeysPerScope is the number of windowKey strings
|
||||
|
||||
+16
-1
@@ -455,7 +455,22 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// 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 {
|
||||
//
|
||||
// Checked on r.URL.Path alone, BEFORE the method check: a GET or
|
||||
// DELETE to exactly "/mcp" is a real, well-known route addressed with
|
||||
// the wrong verb (SF8, review round 2, 2026-08-21) — 405 with an
|
||||
// Allow header naming the one method this route accepts, matching
|
||||
// RFC 9110 §15.5.6, rather than falling through to passthroughUnknown/
|
||||
// 404 the way a path this handler has never heard of would. No
|
||||
// auth check runs first: which HTTP method a route accepts is
|
||||
// unauthenticated request-shape information, not something worth
|
||||
// gating behind a valid API key.
|
||||
if r.URL.Path == federatedMCPPath {
|
||||
if r.Method != http.MethodPost {
|
||||
sw.Header().Set("Allow", http.MethodPost)
|
||||
writeOAIError(sw, http.StatusMethodNotAllowed, "invalid_request_error", "method not allowed")
|
||||
return
|
||||
}
|
||||
u, grp, ok := g.auth.identify(r)
|
||||
g.logAuthEvent(ok, authEventUserName(u), r)
|
||||
if !ok {
|
||||
|
||||
+418
-100
@@ -4,11 +4,14 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// federatedMCPPath is the bare "/mcp" route this file implements —
|
||||
@@ -34,6 +37,35 @@ const (
|
||||
jsonrpcInternalError = -32603
|
||||
)
|
||||
|
||||
// mcpSessionHeader is the MCP Streamable HTTP transport's session-id
|
||||
// header — set by a backend server's "initialize" response to open a
|
||||
// session, and echoed by a client (this gateway, on the handshake-fallback
|
||||
// retry) on every subsequent call that session covers. See
|
||||
// mcpBackendCall's own doc comment for when this gateway ever sends one.
|
||||
const mcpSessionHeader = "Mcp-Session-Id"
|
||||
|
||||
// Per-request budgets a federated call to a single backend server is
|
||||
// allowed: the WHOLE per-server operation — a bare attempt, and, only when
|
||||
// that signals a session is required, one initialize + one retry + a
|
||||
// best-effort session close — must fit inside this window, not each
|
||||
// individual HTTP request separately (mcpBackendCall's own doc comment).
|
||||
// tools/list's fan-out runs one of these per allowed server, concurrently,
|
||||
// so the shorter budget bounds a single client request's worst case to
|
||||
// toolsListBackendTimeout regardless of how many servers are configured;
|
||||
// tools/call talks to exactly one server and gets the longer budget a
|
||||
// slower tool invocation (a web search, a browser action) may need.
|
||||
const (
|
||||
toolsListBackendTimeout = 20 * time.Second
|
||||
toolsCallBackendTimeout = 120 * time.Second
|
||||
// mcpSessionCloseTimeout bounds the best-effort DELETE that ends a
|
||||
// handshake-fallback session (mcpBackendCall). It runs on its own
|
||||
// context.Background()-derived deadline, deliberately NOT the parent
|
||||
// call's own (possibly already near-exhausted) toolsList/
|
||||
// toolsCallBackendTimeout budget: cleanup must never race, or get
|
||||
// starved by, the real call it is cleaning up after.
|
||||
mcpSessionCloseTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
// 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).
|
||||
@@ -42,6 +74,15 @@ const (
|
||||
// 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.
|
||||
//
|
||||
// JSONRPC itself is decoded but never validated against the literal
|
||||
// string "2.0" — deliberately permissive: real MCP clients this gateway
|
||||
// has actually seen (and the mock backends this repo's own tests drive)
|
||||
// never send anything else, rejecting a request purely for carrying a
|
||||
// different (or absent) "jsonrpc" field would reject a client this
|
||||
// gateway can otherwise serve correctly for a purity check no consumer
|
||||
// has ever needed, and the field is not read for any routing or safety
|
||||
// decision this file makes.
|
||||
type jsonrpcRequest struct {
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
@@ -78,10 +119,10 @@ type jsonrpcResponse struct {
|
||||
// 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).
|
||||
// body, an unauthenticated/forbidden caller, a rate limit, a disallowed
|
||||
// HTTP method — 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")
|
||||
@@ -107,18 +148,55 @@ func writeJSONRPCErrorResponse(w http.ResponseWriter, id json.RawMessage, code i
|
||||
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.
|
||||
// isLegacySSETransportURL reports whether rawURL's path ends in "/sse" —
|
||||
// the MCP prior HTTP+SSE transport's own URL convention (superseded by
|
||||
// Streamable HTTP), where the CLIENT — not an intermediary — must drive a
|
||||
// persistent SSE session: a GET establishing the event stream, POSTs
|
||||
// correlated against a server-assigned endpoint the stream itself
|
||||
// announces. mcpBackendCall's simple "POST a JSON-RPC object, read one
|
||||
// response" model cannot speak that transport at all; there is no
|
||||
// handshake or retry that fixes this, unlike the session-required case
|
||||
// MF3 handles. Verified live against production traffic (review round 2,
|
||||
// 2026-08-21): a real server on this transport (URL ending "/sse")
|
||||
// answers a Streamable-HTTP-shaped POST with HTTP 400 "Missing sessionId"
|
||||
// — a different failure MODE than the session-required 4xx/JSON-RPC-error
|
||||
// case mcpBackendCall's handshake retry targets, and one no retry of any
|
||||
// kind resolves. Such a server is excluded from federation's tools/list
|
||||
// fan-out and tools/call resolution (allowedMCPServerNames) before ever
|
||||
// reaching mcpBackendCall — it stays fully reachable via the per-server
|
||||
// proxy (/mcp/{name}/...), where the real client drives the transport
|
||||
// end to end, exactly as it always has.
|
||||
//
|
||||
// A rawURL that fails to parse returns false (not excluded) rather than
|
||||
// panicking or erroring: every configured target URL already passed
|
||||
// validateTargetURLs at construction (mcp_a2a.go), so this branch is
|
||||
// unreachable in practice, and "don't exclude" is the safe default for an
|
||||
// unreachable branch — the same convention sanitizeBaseURL applies to its
|
||||
// own unparsable-URL case (admin.go).
|
||||
func isLegacySSETransportURL(rawURL string) bool {
|
||||
u, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return strings.HasSuffix(u.Path, "/sse")
|
||||
}
|
||||
|
||||
// allowedMCPServerNames returns the sorted names of every MCP server BOTH
|
||||
// grp.allowsMCP permits AND that speaks a transport federation can
|
||||
// actually use (isLegacySSETransportURL excludes the rest) — 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, and a legacy-SSE-transport server is invisible to
|
||||
// federation regardless of group access (it is not a federation-eligible
|
||||
// server for ANY caller, not an authorization decision).
|
||||
func allowedMCPServerNames(cfg *Config, grp *group) []string {
|
||||
names := make([]string, 0, len(cfg.MCPServers))
|
||||
for name := range cfg.MCPServers {
|
||||
if grp.allowsMCP(name) {
|
||||
for name, tc := range cfg.MCPServers {
|
||||
if grp.allowsMCP(name) && !isLegacySSETransportURL(tc.URL) {
|
||||
names = append(names, name)
|
||||
}
|
||||
}
|
||||
@@ -128,7 +206,8 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string {
|
||||
|
||||
// handleMCPFederated implements POST /mcp (federatedMCPPath): the
|
||||
// federated JSON-RPC endpoint this file's own doc comment explains the
|
||||
// need for.
|
||||
// need for. ServeHTTP (llmgateway.go) answers any OTHER method on this
|
||||
// exact path with 405 (Allow: POST) before this function is ever called.
|
||||
//
|
||||
// Unlike handleTargetProxy, this is not a byte-level reverse proxy — it
|
||||
// cannot be: aggregating N servers' tools/list results into one response,
|
||||
@@ -139,8 +218,10 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string {
|
||||
// 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
|
||||
// persisted across separate CLIENT requests (mcpBackendCall's own doc
|
||||
// comment covers the one-call, per-backend-server session a handshake
|
||||
// fallback DOES open and close, entirely inside a single incoming
|
||||
// request), 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
|
||||
@@ -153,11 +234,9 @@ func allowedMCPServerNames(cfg *Config, grp *group) []string {
|
||||
//
|
||||
// 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.
|
||||
// other metered route — with per-target attribution added separately
|
||||
// inside mcpFederatedToolsList/mcpFederatedToolsCall (their own doc
|
||||
// comments cover exactly what "attributed" means for each).
|
||||
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")
|
||||
@@ -210,29 +289,55 @@ 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"
|
||||
// defaultMCPProtocolVersion is reported when the caller's own "initialize"
|
||||
// request omits protocolVersion, or names one outside
|
||||
// knownMCPProtocolVersions. 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 names a recognized one.
|
||||
const defaultMCPProtocolVersion = "2025-11-25"
|
||||
|
||||
// knownMCPProtocolVersions is the small, explicit allowlist of protocol
|
||||
// revisions mcpFederatedInitialize will ever echo back verbatim — every
|
||||
// dated revision this plugin's own reading of the MCP spec's version
|
||||
// history recognizes at the time this constant was last reviewed.
|
||||
// Echoing an ARBITRARY caller-supplied string back as this gateway's own
|
||||
// declared protocolVersion would let a malformed or malicious client make
|
||||
// this response claim compatibility with a revision this plugin has never
|
||||
// actually been evaluated against; the allowlist keeps the "echo the
|
||||
// caller's choice" behavior (mcpFederatedInitialize's own doc comment)
|
||||
// bounded to values this plugin actually knows about, falling back to
|
||||
// defaultMCPProtocolVersion for anything else — including a genuinely
|
||||
// newer, valid revision this constant simply predates; that is a
|
||||
// known, accepted staleness risk of hardcoding a version list at all,
|
||||
// not a correctness bug.
|
||||
var knownMCPProtocolVersions = map[string]bool{
|
||||
"2024-11-05": true,
|
||||
"2025-03-26": true,
|
||||
"2025-06-18": true,
|
||||
"2025-11-25": true,
|
||||
}
|
||||
|
||||
// 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.
|
||||
// protocolVersion echoes the caller's own requested value when it names a
|
||||
// revision in knownMCPProtocolVersions: version negotiation is normally
|
||||
// the caller's choice to make, but echoing an arbitrary unrecognized
|
||||
// string back as this gateway's own declared capability would overclaim
|
||||
// compatibility this plugin cannot back up (knownMCPProtocolVersions' own
|
||||
// doc comment) — falling back to defaultMCPProtocolVersion in that case,
|
||||
// and whenever the caller omits protocolVersion entirely.
|
||||
// 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
|
||||
protocolVersion := defaultMCPProtocolVersion
|
||||
if knownMCPProtocolVersions[params.ProtocolVersion] {
|
||||
protocolVersion = params.ProtocolVersion
|
||||
}
|
||||
|
||||
writeJSONRPCResult(w, req.ID, map[string]any{
|
||||
@@ -262,39 +367,43 @@ type mcpToolsListResult struct {
|
||||
Tools []mcpTool `json:"tools"`
|
||||
}
|
||||
|
||||
// mcpBackendCall issues one JSON-RPC 2.0 request to targetURL — an MCP
|
||||
// doBackendJSONRPC issues one JSON-RPC 2.0 POST 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.
|
||||
// handleTargetProxy reverse-proxies every method to — carrying sessionID
|
||||
// as the Mcp-Session-Id request header when non-empty. It is the one place
|
||||
// that actually builds and sends an HTTP request for this file; every
|
||||
// other function in the handshake-fallback machinery (mcpBackendCall,
|
||||
// mcpBackendHandshake, mcpBackendCloseSession's own DELETE aside) goes
|
||||
// through this.
|
||||
//
|
||||
// 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.
|
||||
// Return values: resp is nil whenever err is non-nil. status is the
|
||||
// backend's real HTTP status when one was received at all (0 for a
|
||||
// request-build or network failure that never got a response), reported
|
||||
// even alongside a non-nil err — mcpBackendCall's own handshake-fallback
|
||||
// decision (a JSON-RPC error OR an HTTP 4xx) needs the status specifically
|
||||
// in the case where a non-2xx response has no valid JSON-RPC body at all
|
||||
// (err is then the "upstream returned HTTP %d" error below, not a
|
||||
// parse/validity error). respSessionID is the backend's own
|
||||
// Mcp-Session-Id RESPONSE header — read on every call, not only
|
||||
// "initialize", so mcpBackendHandshake needs no special-cased second
|
||||
// code path for the one call that actually cares about it.
|
||||
//
|
||||
// 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) {
|
||||
// A non-2xx status is always an error (MF2, review round 2, 2026-08-21):
|
||||
// the earlier version of this file had no such check at all, which let a
|
||||
// non-JSON-RPC-shaped error page (a plain-text 500, an nginx 502 HTML
|
||||
// page) fall through to parseBackendJSONRPC and fail there anyway, but
|
||||
// only by accident of that body not happening to be valid JSON — a body
|
||||
// that WAS valid JSON but not a real JSON-RPC response (see
|
||||
// parseBackendJSONRPC's own validity check) would have silently produced
|
||||
// an all-zero-fields response otherwise.
|
||||
func (g *Gateway) doBackendJSONRPC(ctx context.Context, targetURL, method string, params any, sessionID string) (resp *jsonrpcResponse, status int, respSessionID string, err error) {
|
||||
paramsRaw, err := json.Marshal(params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, 0, "", err
|
||||
}
|
||||
reqBody, err := json.Marshal(jsonrpcRequest{JSONRPC: jsonrpcVersion, ID: json.RawMessage(`"federated"`), Method: method, Params: paramsRaw})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, 0, "", err
|
||||
}
|
||||
|
||||
// gosec G704 (SSRF via taint analysis): targetURL is an
|
||||
@@ -307,22 +416,152 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
|
||||
// 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
|
||||
return nil, 0, "", err
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "application/json, text/event-stream")
|
||||
if sessionID != "" {
|
||||
httpReq.Header.Set(mcpSessionHeader, sessionID)
|
||||
}
|
||||
|
||||
resp, err := g.targetClient.Do(httpReq)
|
||||
httpResp, err := g.targetClient.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, 0, "", err
|
||||
}
|
||||
defer httpResp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure
|
||||
|
||||
status = httpResp.StatusCode
|
||||
respSessionID = httpResp.Header.Get(mcpSessionHeader)
|
||||
|
||||
respBody, err := io.ReadAll(io.LimitReader(httpResp.Body, maxRequestBytes))
|
||||
if err != nil {
|
||||
return nil, status, respSessionID, err
|
||||
}
|
||||
if status < 200 || status >= 300 {
|
||||
return nil, status, respSessionID, fmt.Errorf("upstream returned HTTP %d", status)
|
||||
}
|
||||
|
||||
parsed, err := parseBackendJSONRPC(httpResp.Header.Get("Content-Type"), respBody)
|
||||
if err != nil {
|
||||
return nil, status, respSessionID, err
|
||||
}
|
||||
return parsed, status, respSessionID, nil
|
||||
}
|
||||
|
||||
// mcpBackendHandshake issues a fresh, session-less "initialize" call
|
||||
// against targetURL — minimal clientInfo, this gateway's own default
|
||||
// protocol version — for mcpBackendCall's handshake-fallback retry.
|
||||
// sessionID is "" when the backend's response carries no Mcp-Session-Id
|
||||
// header of its own (a legitimate, spec-permitted case for a server that
|
||||
// turns out to be stateless-tolerant after all); err is non-nil only when
|
||||
// the initialize call itself failed outright (network error, non-2xx, a
|
||||
// JSON-RPC error response to initialize itself) — there is no retry of a
|
||||
// failed handshake.
|
||||
func (g *Gateway) mcpBackendHandshake(ctx context.Context, targetURL string) (sessionID string, err error) {
|
||||
initParams := map[string]any{
|
||||
"protocolVersion": defaultMCPProtocolVersion,
|
||||
"capabilities": map[string]any{},
|
||||
"clientInfo": map[string]any{"name": "traefik-llmgateway", "version": pluginVersion},
|
||||
}
|
||||
resp, _, respSessionID, err := g.doBackendJSONRPC(ctx, targetURL, "initialize", initParams, "")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if resp.Error != nil {
|
||||
return "", fmt.Errorf("initialize: %s (code %d)", resp.Error.Message, resp.Error.Code)
|
||||
}
|
||||
return respSessionID, nil
|
||||
}
|
||||
|
||||
// mcpBackendCloseSession best-effort DELETEs targetURL with sessionID
|
||||
// attached (the MCP Streamable HTTP transport's own session-end
|
||||
// convention) — failure is logged, never surfaced: this runs after the
|
||||
// caller's own real response has already been read and returned, so
|
||||
// nothing downstream is waiting on it, and a backend that does not support
|
||||
// session termination at all (silently ignoring the DELETE, or answering
|
||||
// 404/405 for it) must never turn into a federation-visible error over
|
||||
// pure cleanup.
|
||||
func (g *Gateway) mcpBackendCloseSession(ctx context.Context, targetURL, sessionID string) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, targetURL, nil) //nolint:gosec // same operator-configured target URL as doBackendJSONRPC
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
req.Header.Set(mcpSessionHeader, sessionID)
|
||||
|
||||
resp, err := g.targetClient.Do(req)
|
||||
if err != nil {
|
||||
g.logf("federated mcp: best-effort session close failed for %q: %v", targetURL, err)
|
||||
return
|
||||
}
|
||||
_ = resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure
|
||||
}
|
||||
|
||||
// mcpBackendCall issues a JSON-RPC 2.0 request for method to targetURL —
|
||||
// bare first, matching the per-request "no persisted CLIENT session"
|
||||
// design (handleMCPFederated's own doc comment) — and falls back to a
|
||||
// ONE-SHOT initialize/retry/close handshake ONLY when the bare attempt
|
||||
// signals the server actually requires a session: a JSON-RPC-level error
|
||||
// response, or an HTTP 4xx status.
|
||||
//
|
||||
// This ruling is backed by a live probe against all 11 production servers
|
||||
// behind this gateway (review round 2, 2026-08-21): 8 answer a bare,
|
||||
// session-less call outright (brave-search, captcha-solver, google-maps,
|
||||
// lightpanda, time, weather, web-search, wikipedia); readitall and fetch
|
||||
// both reject the bare call — readitall: "method tools/list is invalid
|
||||
// during..."; fetch: -32602 with or without params — and need exactly
|
||||
// this handshake to work; a third class (a legacy HTTP+SSE-transport
|
||||
// server, verified as playwright, HTTP 400 "Missing sessionId") is
|
||||
// excluded from federation entirely before ever reaching this function
|
||||
// (isLegacySSETransportURL, allowedMCPServerNames), since no handshake
|
||||
// fixes a transport this function was never built to speak.
|
||||
//
|
||||
// A network failure, timeout, or 5xx never triggers the fallback:
|
||||
// retrying those would spend a second timeout budget against a server
|
||||
// that was never going to answer either way, bare or not — only a
|
||||
// response that specifically signals "you're missing a session" is worth
|
||||
// a second attempt. ctx governs the WHOLE operation (bare attempt,
|
||||
// optional handshake, optional retry, best-effort close all share it,
|
||||
// except the close itself — see mcpSessionCloseTimeout's own doc
|
||||
// comment); callers derive it from context.WithTimeout(r.Context(), ...)
|
||||
// with toolsListBackendTimeout or toolsCallBackendTimeout (MF1, review
|
||||
// round 2, 2026-08-21).
|
||||
func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string, params any) (*jsonrpcResponse, error) {
|
||||
resp, status, _, err := g.doBackendJSONRPC(ctx, targetURL, method, params, "")
|
||||
if err == nil && resp.Error == nil {
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
retryEligible := (err == nil && resp.Error != nil) || (status >= 400 && status < 500)
|
||||
if !retryEligible {
|
||||
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
|
||||
sessionID, initErr := g.mcpBackendHandshake(ctx, targetURL)
|
||||
if initErr != nil {
|
||||
// The handshake itself failed: surface the bare attempt's own
|
||||
// outcome when it has one (a JSON-RPC error response is more
|
||||
// informative to a caller than a generic "handshake failed"),
|
||||
// falling back to initErr only when the bare attempt produced no
|
||||
// usable response at all (its own err was non-nil — the HTTP-4xx-
|
||||
// with-no-JSON-RPC-body case).
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
return parseBackendJSONRPC(resp.Header.Get("Content-Type"), respBody)
|
||||
if sessionID != "" {
|
||||
defer func() {
|
||||
closeCtx, cancel := context.WithTimeout(context.Background(), mcpSessionCloseTimeout)
|
||||
defer cancel()
|
||||
g.mcpBackendCloseSession(closeCtx, targetURL, sessionID)
|
||||
}()
|
||||
}
|
||||
|
||||
retryResp, _, _, retryErr := g.doBackendJSONRPC(ctx, targetURL, method, params, sessionID)
|
||||
if retryErr != nil {
|
||||
return nil, retryErr
|
||||
}
|
||||
return retryResp, nil
|
||||
}
|
||||
|
||||
// parseBackendJSONRPC decodes body as a backend MCP server's JSON-RPC
|
||||
@@ -330,8 +569,21 @@ func (g *Gateway) mcpBackendCall(ctx context.Context, targetURL, method string,
|
||||
// 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.
|
||||
// LAST such line is parsed (lastSSEDataLine's own doc comment covers what
|
||||
// this does and does not correlate).
|
||||
//
|
||||
// A body that parses as valid JSON but carries NEITHER "result" NOR
|
||||
// "error" is also rejected (MF2, review round 2, 2026-08-21): a backend
|
||||
// answering with some other JSON shape entirely (a plain `{}`, a
|
||||
// non-compliant `{"status":"ok"}`) would otherwise decode into a
|
||||
// jsonrpcResponse with every field at its zero value and LOOK like a
|
||||
// legitimate, empty-but-successful JSON-RPC response — silently producing
|
||||
// a client-visible `{"jsonrpc":"2.0","id":1}` with neither a result to
|
||||
// read nor an error to report. Every caller of this function treats the
|
||||
// returned error identically to a network failure (fails the call,
|
||||
// mcpFederatedToolsList skips the server rather than crediting it with an
|
||||
// empty success), so this is caught here once rather than at every call
|
||||
// site.
|
||||
func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, error) {
|
||||
if strings.Contains(strings.ToLower(contentType), "text/event-stream") {
|
||||
body = lastSSEDataLine(body)
|
||||
@@ -340,6 +592,9 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err
|
||||
if err := json.Unmarshal(body, &out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if out.Result == nil && out.Error == nil {
|
||||
return nil, fmt.Errorf("invalid JSON-RPC response: neither result nor error present")
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
@@ -348,6 +603,19 @@ func parseBackendJSONRPC(contentType string, body []byte) (*jsonrpcResponse, err
|
||||
// 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).
|
||||
//
|
||||
// This does NOT implement per-request id correlation across multiple
|
||||
// concurrent SSE-delivered messages the way a real MCP client's
|
||||
// Streamable-HTTP transport layer would (matching a specific response to
|
||||
// the request that triggered it by id, ignoring unrelated
|
||||
// server-initiated notifications interleaved on the same stream) — it
|
||||
// takes the textually LAST "data:" line, full stop. This is a deliberate,
|
||||
// documented limitation, acceptable for mcpBackendCall's own one-request-
|
||||
// in-flight-at-a-time usage of this stream (it never has two outstanding
|
||||
// calls sharing one SSE response to correlate between): a backend that
|
||||
// interleaves an unrelated notification AFTER its real response on the
|
||||
// same stream would be misread by this function. No production server
|
||||
// probed for this round exhibited that behavior.
|
||||
func lastSSEDataLine(body []byte) []byte {
|
||||
var last []byte
|
||||
for _, line := range bytes.Split(body, []byte("\n")) {
|
||||
@@ -366,22 +634,28 @@ func lastSSEDataLine(body []byte) []byte {
|
||||
|
||||
// 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.
|
||||
// concurrently — each under its own toolsListBackendTimeout budget
|
||||
// (context.WithTimeout off r.Context(), MF1) — 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/invalid result
|
||||
// is logged and skipped, not surfaced as a whole-call failure — UNLESS
|
||||
// every single attempted server failed, in which case an empty tools list
|
||||
// would misleadingly look like "this caller's group has no MCP access" —
|
||||
// see the loud-failure branch below.
|
||||
//
|
||||
// countTargetRequest attributes one request to EVERY server actually
|
||||
// contacted here (handleMCPFederated's own doc comment) — a single
|
||||
// countTargetRequests (limits.go) attributes one request to EVERY server
|
||||
// ATTEMPTED here, in ONE batched call after wg.Wait() — "attempted", not
|
||||
// "reached" or "succeeded": a server this gateway dialed and got a
|
||||
// response (or a timeout, or a connection refusal) from still had a real
|
||||
// request sent to it and a real slot of this gateway's outbound capacity
|
||||
// spent on it, which is what these counters exist to track. A single
|
||||
// federated tools/list call can move several targets' own counters, not
|
||||
// just one.
|
||||
// just one — unlike handleTargetProxy's always-exactly-one-target shape
|
||||
// (mcp_a2a.go).
|
||||
func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request, req jsonrpcRequest, grp *group) {
|
||||
names := allowedMCPServerNames(g.cfg, grp)
|
||||
|
||||
@@ -395,9 +669,10 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request,
|
||||
wg.Add(1)
|
||||
go func(name string) {
|
||||
defer wg.Done()
|
||||
ctx, cancel := context.WithTimeout(r.Context(), toolsListBackendTimeout)
|
||||
defer cancel()
|
||||
targetURL := g.cfg.MCPServers[name].URL
|
||||
resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/list", struct{}{})
|
||||
g.limiter.countTargetRequest(targetKindMCP, name)
|
||||
resp, err := g.mcpBackendCall(ctx, targetURL, "tools/list", struct{}{})
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
@@ -417,12 +692,32 @@ func (g *Gateway) mcpFederatedToolsList(w http.ResponseWriter, r *http.Request,
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if len(names) > 0 {
|
||||
attempted := make([]limitScope, len(names))
|
||||
for i, name := range names {
|
||||
attempted[i] = limitScope{kind: targetKindMCP, id: name}
|
||||
}
|
||||
g.limiter.countTargetRequests(attempted)
|
||||
}
|
||||
|
||||
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 })
|
||||
|
||||
// Every attempted server failed: an empty tools list here would read
|
||||
// as "this caller's group has legitimate no-access", indistinguishable
|
||||
// from the case where names itself was empty (no allowed servers at
|
||||
// all — that IS a correct empty result, left alone below). A total
|
||||
// outage is a real infrastructure problem the caller needs to see as
|
||||
// one, not a silently degraded empty success (MF3, review round 2,
|
||||
// 2026-08-21).
|
||||
if len(names) > 0 && len(failed) == len(names) {
|
||||
writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "no MCP server reachable")
|
||||
return
|
||||
}
|
||||
|
||||
sort.Slice(merged, func(i, j int) bool { return merged[i].Name < merged[j].Name })
|
||||
writeJSONRPCResult(w, req.ID, mcpToolsListResult{Tools: merged})
|
||||
}
|
||||
|
||||
@@ -441,9 +736,10 @@ type mcpToolCallParams struct {
|
||||
// "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.
|
||||
// configured server the caller's group simply cannot reach, or one that
|
||||
// speaks a transport federation excludes entirely, since allowedNames
|
||||
// (allowedMCPServerNames) already excludes both; this is what makes
|
||||
// tools/call authz- and transport-aware without a second, separate check.
|
||||
func resolveFederatedTool(fullName string, allowedNames []string) (serverName, toolName string, ok bool) {
|
||||
best := ""
|
||||
for _, name := range allowedNames {
|
||||
@@ -458,13 +754,20 @@ func resolveFederatedTool(fullName string, allowedNames []string) (serverName, t
|
||||
}
|
||||
|
||||
// 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).
|
||||
// server (resolveFederatedTool), forwards the call there — under its own
|
||||
// toolsCallBackendTimeout budget (context.WithTimeout off r.Context(),
|
||||
// MF1) — 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, one this caller's group cannot
|
||||
// reach, or one on a transport federation excludes entirely).
|
||||
//
|
||||
// countTargetRequest (limits.go) attributes exactly one request to the
|
||||
// single resolved server, once the call was actually attempted against
|
||||
// it — the same "attempted, not necessarily succeeded" accounting
|
||||
// mcpFederatedToolsList's own doc comment explains for its own fan-out.
|
||||
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 {
|
||||
@@ -478,13 +781,28 @@ func (g *Gateway) mcpFederatedToolsCall(w http.ResponseWriter, r *http.Request,
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), toolsCallBackendTimeout)
|
||||
defer cancel()
|
||||
targetURL := g.cfg.MCPServers[serverName].URL
|
||||
resp, err := g.mcpBackendCall(r.Context(), targetURL, "tools/call", mcpToolCallParams{Name: toolName, Arguments: params.Arguments})
|
||||
resp, err := g.mcpBackendCall(ctx, 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
|
||||
}
|
||||
|
||||
// Defensive final guard (MF2, review round 2, 2026-08-21): mcpBackendCall
|
||||
// can only return a nil-err resp whose Result/Error are both nil if some
|
||||
// future code path bypasses doBackendJSONRPC's own parseBackendJSONRPC
|
||||
// validity check (which already rejects this shape as an error) — this
|
||||
// is unreachable today, kept as the explicit guard the review round
|
||||
// asked for rather than trusting that invariant silently.
|
||||
if resp.Result == nil && resp.Error == nil {
|
||||
g.logf("federated tools/call: server %q returned a JSON-RPC response with neither result nor error", serverName)
|
||||
writeJSONRPCErrorResponse(w, req.ID, jsonrpcInternalError, "invalid upstream response")
|
||||
return
|
||||
}
|
||||
|
||||
writeJSONRPCEnvelope(w, jsonrpcResponse{ID: req.ID, Result: resp.Result, Error: resp.Error})
|
||||
}
|
||||
|
||||
@@ -4,8 +4,11 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -181,6 +184,84 @@ func newMockJSONRPCServer(t *testing.T, tools []mcpTool) *mockJSONRPCServer {
|
||||
return m
|
||||
}
|
||||
|
||||
// sessionRequiredMockServer simulates a backend that rejects a bare
|
||||
// (session-less) call with a JSON-RPC error but honors a proper
|
||||
// initialize -> session -> retry handshake — modeling readitall/fetch's
|
||||
// real, live-probed behavior (mcpBackendCall's own doc comment). It
|
||||
// records the exact sequence of calls it received (method plus whichever
|
||||
// session header, if any, accompanied it) so a test can assert the
|
||||
// handshake shape precisely: bare attempt, initialize, retried real call
|
||||
// carrying the session header, best-effort DELETE.
|
||||
type sessionRequiredMockServer struct {
|
||||
srv *httptest.Server
|
||||
sessionID string
|
||||
tools []mcpTool
|
||||
calls []string
|
||||
mu sync.Mutex
|
||||
failInit bool
|
||||
}
|
||||
|
||||
func newSessionRequiredMockServer(t *testing.T, tools []mcpTool) *sessionRequiredMockServer {
|
||||
t.Helper()
|
||||
m := &sessionRequiredMockServer{tools: tools, sessionID: "sess-1"}
|
||||
m.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotSession := r.Header.Get(mcpSessionHeader)
|
||||
|
||||
if r.Method == http.MethodDelete {
|
||||
m.mu.Lock()
|
||||
m.calls = append(m.calls, "DELETE:session="+gotSession)
|
||||
m.mu.Unlock()
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
|
||||
var req jsonrpcRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
t.Fatalf("session-required mock: decode request: %v", err)
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.calls = append(m.calls, req.Method+":session="+gotSession)
|
||||
m.mu.Unlock()
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if req.Method == "initialize" {
|
||||
if m.failInit {
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "initialize refused"}})
|
||||
return
|
||||
}
|
||||
w.Header().Set(mcpSessionHeader, m.sessionID)
|
||||
result, _ := json.Marshal(map[string]any{"protocolVersion": defaultMCPProtocolVersion, "capabilities": map[string]any{}})
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||
return
|
||||
}
|
||||
|
||||
if gotSession != m.sessionID {
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required: call initialize first"}})
|
||||
return
|
||||
}
|
||||
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})
|
||||
_ = 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
|
||||
}
|
||||
|
||||
func (m *sessionRequiredMockServer) callSequence() []string {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return append([]string(nil), m.calls...)
|
||||
}
|
||||
|
||||
// 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
|
||||
@@ -703,3 +784,699 @@ func TestHandleMCPFederated_RateLimited_Returns429(t *testing.T) {
|
||||
t.Fatalf("second request status = %d, want 429, body=%s", rec2.Code, rec2.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// --- MF1: per-backend-call timeouts ---
|
||||
|
||||
func TestBackendTimeoutConstants(t *testing.T) {
|
||||
if toolsListBackendTimeout != 20*time.Second {
|
||||
t.Errorf("toolsListBackendTimeout = %v, want 20s", toolsListBackendTimeout)
|
||||
}
|
||||
if toolsCallBackendTimeout != 120*time.Second {
|
||||
t.Errorf("toolsCallBackendTimeout = %v, want 120s", toolsCallBackendTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext
|
||||
// proves the per-server context.WithTimeout(r.Context(), toolsListBackendTimeout)
|
||||
// (MF1) actually derives its deadline from the INCOMING request's own
|
||||
// context, not a freestanding timer: context.WithTimeout always resolves
|
||||
// to the EARLIER of its parent's existing deadline and its own duration,
|
||||
// so giving the incoming request a context with a much shorter deadline
|
||||
// than toolsListBackendTimeout (20s) and confirming the whole call
|
||||
// returns in well under a second — against a backend that hangs
|
||||
// indefinitely — proves the wiring without ever waiting anywhere near the
|
||||
// real 20s budget in this test suite.
|
||||
func TestHandleMCPFederated_ToolsList_SlowBackend_BoundedByRequestContext(t *testing.T) {
|
||||
block := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
<-block
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
t.Cleanup(func() { close(block) }) // must unblock the handler BEFORE srv.Close (which waits for it) — registered after, so it runs first (t.Cleanup is LIFO)
|
||||
|
||||
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true) // alphaOnly: only the hanging server is reachable
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")})
|
||||
shortCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
req = req.WithContext(shortCtx)
|
||||
|
||||
start := time.Now()
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if elapsed > 3*time.Second {
|
||||
t.Fatalf("elapsed = %v, want well under toolsListBackendTimeout (20s) — the per-server timeout must derive from r.Context(), not ignore it", elapsed)
|
||||
}
|
||||
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 internal error (the only allowed server timed out, degrading to 'no MCP server reachable')", got.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsCall_SlowBackend_BoundedByRequestContext
|
||||
// mirrors the tools/list timeout test above for tools/call's own,
|
||||
// separate context.WithTimeout(r.Context(), toolsCallBackendTimeout) call
|
||||
// site.
|
||||
func TestHandleMCPFederated_ToolsCall_SlowBackend_BoundedByRequestContext(t *testing.T) {
|
||||
block := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
<-block
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
t.Cleanup(func() { close(block) })
|
||||
|
||||
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_lookup"})
|
||||
req := newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams})
|
||||
shortCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
req = req.WithContext(shortCtx)
|
||||
|
||||
start := time.Now()
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if elapsed > 3*time.Second {
|
||||
t.Fatalf("elapsed = %v, want well under toolsCallBackendTimeout (120s) — the per-call timeout must derive from r.Context(), not ignore it", elapsed)
|
||||
}
|
||||
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 internal error (the resolved server timed out)", got.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// --- MF2: non-2xx status and empty-envelope guards ---
|
||||
|
||||
// TestHandleMCPFederated_ToolsCall_BackendNon2xxStatus_IsError proves a
|
||||
// non-2xx upstream response is always an error, even when its body
|
||||
// happens to be a validly-shaped JSON-RPC success object — a 5xx (or any
|
||||
// non-2xx) status is a transport-level failure the body's own content can
|
||||
// never override.
|
||||
func TestHandleMCPFederated_ToolsCall_BackendNon2xxStatus_IsError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","result":{}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
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("1"), 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 internal error — a 5xx status must never be rescued by a well-shaped body", got.Error)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsCall_BackendEmptyEnvelope_ConvertsToInternalError
|
||||
// is the reviewer's named silent-failure scenario (MF2): a backend answers
|
||||
// 200 with valid JSON that is NOT a real JSON-RPC response (neither
|
||||
// "result" nor "error") — without the guard, this decodes into an
|
||||
// all-zero-fields jsonrpcResponse and the client would silently receive
|
||||
// {"jsonrpc":"2.0","id":1} with no indication anything went wrong.
|
||||
func TestHandleMCPFederated_ToolsCall_BackendEmptyEnvelope_ConvertsToInternalError(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"federated","status":"ok"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
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("1"), 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.Fatal("want a JSON-RPC error, not a silent empty success — this is the reviewer's exact silent-failure scenario")
|
||||
}
|
||||
if got.Error.Code != jsonrpcInternalError {
|
||||
t.Errorf("error code = %d, want %d", got.Error.Code, jsonrpcInternalError)
|
||||
}
|
||||
if len(got.Result) != 0 {
|
||||
t.Errorf(`result = %s, want empty — must never emit {"jsonrpc":"2.0","id":1} silently`, got.Result)
|
||||
}
|
||||
}
|
||||
|
||||
// --- MF3: bare-first-with-fallback handshake retry ---
|
||||
|
||||
// TestMcpBackendCall_HandshakeFallbackTrigger proves the fallback fires
|
||||
// ONLY on a JSON-RPC-level error or an HTTP 4xx status from the bare
|
||||
// attempt — never on a 5xx — table-driven over exactly the two trigger
|
||||
// shapes and the one non-trigger shape (a network failure is covered
|
||||
// separately below, since it needs a different mock shape entirely: no
|
||||
// server to answer at all).
|
||||
func TestMcpBackendCall_HandshakeFallbackTrigger(t *testing.T) {
|
||||
cases := []struct {
|
||||
bareHandler http.HandlerFunc
|
||||
name string
|
||||
wantHandshake bool
|
||||
}{
|
||||
{
|
||||
name: "bare JSON-RPC error triggers the handshake",
|
||||
bareHandler: func(w http.ResponseWriter, r *http.Request) {
|
||||
var req jsonrpcRequest
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}})
|
||||
},
|
||||
wantHandshake: true,
|
||||
},
|
||||
{
|
||||
name: "HTTP 403 with no JSON-RPC body triggers the handshake",
|
||||
bareHandler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
},
|
||||
wantHandshake: true,
|
||||
},
|
||||
{
|
||||
name: "HTTP 500 never triggers the handshake",
|
||||
bareHandler: func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
},
|
||||
wantHandshake: false,
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var callCount atomic.Int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if callCount.Add(1) == 1 {
|
||||
tc.bareHandler(w, r)
|
||||
return
|
||||
}
|
||||
if r.Method == http.MethodDelete {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
var req jsonrpcRequest
|
||||
_ = json.Unmarshal(body, &req)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
result, _ := json.Marshal(map[string]any{})
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newFederationTestConfig("http://alpha.invalid", "http://beta.invalid", false)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
|
||||
_, _ = gw.mcpBackendCall(context.Background(), srv.URL, "tools/list", struct{}{})
|
||||
|
||||
gotCalls := callCount.Load()
|
||||
if tc.wantHandshake && gotCalls < 2 {
|
||||
t.Errorf("calls = %d, want >=2 (bare attempt + initialize handshake)", gotCalls)
|
||||
}
|
||||
if !tc.wantHandshake && gotCalls != 1 {
|
||||
t.Errorf("calls = %d, want exactly 1 (bare attempt only, no handshake retry against a 5xx)", gotCalls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake proves a
|
||||
// connection-level failure (nothing listening at all) never triggers the
|
||||
// handshake fallback either — the same "not retry-eligible" rule as a
|
||||
// 5xx, verified against a genuinely different failure shape (doBackendJSONRPC
|
||||
// never even gets an HTTP status back).
|
||||
func TestMcpBackendCall_NetworkFailure_NeverAttemptsHandshake(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)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
|
||||
_, err = gw.mcpBackendCall(context.Background(), "http://127.0.0.1:1", "tools/list", struct{}{})
|
||||
if err == nil {
|
||||
t.Fatal("want an error (nothing listening on 127.0.0.1:1)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry
|
||||
// is MF3's core proof: a server that rejects a bare tools/list call
|
||||
// (readitall/fetch's real, live-probed behavior — sessionRequiredMockServer's
|
||||
// own doc comment) still ends up contributing its tools to the federated
|
||||
// aggregate, via exactly the documented bare -> initialize -> retry ->
|
||||
// close sequence.
|
||||
func TestHandleMCPFederated_ToolsList_HandshakeFallback_SucceedsAfterSessionRetry(t *testing.T) {
|
||||
strict := newSessionRequiredMockServer(t, []mcpTool{{Name: "search"}})
|
||||
cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true) // alphaOnly: only "alpha" (the strict server) is reachable
|
||||
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 (the handshake fallback should have recovered)", 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_search" {
|
||||
t.Fatalf("tools = %+v, want exactly alpha_search", result.Tools)
|
||||
}
|
||||
|
||||
seq := strict.callSequence()
|
||||
if len(seq) != 4 {
|
||||
t.Fatalf("call sequence = %v, want exactly 4 calls (bare attempt, initialize, retried real call, best-effort close — the DELETE runs synchronously inside mcpBackendCall's own defer, before it returns)", seq)
|
||||
}
|
||||
if seq[0] != "tools/list:session=" {
|
||||
t.Errorf("call[0] = %q, want a bare tools/list attempt with no session header", seq[0])
|
||||
}
|
||||
if seq[1] != "initialize:session=" {
|
||||
t.Errorf("call[1] = %q, want an initialize handshake with no session header", seq[1])
|
||||
}
|
||||
if seq[2] != "tools/list:session=sess-1" {
|
||||
t.Errorf("call[2] = %q, want the retried tools/list carrying the session header the handshake returned", seq[2])
|
||||
}
|
||||
if seq[3] != "DELETE:session=sess-1" {
|
||||
t.Errorf("call[3] = %q, want the best-effort session close carrying the same session header", seq[3])
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsCall_HandshakeFallback_SucceedsAfterSessionRetry
|
||||
// mirrors the tools/list handshake test above for tools/call's own
|
||||
// single-target path, additionally asserting the best-effort session
|
||||
// close (DELETE) actually happened.
|
||||
func TestHandleMCPFederated_ToolsCall_HandshakeFallback_SucceedsAfterSessionRetry(t *testing.T) {
|
||||
strict := newSessionRequiredMockServer(t, nil)
|
||||
cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
callParams, _ := json.Marshal(mcpToolCallParams{Name: "alpha_search"})
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage(`"client-7"`), 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 (the handshake fallback should have recovered)", got.Error)
|
||||
}
|
||||
if string(got.ID) != `"client-7"` {
|
||||
t.Errorf("id = %s, want the client's own id echoed back", got.ID)
|
||||
}
|
||||
|
||||
seq := strict.callSequence()
|
||||
if !slices.Contains(seq, "tools/call:session=") {
|
||||
t.Errorf("call sequence = %v, want a bare tools/call attempt with no session header", seq)
|
||||
}
|
||||
if !slices.Contains(seq, "initialize:session=") {
|
||||
t.Errorf("call sequence = %v, want an initialize handshake", seq)
|
||||
}
|
||||
if !slices.Contains(seq, "tools/call:session=sess-1") {
|
||||
t.Errorf("call sequence = %v, want the retried tools/call carrying the session header", seq)
|
||||
}
|
||||
if !slices.Contains(seq, "DELETE:session=sess-1") {
|
||||
t.Errorf("call sequence = %v, want a best-effort DELETE closing the session", seq)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsList_HandshakeFallback_GivesUpWhenInitializeAlsoFails
|
||||
// proves a server that rejects the handshake's own initialize call too
|
||||
// ends up in the failed/skipped list (not a crash, not a hang) — and,
|
||||
// since it is the only allowed server, the aggregate degrades all the way
|
||||
// to the "no MCP server reachable" loud failure (MF3), not an empty
|
||||
// success.
|
||||
func TestHandleMCPFederated_ToolsList_HandshakeFallback_GivesUpWhenInitializeAlsoFails(t *testing.T) {
|
||||
strict := newSessionRequiredMockServer(t, []mcpTool{{Name: "search"}})
|
||||
strict.failInit = true
|
||||
cfg := newFederationTestConfig(strict.srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
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 || got.Error.Code != jsonrpcInternalError || got.Error.Message != "no MCP server reachable" {
|
||||
t.Errorf("error = %+v, want internal error %q", got.Error, "no MCP server reachable")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsList_AllServersFail_ReturnsInternalError
|
||||
// proves that when every attempted server fails, the response is a loud
|
||||
// JSON-RPC error, not a silently-empty tools list that would look
|
||||
// identical to "this caller's group has no MCP access at all".
|
||||
func TestHandleMCPFederated_ToolsList_AllServersFail_ReturnsInternalError(t *testing.T) {
|
||||
cfg := newFederationTestConfig("http://127.0.0.1:1", "http://127.0.0.1:2", 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 || got.Error.Code != jsonrpcInternalError || got.Error.Message != "no MCP server reachable" {
|
||||
t.Errorf("error = %+v, want internal error %q", got.Error, "no MCP server reachable")
|
||||
}
|
||||
}
|
||||
|
||||
// --- MF3: legacy HTTP+SSE transport exclusion ---
|
||||
|
||||
func TestIsLegacySSETransportURL(t *testing.T) {
|
||||
cases := []struct {
|
||||
url string
|
||||
want bool
|
||||
}{
|
||||
{"http://playwright.internal:8080/sse", true},
|
||||
{"http://playwright.internal:8080/sse?foo=bar", true},
|
||||
{"http://mcp.internal:8080", false},
|
||||
{"http://mcp.internal:8080/mcp", false},
|
||||
{"http://mcp.internal:8080/sse/extra", false},
|
||||
{"://not a valid url", false}, // malformed -> safe default: not excluded
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.url, func(t *testing.T) {
|
||||
if got := isLegacySSETransportURL(tc.url); got != tc.want {
|
||||
t.Errorf("isLegacySSETransportURL(%q) = %v, want %v", tc.url, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllowedMCPServerNames_ExcludesLegacySSETransport(t *testing.T) {
|
||||
cfg := CreateConfig()
|
||||
cfg.MCPServers = map[string]*TargetConfig{
|
||||
"alpha": {URL: "http://alpha.internal"},
|
||||
"playwright": {URL: "http://playwright.internal/sse"},
|
||||
"beta": {URL: "http://beta.internal"},
|
||||
}
|
||||
grp := &group{name: "g"} // unrestricted mcpServers: matches everything except the transport exclusion
|
||||
got := allowedMCPServerNames(cfg, grp)
|
||||
want := []string{"alpha", "beta"}
|
||||
if !slices.Equal(got, want) {
|
||||
t.Errorf("allowedMCPServerNames = %v, want %v (playwright excluded: legacy HTTP+SSE transport)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsList_ExcludesLegacySSEServer_NeverContacted
|
||||
// proves a legacy-SSE-transport server is not merely omitted from the
|
||||
// aggregate's RESULT, but never contacted by the fan-out at all.
|
||||
func TestHandleMCPFederated_ToolsList_ExcludesLegacySSEServer_NeverContacted(t *testing.T) {
|
||||
alpha := newMockJSONRPCServer(t, []mcpTool{{Name: "lookup"}})
|
||||
legacy := newMockJSONRPCServer(t, []mcpTool{{Name: "browse"}})
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.MCPServers = map[string]*TargetConfig{
|
||||
"alpha": {URL: alpha.srv.URL},
|
||||
"playwright": {URL: legacy.srv.URL + "/sse"},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/list", ID: json.RawMessage("1")}))
|
||||
var got jsonrpcResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
var result mcpToolsListResult
|
||||
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||
t.Fatalf("decode result: %v", err)
|
||||
}
|
||||
if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_lookup" {
|
||||
t.Fatalf("tools = %+v, want exactly alpha_lookup (playwright excluded)", result.Tools)
|
||||
}
|
||||
if legacy.called.Load() != 0 {
|
||||
t.Error("the legacy-SSE-transport server must never be contacted by tools/list fan-out")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleMCPFederated_ToolsCall_LegacySSEServer_UnresolvableLikeUnconfigured
|
||||
// proves tools/call cannot route to a legacy-SSE-transport server either —
|
||||
// unresolvable, same as an unconfigured name, and never contacted.
|
||||
func TestHandleMCPFederated_ToolsCall_LegacySSEServer_UnresolvableLikeUnconfigured(t *testing.T) {
|
||||
legacy := newMockJSONRPCServer(t, nil)
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}}
|
||||
cfg.MCPServers = map[string]*TargetConfig{"playwright": {URL: legacy.srv.URL + "/sse"}}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}}
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
callParams, _ := json.Marshal(mcpToolCallParams{Name: "playwright_click"})
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "tools/call", ID: json.RawMessage("1"), Params: callParams}))
|
||||
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 invalid params (legacy-SSE server must be unresolvable via tools/call, same as unconfigured)", got.Error)
|
||||
}
|
||||
if legacy.called.Load() != 0 {
|
||||
t.Error("the legacy-SSE-transport server must never be contacted")
|
||||
}
|
||||
}
|
||||
|
||||
// --- SF6: protocol-version allowlist ---
|
||||
|
||||
func TestHandleMCPFederated_Initialize_AllowlistedVersionsEchoedVerbatim(t *testing.T) {
|
||||
for version := range knownMCPProtocolVersions {
|
||||
t.Run(version, func(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": version})
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1"), Params: params}))
|
||||
|
||||
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 != version {
|
||||
t.Errorf("protocolVersion = %q, want the allowlisted caller-supplied version %q echoed verbatim", result.ProtocolVersion, version)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleMCPFederated_Initialize_UnrecognizedVersion_FallsBackToDefault(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": "1999-01-01"})
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newFederatedRequest(t, "sk-alice", jsonrpcRequest{JSONRPC: jsonrpcVersion, Method: "initialize", ID: json.RawMessage("1"), Params: params}))
|
||||
|
||||
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 for an unrecognized caller-supplied version", result.ProtocolVersion, defaultMCPProtocolVersion)
|
||||
}
|
||||
}
|
||||
|
||||
// --- SF8: explicit 405 for a disallowed method on exactly /mcp ---
|
||||
|
||||
func TestHandleMCPFederated_DisallowedMethod_Returns405WithAllowHeader(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)
|
||||
}
|
||||
|
||||
for _, method := range []string{http.MethodGet, http.MethodDelete, http.MethodPut} {
|
||||
t.Run(method, func(t *testing.T) {
|
||||
req := httptest.NewRequest(method, federatedMCPPath, nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusMethodNotAllowed {
|
||||
t.Fatalf("status = %d, want 405, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if got := rec.Header().Get("Allow"); got != http.MethodPost {
|
||||
t.Errorf("Allow header = %q, want %q", got, http.MethodPost)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("no auth still 405 (method check precedes auth)", func(t *testing.T) {
|
||||
req := httptest.NewRequest(http.MethodGet, federatedMCPPath, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusMethodNotAllowed {
|
||||
t.Fatalf("status = %d, want 405 even without a valid API key, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller proves
|
||||
// mcpBackendCloseSession's own "best-effort... failure logged, never
|
||||
// surfaced" contract: the server accepts the handshake and the retried
|
||||
// real call normally, but drops the connection outright on the
|
||||
// best-effort DELETE — the caller must still see the real call's own
|
||||
// successful result, not a failure manufactured by cleanup.
|
||||
func TestMcpBackendCall_SessionCloseFailure_NeverSurfacedToCaller(t *testing.T) {
|
||||
const sessionID = "sess-close-fail"
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodDelete {
|
||||
// A network-level failure on close, not merely a non-2xx
|
||||
// status: mcpBackendCloseSession never even inspects the
|
||||
// response status, only whether Do() itself returned an
|
||||
// error — hijacking and dropping the raw connection is what
|
||||
// actually exercises that path.
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
t.Fatal("mock server: ResponseWriter does not support Hijack")
|
||||
}
|
||||
conn, _, err := hj.Hijack()
|
||||
if err != nil {
|
||||
t.Fatalf("hijack: %v", err)
|
||||
}
|
||||
_ = conn.Close()
|
||||
return
|
||||
}
|
||||
|
||||
var req jsonrpcRequest
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch {
|
||||
case req.Method == "initialize":
|
||||
w.Header().Set(mcpSessionHeader, sessionID)
|
||||
result, _ := json.Marshal(map[string]any{})
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||
case req.Method == "tools/list" && r.Header.Get(mcpSessionHeader) == "":
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Error: &jsonrpcError{Code: -32000, Message: "session required"}})
|
||||
default:
|
||||
result, _ := json.Marshal(mcpToolsListResult{Tools: []mcpTool{{Name: "search"}}})
|
||||
_ = json.NewEncoder(w).Encode(jsonrpcResponse{ID: req.ID, Result: result})
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newFederationTestConfig(srv.URL, "http://beta.invalid", true)
|
||||
h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
|
||||
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 failed best-effort session close must never surface as a call failure", got.Error)
|
||||
}
|
||||
var result mcpToolsListResult
|
||||
if err := json.Unmarshal(got.Result, &result); err != nil {
|
||||
t.Fatalf("decode result: %v", err)
|
||||
}
|
||||
if len(result.Tools) != 1 || result.Tools[0].Name != "alpha_search" {
|
||||
t.Fatalf("tools = %+v, want alpha_search despite the session-close failure", result.Tools)
|
||||
}
|
||||
}
|
||||
|
||||
+19
-10
@@ -244,10 +244,15 @@ func (g *Gateway) runUnified(w http.ResponseWriter, r *http.Request, u *user, gr
|
||||
// contract (admin.go) that the dashboard's usage table already promised.
|
||||
//
|
||||
// Callers metering actual LLM traffic wrap this result in withTotalScope
|
||||
// before passing it to checkAndCount/account; handleAdminAPI (admin.go)
|
||||
// calls this directly, without withTotalScope, so an admin request's own
|
||||
// req/min-req/day accounting never contributes to the total-scope
|
||||
// LLM-traffic series (v0.2 data-layer task).
|
||||
// before passing it to checkAndCount/account. handleAdminAPI (admin.go)
|
||||
// never calls buildLimitScopes at all — its own buildAdminUsage builds
|
||||
// {kind, id, limits} literals directly from authStore.snapshot's user/group
|
||||
// listing, for currentUsage's read-only purposes — and, separately, never
|
||||
// calls checkAndCount either (handleAdminAPI's own doc comment: admin
|
||||
// traffic must never move req/min-req/day statistics). Fixed a stale
|
||||
// version of this comment (review round 2, 2026-08-21) that claimed
|
||||
// handleAdminAPI "calls this directly, without withTotalScope" — it never
|
||||
// called this function at all, on any version of this file.
|
||||
func buildLimitScopes(u *user, grp *group) []limitScope {
|
||||
return []limitScope{
|
||||
{limits: u.limits, kind: "user", id: u.name},
|
||||
@@ -259,12 +264,16 @@ func buildLimitScopes(u *user, grp *group) []limitScope {
|
||||
// (totalScopeKind/totalScopeID, limits.go) appended — the single helper
|
||||
// every metered route (runUnified above; resolveMediaRequest,
|
||||
// routes_media.go; handlePassthrough, routes_passthrough.go;
|
||||
// handleTargetProxy, mcp_a2a.go) calls around its own buildLimitScopes
|
||||
// result, so none of them can forget it and none of them duplicate the
|
||||
// scope literal. Deliberately not folded into buildLimitScopes itself:
|
||||
// that helper is also called directly by handleAdminAPI (admin.go) for
|
||||
// its own req/min-req/day accounting, and admin traffic must never
|
||||
// contribute to the total LLM-traffic scope.
|
||||
// handleTargetProxy, mcp_a2a.go; handleMCPFederated, mcp_federation.go)
|
||||
// calls around its own buildLimitScopes result, so none of them can forget
|
||||
// it and none of them duplicate the scope literal. Deliberately not folded
|
||||
// into buildLimitScopes itself: buildLimitScopes' own doc comment now
|
||||
// explains why admin traffic needs no special-casing here at all —
|
||||
// handleAdminAPI never calls either function, so there was never a real
|
||||
// "admin traffic must not contribute to total" case for this split to
|
||||
// guard against; the split is kept anyway because buildLimitScopes'
|
||||
// {user, group} pair and the synthetic total scope are conceptually
|
||||
// different additions, worth two names.
|
||||
func withTotalScope(scopes []limitScope) []limitScope {
|
||||
return append(scopes, limitScope{kind: totalScopeKind, id: totalScopeID, limits: nil})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user