diff --git a/CLAUDE.md b/CLAUDE.md index e3dc10a..d1ea4e0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -18,7 +18,7 @@ Self-hosted MCP manager/gateway: one streamable-HTTP `/mcp` endpoint federating - `domain/policy.ts` — `PolicyService`: toolEnabled ∧ (override(allow) ∨ (tier ≤ maxTier ∧ ¬deny)); maxTier = per-upstream grant ?? role default. Same function gates tools/list AND tools/call. `allowsFor(principal, entry)` = envelope ∧ personal prefs (deny-only rows in `user_prefs`; "enable" deletes the row — narrowing can never widen) - `auth/` — `static-tokens.ts` (timing-safe bearer match), `oidc.ts` (jose JWKS resource-server verifier for inbound *access* tokens), `login.ts` (interactive login: openid-client cookie+PKCE confidential-client flow consuming an *id-token*; signed identity-only session cookie, HMAC + freshness; `safeReturnTo`), `authz-server.ts` (OAuth AS facade: RFC 8414 metadata, RFC 7591 DCR for public clients, single-use hashed 60s codes + PKCE S256, HS256 gateway JWTs keyed by `GATEWAY_JWT_SECRET` (default derived from `SESSION_SECRET`), rotating refresh tokens — 30d sliding, family-revoked on replay, client-bound consume that can't burn a live token — register rate limit; clients managed via `/api/oauth-clients` + Users tab), `prm.ts` (RFC 9728 doc + WWW-Authenticate; lists the gateway itself as AS when login is configured, else the raw IdP), `directory.ts` (app-only Graph search of Entra users/groups via the login app's own creds — powers the admin UI group-mapping typeahead at `/api/directory/search`; null for non-Entra issuers → UI degrades to paste-an-id), `principal.ts` (session binding key). Four inbound auth paths in `createAuthResolver`: static token, gateway-issued JWT (routed by unverified `iss == PUBLIC_URL`, then fully verified), OIDC bearer, and the cookie session — the cookie/JWT carry only identity and the role is re-resolved every request (persisted at callback via `setUserRole`), so a session id never carries privilege. `loginUpsert()` is shared by the bearer + callback paths so they can't drift. `/oauth/authorize` brokers user auth to Entra by piggybacking the interactive login: the pending request rides in the signed transient cookie and `/auth/callback` mints the code. - `secrets/` — `SecretStore` interface (scheme-tagged: `bao` | `kv`), `openbao.ts` (KV v2, AppRole or token, 5-min cache), `keyvault.ts` (Azure Key Vault, `DefaultAzureCredential`, lazy SDK import, same 5-min cache; `put(path, field)` writes `path-field`), `memory.ts` (tests). Refs: `bao:path#field` / `kv:secret-name`; env refs: `${VAR}` — all resolved only at upstream connect time. One store at a time (`BAO_ADDR` xor `KEY_VAULT_URI`) -- `upstream/connection.ts` — one pooled SDK `Client` per upstream; header/env injection; backoff reconnect (1s→60s) + `onRecovered`; retry-once on dropped transport +- `upstream/connection.ts` — one pooled SDK `Client` per upstream; header/env injection; backoff reconnect (1s→60s) + `onRecovered`; retry-once on dropped transport AND on server-side session expiry (upstream 404 "unknown session" → transparent re-initialize + retry, per MCP spec) - `upstream/manager.ts` also pools **per-principal links** for `sessionMode:"per-user"` upstreams: spec clone with the caller's credential REFS layered over headers/env (still resolved via the secret store at connect — anti-passthrough intact); catalog discovery stays on the shared link; personal pool flushed on upstream upsert/remove. `requirePersonalCredentials` refuses the shared fallback - `upstream/manager.ts` — policy-free catalog owner; hot `upsertUpstream`/`removeUpstream`; `summaries()` for the UI - `mcp/gateway-server.ts` — low-level SDK `Server` per session, closes over the Principal; unknown and forbidden tools get the same error (no existence oracle) diff --git a/packages/gateway/src/upstream/connection.test.ts b/packages/gateway/src/upstream/connection.test.ts new file mode 100644 index 0000000..0186cbb --- /dev/null +++ b/packages/gateway/src/upstream/connection.test.ts @@ -0,0 +1,122 @@ +/** + * Integration test for UpstreamConnection against a real streamable-HTTP MCP + * server — specifically the stale-session path: the upstream expires/forgets + * the HTTP session server-side (server restart, TTL) while the local + * transport still looks healthy. Per the MCP spec the server answers 404 and + * the client must re-initialize; the connection does that transparently and + * retries the call once. + */ + +import { createServer, type IncomingMessage, type Server as HttpServer, type ServerResponse } from "node:http"; +import { randomUUID } from "node:crypto"; +import type { AddressInfo } from "node:net"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; +import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"; +import { UpstreamConnection } from "./connection.js"; +import type { UpstreamSpec } from "../config.js"; + +/** Minimal stateful streamable-HTTP MCP server whose sessions we can expire. */ +class FakeUpstream { + private readonly sessions = new Map(); + private http!: HttpServer; + initializeCount = 0; + + async start(): Promise { + this.http = createServer((req, res) => { + this.handle(req, res).catch((err) => { + res.writeHead(500).end(String(err)); + }); + }); + await new Promise((resolve) => this.http.listen(0, resolve)); + return `http://localhost:${(this.http.address() as AddressInfo).port}/mcp`; + } + + /** Simulate a server-side restart/TTL: all session ids become unknown. */ + expireAllSessions(): void { + this.sessions.clear(); + } + + async stop(): Promise { + await new Promise((resolve) => this.http.close(() => resolve())); + } + + private async handle(req: IncomingMessage, res: ServerResponse): Promise { + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(chunk as Buffer); + const raw = Buffer.concat(chunks).toString("utf8"); + const body: unknown = raw ? JSON.parse(raw) : undefined; + + const sessionId = req.headers["mcp-session-id"] as string | undefined; + const existing = sessionId ? this.sessions.get(sessionId) : undefined; + if (existing) { + await existing.handleRequest(req, res, body); + return; + } + + if (req.method === "POST" && isInitializeRequest(body)) { + this.initializeCount += 1; + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: (id) => this.sessions.set(id, transport), + }); + const mcp = new McpServer({ name: "fake-upstream", version: "0.0.0" }); + mcp.tool("ping", async () => ({ content: [{ type: "text" as const, text: "pong" }] })); + await mcp.connect(transport); + await transport.handleRequest(req, res, body); + return; + } + + // Unknown or expired session — what real servers (and the SDK) return. + res.writeHead(404, { "Content-Type": "application/json" }).end( + JSON.stringify({ error: "Unknown or expired MCP session. Re-initialize." }) + ); + } +} + +const spec = (url: string): UpstreamSpec => + ({ + id: "fake", + namespace: "fake", + transport: "http", + url, + headers: {}, + enabled: true, + }) as UpstreamSpec; + +describe("UpstreamConnection — stale streamable-HTTP sessions", () => { + const upstream = new FakeUpstream(); + let connection: UpstreamConnection; + + beforeAll(async () => { + const url = await upstream.start(); + connection = new UpstreamConnection(spec(url)); + await connection.connect(); + }); + + afterAll(async () => { + await connection.close(); + await upstream.stop(); + }); + + it("calls tools over the pooled session", async () => { + const result = await connection.callTool("ping", {}); + expect(result.content).toEqual([{ type: "text", text: "pong" }]); + expect(upstream.initializeCount).toBe(1); + }); + + it("re-initializes and retries once when the upstream expired the session", async () => { + upstream.expireAllSessions(); + const result = await connection.callTool("ping", {}); + expect(result.content).toEqual([{ type: "text", text: "pong" }]); + expect(upstream.initializeCount).toBe(2); // a fresh session was minted + expect(connection.connected).toBe(true); + }); + + it("does not re-initialize for genuine tool errors", async () => { + const result = await connection.callTool("does-not-exist", {}); + expect(result.isError).toBe(true); // the upstream's error passes through + expect(upstream.initializeCount).toBe(2); // no needless session churn + }); +}); diff --git a/packages/gateway/src/upstream/connection.ts b/packages/gateway/src/upstream/connection.ts index ed98c52..094197a 100644 --- a/packages/gateway/src/upstream/connection.ts +++ b/packages/gateway/src/upstream/connection.ts @@ -19,7 +19,10 @@ */ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; -import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; +import { + StreamableHTTPClientTransport, + StreamableHTTPError, +} from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { StdioClientTransport, getDefaultEnvironment, @@ -40,6 +43,19 @@ import { SERVER_NAME, SERVER_VERSION } from "../version.js"; const BACKOFF_INITIAL_MS = 1_000; const BACKOFF_MAX_MS = 60_000; +/** + * The upstream forgot our streamable-HTTP session while the local transport + * still looks healthy — servers answer 404 (spec: "re-initialize") when they + * restart or expire sessions server-side. The message fallback catches the + * same condition when a proxy or wrapper re-throws it untyped. + */ +const isStaleSession = (err: unknown): boolean => { + if (err instanceof StreamableHTTPError && err.code === 404) return true; + return /unknown or expired mcp session|session not found/i.test( + err instanceof Error ? err.message : String(err) + ); +}; + export interface UpstreamStatus { connected: boolean; lastError: string | null; @@ -186,22 +202,42 @@ export class UpstreamConnection { /** * Call a tool. If the transport dropped (upstream restart, lost HTTP - * session), reconnect and retry exactly once. + * session) or the upstream expired our session server-side (404 per the + * MCP spec), re-initialize and retry exactly once. */ async callTool(name: string, args: Record): Promise { const client = await this.requireClient(); try { return (await client.callTool({ name, arguments: args })) as CallToolResult; } catch (err) { - if (this.connected) throw err; // upstream answered with an error — not a transport drop - console.error( - `[upstream:${this.spec.id}] call "${name}" hit a dropped connection — retrying once` - ); + if (this.connected && !isStaleSession(err)) { + throw err; // upstream answered with an error — not a session/transport problem + } + if (this.connected) { + console.error( + `[upstream:${this.spec.id}] call "${name}" hit an expired upstream session — re-initializing` + ); + await this.resetClient(); + } else { + console.error( + `[upstream:${this.spec.id}] call "${name}" hit a dropped connection — retrying once` + ); + } const fresh = await this.requireClient(); return (await fresh.callTool({ name, arguments: args })) as CallToolResult; } } + /** Drop the pooled client so the next call initializes a fresh session. */ + private async resetClient(): Promise { + const client = this.client; + this.client = null; + if (client) { + client.onclose = undefined; // deliberate reset — no backoff reconnect + await client.close().catch(() => undefined); + } + } + private async requireClient(): Promise { if (!this.client) await this.connect(); if (!this.client) throw new Error(`upstream "${this.spec.id}" is not connected`);