upstream: re-initialize when the server expired our session (404)

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)

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<string, StreamableHTTPServerTransport>();
+  private http!: HttpServer;
+  initializeCount = 0;
+
+  async start(): Promise<string> {
+    this.http = createServer((req, res) => {
+      this.handle(req, res).catch((err) => {
+        res.writeHead(500).end(String(err));
+      });
+    });
+    await new Promise<void>((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<void> {
+    await new Promise<void>((resolve) => this.http.close(() => resolve()));
+  }
+
+  private async handle(req: IncomingMessage, res: ServerResponse): Promise<void> {
+    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
+  });
+});

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<string, unknown>): Promise<CallToolResult> {
     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<void> {
+    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<Client> {
     if (!this.client) await this.connect();
     if (!this.client) throw new Error(`upstream "${this.spec.id}" is not connected`);