@mcp-ts/sdk 2.4.1 → 2.4.3

package/src/server/mcp/multi-session-client.ts

@@ -1,54 +1,90 @@
-
-
 import { MCPClient } from './oauth-client.js';
 import { sessions, type Session } from '../storage/index.js';
+import type { ToolClientProvider } from '../../shared/types.js';
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
 
 const DEFAULT_TIMEOUT_MS = 15000;
 const DEFAULT_MAX_RETRIES = 2;
 const DEFAULT_RETRY_DELAY_MS = 1000;
 const CONNECTION_BATCH_SIZE = 5;
 
-/**
- * Manages multiple MCP connections for a single user.
- * Allows aggregating tools from all connected servers.
- */
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
 export interface MultiSessionOptions {
     /**
-     * Connection timeout in milliseconds
+     * Connection timeout in milliseconds.
      * @default 15000
      */
     timeout?: number;
+
     /**
-     * Maximum number of retry attempts
+     * Maximum number of retry attempts per session.
      * @default 2
      */
     maxRetries?: number;
+
     /**
-     * Delay between retries in milliseconds
+     * Delay between retry attempts in milliseconds.
      * @default 1000
      */
     retryDelay?: number;
+
+    /**
+     * Custom session provider. When provided, `connect()` calls this
+     * instead of querying the storage backend. Useful when sessions are
+     * synced externally (e.g. via Supabase Realtime, WebSocket, or an
+     * in-memory cache).
+     *
+     * Default: reads from the storage backend via `sessions.list(userId)`.
+     */
+    sessionProvider?: () => Promise<Session[]>;
+
+    /**
+     * Called after a session is successfully connected.
+     */
+    onSessionConnected?: (sessionId: string, client: MCPClient) => void;
+
+    /**
+     * Called when a session is evicted from the in-memory client list
+     * because it no longer exists in the active sessions list.
+     */
+    onSessionEvicted?: (sessionId: string) => void;
+
+    /**
+     * Called when all retry attempts for a session have been exhausted.
+     */
+    onSessionFailed?: (sessionId: string, error: unknown) => void;
 }
 
+// ---------------------------------------------------------------------------
+// MultiSessionClient
+// ---------------------------------------------------------------------------
+
 /**
  * Manages multiple MCP client connections for a single user.
  *
- * On a traditional long-running server, you can cache this instance per user
- * so the connections stay alive between requests. On serverless, a new instance
- * will be created per invocation, but the underlying session data is always
- * read from the storage backend so nothing is lost between calls.
+ * Cached instances stay alive between requests on long-running servers.
+ * On serverless, the underlying session data persists in storage — the
+ * client rehydrates from `sessionProvider` (or the default storage backend)
+ * each time `connect()` is called.
+ *
+ * @implements {ToolClientProvider}
  */
-export class MultiSessionClient {
+export class MultiSessionClient implements ToolClientProvider {
     private clients: MCPClient[] = [];
+    private connectionPromises = new Map<string, Promise<void>>();
     private userId: string;
-    private options: MultiSessionOptions;
+    private options: Required<Pick<MultiSessionOptions, 'timeout' | 'maxRetries' | 'retryDelay'>> &
+        Pick<MultiSessionOptions, 'sessionProvider' | 'onSessionConnected' | 'onSessionEvicted' | 'onSessionFailed'>;
 
     /**
-     * Creates a new MultiSessionClient for the given user userId.
-     *
-     * @param userId - A unique string identifying the user (e.g. user ID or email).
-     * @param options  - Optional tuning for connection timeout, retry count, and retry delay.
-     *                   Falls back to sensible defaults if not provided.
+     * @param userId - Unique identifier for the user (e.g. user ID or email).
+     * @param options - Optional tuning and lifecycle hooks.
      */
     constructor(userId: string, options: MultiSessionOptions = {}) {
         this.userId = userId;
@@ -56,36 +92,123 @@ export class MultiSessionClient {
             timeout: DEFAULT_TIMEOUT_MS,
             maxRetries: DEFAULT_MAX_RETRIES,
             retryDelay: DEFAULT_RETRY_DELAY_MS,
-            ...options
+            ...options,
         };
     }
 
+    // -----------------------------------------------------------------------
+    // Public API
+    // -----------------------------------------------------------------------
+
     /**
-     * Fetches all sessions for this userId from storage and returns only the
-     * ones that are ready to connect.
+     * Fetches active sessions and establishes connections to all of them.
      *
-     * A session is considered connectable when:
-     * - It has a `serverId`, `serverUrl`, and `callbackUrl` (i.e. it was fully initialized)
-     * - Its status is `active`. Pending sessions are skipped here
-     *   and let the OAuth flow complete separately before we try to reconnect them.
+     * Call this once after creating the client. On long-running servers you
+     * can cache the `MultiSessionClient` instance and call `connect()` on
+     * each request — already-connected sessions are skipped internally.
      */
-    private async getActiveSessions(): Promise<Session[]> {
-        const sessionList = await sessions.list(this.userId);
-        const valid = sessionList.filter(s =>
+    async connect(): Promise<void> {
+        const sessions = await this.fetchActiveSessions();
+        const activeSessionIds = new Set(sessions.map(s => s.sessionId));
+
+        // Evict clients whose sessions no longer appear in the active list.
+        // Without this, a stale MCPClient with a live SDK transport would
+        // stay in the clients array after its session was deleted externally
+        // (e.g. via the SSE handler's clearSession()), causing OAuth refreshes
+        // to fail with a FK violation when they try to patchCredentials()
+        // against the now-deleted session.
+        for (const client of this.clients) {
+            if (!activeSessionIds.has(client.getSessionId())) {
+                this.options.onSessionEvicted?.(client.getSessionId());
+            }
+        }
+        this.clients = this.clients.filter(c => activeSessionIds.has(c.getSessionId()));
+
+        await this.connectInBatches(sessions);
+    }
+
+    /**
+     * Drops all cached `MCPClient` instances and reconnects fresh from storage.
+     *
+     * Call this when downstream MCP servers have expired their transport sessions
+     * (e.g. after a remote server restart) and subsequent tool calls return
+     * "Session not found. Reconnect without session header." errors.
+     *
+     * OAuth tokens are preserved in the storage backend — no re-authentication
+     * is required. Only the in-memory transport sessions are cleared.
+     */
+    async reconnect(): Promise<void> {
+        await this.disconnect();
+        await this.connect();
+    }
+
+    /**
+     * Returns all currently connected `MCPClient` instances.
+     *
+     * Use this to enumerate available tools across all connected servers,
+     * or to route a tool call to the right client by `serverId`.
+     */
+    getClients(): MCPClient[] {
+        return this.clients;
+    }
+
+    /**
+     * Removes and disconnects a single session by ID.
+     *
+     * @returns `true` if the session was found and removed, `false` if not found.
+     */
+    async removeSession(sessionId: string): Promise<boolean> {
+        const idx = this.clients.findIndex(c => c.getSessionId() === sessionId);
+        if (idx === -1) return false;
+        const [client] = this.clients.splice(idx, 1);
+        await client.disconnect();
+        return true;
+    }
+
+    /**
+     * Gracefully disconnects all active MCP clients and clears the internal list.
+     *
+     * For Streamable HTTP sessions, each client sends an HTTP DELETE to its MCP
+     * endpoint per the spec before closing locally. All disconnects run in
+     * parallel so shutdown is not serialised across many sessions.
+     *
+     * Call this during server shutdown or when a user logs out to free up
+     * underlying transport resources (SSE streams, HTTP connections, etc.).
+     */
+    async disconnect(): Promise<void> {
+        await Promise.all(this.clients.map((client) => client.disconnect()));
+        this.clients = [];
+    }
+
+    // -----------------------------------------------------------------------
+    // Internals
+    // -----------------------------------------------------------------------
+
+    /**
+     * Resolves the list of sessions to connect.
+     *
+     * Uses the custom `sessionProvider` when provided, otherwise falls back
+     * to querying the storage backend via `sessions.list(userId)`.
+     */
+    private async fetchActiveSessions(): Promise<Session[]> {
+        const sessionList = this.options.sessionProvider
+            ? await this.options.sessionProvider()
+            : await sessions.list(this.userId);
+
+        return sessionList.filter(s =>
             s.serverId &&
             s.serverUrl &&
             s.callbackUrl &&
             s.status === 'active'
         );
-        return valid;
     }
 
     /**
-     * Connects to a list of sessions in controlled batches of `CONNECTION_BATCH_SIZE`.
+     * Connects a list of sessions in controlled batches.
      *
-     * Batching prevents overwhelming the event loop or external servers when a user
-     * has many active MCP sessions (e.g. 20+ servers). Within each batch, sessions
-     * are connected concurrently using `Promise.all` for speed.
+     * Batching prevents overwhelming the event loop or external servers when
+     * a user has many active MCP sessions. Within each batch, sessions are
+     * connected concurrently using `Promise.all`.
      */
     private async connectInBatches(sessions: Session[]): Promise<void> {
         for (let i = 0; i < sessions.length; i += CONNECTION_BATCH_SIZE) {
@@ -94,38 +217,30 @@ export class MultiSessionClient {
         }
     }
 
-    private connectionPromises = new Map<string, Promise<void>>();
-
     /**
-     * Connects a single session, with built-in deduplication to prevent race conditions.
+     * Connects a single session, with deduplication to prevent race conditions.
      *
-     * - If a client for this session already exists and is connected, returns immediately.
-     * - If the existing client entry is no longer connected (e.g. it was explicitly
-     *   disconnected), it is evicted so that `establishConnectionWithRetries` creates a
-     *   fresh transport — preventing "Client already connected" errors from the SDK.
-     * - If a connection attempt for this session is already in-flight (e.g. from a
-     *   concurrent call), it joins the existing promise instead of starting a new one.
-     *   This is the key concurrency lock — the `connectionPromises` map acts as a
-     *   per-session mutex so we never spin up two physical connections for the same session.
-     * - On completion (success or failure), the promise is cleaned up from the map.
+     * - If a client for this session already exists and is connected, returns
+     *   immediately.
+     * - If the existing client entry is no longer connected (e.g. explicit
+     *   disconnect), it is evicted so a fresh transport is created.
+     * - If a connection attempt for this session is already in-flight, the
+     *   existing promise is reused as a per-session mutex.
+     * - On completion (success or failure), the promise is cleaned up from
+     *   the connectionPromises map.
      */
     private async connectSession(session: Session): Promise<void> {
         const existing = this.clients.find(c => c.getSessionId() === session.sessionId);
 
         if (existing) {
             if (existing.isConnected()) {
-                // Genuinely connected — nothing to do.
                 return;
             }
 
-            // Client entry exists but is no longer connected (explicit disconnect or
-            // a prior failed reconnect attempt). Remove it so the fresh connect below
-            // starts with a clean slate and the underlying SDK Client doesn't complain
-            // about an already-attached transport.
+            this.options.onSessionEvicted?.(existing.getSessionId());
             this.clients = this.clients.filter(c => c !== existing);
         }
 
-        // Avoid concurrent connection attempts for the same session
         if (this.connectionPromises.has(session.sessionId)) {
             return this.connectionPromises.get(session.sessionId)!;
         }
@@ -142,22 +257,19 @@ export class MultiSessionClient {
     }
 
     /**
-     * The core connection loop for a single session.
+     * Core connection loop for a single session with retry logic.
      *
-     * Attempts to establish a physical MCP connection, retrying up to `maxRetries` times
-     * if the connection fails. Each attempt:
-     * 1. Creates a fresh `MCPClient` instance from the session data.
-     * 2. Races the connect call against a timeout promise — if the server doesn't respond
-     *    within `timeoutMs`, the attempt is aborted and counted as a failure.
-     * 3. On success, replaces any stale client entry for this session in the `clients` array.
+     * 1. Creates a fresh `MCPClient` from the session data.
+     * 2. Races `client.connect()` against a timeout.
+     * 3. On success, replaces any stale entry and fires `onSessionConnected`.
      * 4. On failure, waits `retryDelay` ms before the next attempt.
      *
-     * If all attempts are exhausted, logs an error and returns silently (does not throw),
-     * so a single bad server doesn't block the rest of the batch from connecting.
+     * If all attempts are exhausted, logs an error and returns silently so
+     * a single bad server doesn't block the rest of the batch.
      */
     private async establishConnectionWithRetries(session: Session): Promise<void> {
-        const maxRetries = this.options.maxRetries ?? DEFAULT_MAX_RETRIES;
-        const retryDelay = this.options.retryDelay ?? DEFAULT_RETRY_DELAY_MS;
+        const maxRetries = this.options.maxRetries;
+        const retryDelay = this.options.retryDelay;
         let lastError: unknown;
 
         for (let attempt = 0; attempt <= maxRetries; attempt++) {
@@ -173,10 +285,13 @@ export class MultiSessionClient {
                     headers: session.headers,
                 });
 
-                const timeoutMs = this.options.timeout ?? DEFAULT_TIMEOUT_MS;
+                const timeoutMs = this.options.timeout;
                 let timeoutTimer: ReturnType<typeof setTimeout>;
                 const timeoutPromise = new Promise<never>((_, reject) => {
-                    timeoutTimer = setTimeout(() => reject(new Error(`Connection timed out after ${timeoutMs}ms`)), timeoutMs);
+                    timeoutTimer = setTimeout(
+                        () => reject(new Error(`Connection timed out after ${timeoutMs}ms`)),
+                        timeoutMs,
+                    );
                 });
 
                 try {
@@ -185,10 +300,11 @@ export class MultiSessionClient {
                     clearTimeout(timeoutTimer!);
                 }
 
-                // Always replace the disconnected client entry
                 this.clients = this.clients.filter(c => c.getSessionId() !== session.sessionId);
                 this.clients.push(client);
-                return; // successfully connected
+                this.options.onSessionConnected?.(session.sessionId, client);
+                return;
+
             } catch (error) {
                 lastError = error;
                 if (attempt < maxRetries) {
@@ -197,69 +313,11 @@ export class MultiSessionClient {
             }
         }
 
-        console.error(`[MultiSessionClient] Failed to connect to session ${session.sessionId} after ${maxRetries + 1} attempts:`, lastError);
-    }
-
-    /**
-     * The main entry point. Fetches all active sessions for this userId from
-     * storage and establishes connections to all of them in batches.
-     *
-     * Call this once after creating the client. On traditional servers, you can
-     * cache the `MultiSessionClient` instance after calling `connect()` to avoid
-     * re-fetching and re-connecting on every request.
-     */
-    async connect(): Promise<void> {
-        const sessions = await this.getActiveSessions();
-        const activeSessionIds = new Set(sessions.map(s => s.sessionId));
-
-        // Prune stale clients whose sessions no longer exist in storage.
-        // Otherwise a lingering MCPClient with a still-alive SDK Client
-        // transport triggers OAuth refreshes on listTools(), and the
-        // StorageOAuthClientProvider.state() -> patchCredentials() call
-        // fails with a FK violation because the session was deleted.
-        this.clients = this.clients.filter(c => activeSessionIds.has(c.getSessionId()));
-
-        await this.connectInBatches(sessions);
-    }
-
-    /**
-     * Drops all cached `MCPClient` instances and reconnects fresh from storage.
-     *
-     * Call this when downstream MCP servers have expired their transport sessions
-     * (e.g. after a remote server restart) and subsequent tool calls return
-     * "Session not found. Reconnect without session header." errors.
-     *
-     * OAuth tokens are preserved in the storage backend — no re-authentication
-     * is required. Only the in-memory transport sessions are cleared.
-     */
-    async reconnect(): Promise<void> {
-        await this.disconnect();
-        await this.connect();
-    }
-
-    /**
-     * Returns all currently connected `MCPClient` instances.
-     *
-     * Use this to enumerate available tools across all connected servers,
-     * or to route a tool call to the right client by `serverId`.
-     */
-    getClients(): MCPClient[] {
-        return this.clients;
-    }
-
-    /**
-     * Gracefully disconnects all active MCP clients and clears the internal client list.
-     *
-     * For Streamable HTTP sessions, each client sends an HTTP DELETE to its MCP
-     * endpoint per the spec before closing locally. All disconnects run in
-     * parallel so shutdown is not serialised across many sessions.
-     *
-     * Call this during server shutdown or when a user logs out to free up
-     * underlying transport resources (SSE streams, HTTP connections, etc.).
-     */
-    async disconnect(): Promise<void> {
-        await Promise.all(this.clients.map((client) => client.disconnect()));
-        this.clients = [];
+        this.options.onSessionFailed?.(session.sessionId, lastError);
+        console.error(
+            `[MultiSessionClient] Failed to connect to session ${session.sessionId} ` +
+            `after ${maxRetries + 1} attempts:`,
+            lastError,
+        );
     }
 }
-

package/src/server/mcp/oauth-client.ts

@@ -257,7 +257,6 @@ export class MCPClient {
           const hasSessionHeader = init?.headers && new Headers(init.headers as HeadersInit).has('mcp-session-id');
 
           if (response.status === 404 && hasSessionHeader) {
-            this.client = null;
             throw new Error("MCP_SESSION_EXPIRED: Downstream session was not found on the server.");
           }
 
@@ -791,7 +790,9 @@ export class MCPClient {
         params: {},
       };
 
-      const result = await this.client.request(request, ListToolsResultSchema);
+      const result = await this.withRetry(() =>
+        this.client!.request(request, ListToolsResultSchema)
+      );
 
       if (this.serverId) {
         this._onConnectionEvent.fire({
@@ -837,7 +838,9 @@ export class MCPClient {
     };
 
     try {
-      const result = await this.client.request(request, CallToolResultSchema);
+      const result = await this.withRetry(() =>
+        this.client!.request(request, CallToolResultSchema)
+      );
 
       this._onObservabilityEvent.fire({
         type: 'mcp:client:tool_call',
@@ -897,7 +900,9 @@ export class MCPClient {
         params: {},
       };
 
-      const result = await this.client.request(request, ListPromptsResultSchema);
+      const result = await this.withRetry(() =>
+        this.client!.request(request, ListPromptsResultSchema)
+      );
 
       this.emitStateChange('READY');
       this.emitProgress(`Discovered ${result.prompts.length} prompts`);
@@ -931,7 +936,9 @@ export class MCPClient {
       },
     };
 
-    return await this.client.request(request, GetPromptResultSchema);
+    return await this.withRetry(() =>
+      this.client!.request(request, GetPromptResultSchema)
+    );
   }
 
   /**
@@ -952,7 +959,9 @@ export class MCPClient {
         params: {},
       };
 
-      const result = await this.client.request(request, ListResourcesResultSchema);
+      const result = await this.withRetry(() =>
+        this.client!.request(request, ListResourcesResultSchema)
+      );
 
       this.emitStateChange('READY');
       this.emitProgress(`Discovered ${result.resources.length} resources`);
@@ -984,7 +993,34 @@ export class MCPClient {
       },
     };
 
-    return await this.client.request(request, ReadResourceResultSchema);
+    return await this.withRetry(() =>
+      this.client!.request(request, ReadResourceResultSchema)
+    );
+  }
+
+  /**
+   * Wraps an MCP request with automatic transport-session recovery.
+   *
+   * When the downstream MCP server rejects the request with a 404 indicating
+   * the transport session has expired, this method tears down the stale SDK
+   * client and transport, calls {@link reconnect} to negotiate a fresh session,
+   * and retries the request once.
+   *
+   * Non-transient errors (network failures, auth errors, etc.) propagate as-is.
+   */
+  private async withRetry<T>(fn: () => Promise<T>): Promise<T> {
+    try {
+      return await fn();
+    } catch (error) {
+      if (!(error instanceof Error && error.message.includes('MCP_SESSION_EXPIRED'))) throw error;
+      if (this.client) {
+        try { await this.client.close(); } catch {}
+        this.transport = null;
+        this.client = null;
+      }
+      await this.reconnect();
+      return await fn();
+    }
   }
 
   /**
@@ -1000,6 +1036,12 @@ export class MCPClient {
       throw new Error('OAuth provider not initialized');
     }
 
+    // Close the client initialize() may have created — we need a fresh
+    // client that will negotiate a new transport session.
+    if (this.client) {
+      try { await this.client.close(); } catch {}
+    }
+
     this.client = new Client(
       {
         name: MCP_CLIENT_NAME,

package/src/server/storage/index.ts

@@ -141,13 +141,13 @@ async function createStorage(): Promise<SessionStore> {
 
     if (type === 'supabase') {
         const url = process.env.SUPABASE_URL;
-        const key = process.env.SUPABASE_SERVICE_ROLE_KEY || process.env.SUPABASE_ANON_KEY;
+        const key = process.env.SUPABASE_SECRET_KEY || process.env.SUPABASE_ANON_KEY;
         
         if (!url || !key) {
-            console.warn('[mcp-ts][Storage] Explicit selection "supabase" requires SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY.');
+            console.warn('[mcp-ts][Storage] Explicit selection "supabase" requires SUPABASE_URL and SUPABASE_SECRET_KEY.');
         } else {
-            if (!process.env.SUPABASE_SERVICE_ROLE_KEY) {
-                console.warn('[mcp-ts][Storage] ⚠️  Warning: Using "SUPABASE_ANON_KEY" for server-side storage. You may encounter RLS policy violations. "SUPABASE_SERVICE_ROLE_KEY" is recommended.');
+            if (!process.env.SUPABASE_SECRET_KEY) {
+                console.warn('[mcp-ts][Storage] ⚠️  Warning: Using "SUPABASE_ANON_KEY" for server-side storage. You may encounter RLS policy violations. "SUPABASE_SECRET_KEY" is recommended.');
             }
             try {
                 const { createClient } = await import('@supabase/supabase-js');
@@ -210,14 +210,14 @@ async function createStorage(): Promise<SessionStore> {
         return await initializeStorage(new SqliteStorage({ path: process.env.MCP_TS_STORAGE_SQLITE_PATH }));
     }
 
-    if (process.env.SUPABASE_URL && (process.env.SUPABASE_SERVICE_ROLE_KEY || process.env.SUPABASE_ANON_KEY)) {
+    if (process.env.SUPABASE_URL && (process.env.SUPABASE_SECRET_KEY || process.env.SUPABASE_ANON_KEY)) {
         try {
             const { createClient } = await import('@supabase/supabase-js');
             const url = process.env.SUPABASE_URL;
-            const key = process.env.SUPABASE_SERVICE_ROLE_KEY || process.env.SUPABASE_ANON_KEY!;
+            const key = process.env.SUPABASE_SECRET_KEY || process.env.SUPABASE_ANON_KEY!;
             
-            if (!process.env.SUPABASE_SERVICE_ROLE_KEY) {
-                console.warn('[mcp-ts][Storage] ⚠️ Warning: Using "SUPABASE_ANON_KEY" for server-side storage. You may encounter RLS policy violations. "SUPABASE_SERVICE_ROLE_KEY" is recommended.');
+            if (!process.env.SUPABASE_SECRET_KEY) {
+                console.warn('[mcp-ts][Storage] ⚠️ Warning: Using "SUPABASE_ANON_KEY" for server-side storage. You may encounter RLS policy violations. "SUPABASE_SECRET_KEY" is recommended.');
             }
 
             const client = createClient(url, key);