From 6fc38ca16b7634e9083f1fedc8827f967229e2e9 Mon Sep 17 00:00:00 2001 From: Eugene Samotija Date: Sat, 25 Jul 2026 13:32:17 +0300 Subject: [PATCH] upstream: re-initialize when the server expired our session (404) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The retry-once path only covered locally dropped transports. When an upstream expires or forgets the streamable-HTTP session server-side (restart, TTL — hit in production with Publora), the SDK client still looks connected, the server answers 404 "Unknown or expired MCP session", and the gateway propagated that error to the caller. Per the MCP spec a 404 on an existing session means re-initialize: the connection now detects it (StreamableHTTPError code 404, plus a message fallback for untyped rethrows), drops the pooled client without triggering backoff, reconnects and retries the call exactly once. Genuine tool errors still pass through untouched. Integration test runs a real StreamableHTTPServerTransport backend whose session store is cleared mid-test to reproduce the failure. Co-Authored-By: Claude Fable 5 --- CLAUDE.md | 2 +- .../gateway/src/upstream/connection.test.ts | 122 ++++++++++++++++++ packages/gateway/src/upstream/connection.ts | 48 ++++++- 3 files changed, 165 insertions(+), 7 deletions(-) create mode 100644 packages/gateway/src/upstream/connection.test.ts 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`);