feat: add stateless transport manager with cache (#36)

Author: chrisleekr

.vscode/settings.json

@@ -45,6 +45,7 @@
     "CERTDIR",
     "chrislee",
     "chrisleekr",
+    "dbsize",
     "dind",
     "dpop",
     "interruptible",

README.md

@@ -7,6 +7,7 @@ A playground for Model Context Protocol (MCP) server built with TypeScript and S
 - MCP Server implementation: HTTP-Based Streamable transport using `@modelcontextprotocol/sdk` with HTTP transport, session management, and tool execution.
 - OAuth authentication/3rd party authorization: Implements an OAuth server for MCP clients to process 3rd party authorization servers like Auth0, providing Dynamic Application Registration for MCP server.
 - Storage: Provide storage for MCP server to store data like OAuth sessions, tokens, etc.
+- Session Management: MCP server can manage multiple sessions stateless.
 - Tools: `echo`, `system-time`, `streaming`, `project` for demonstration.
 - Prompts: `echo`
 

src/core/server/http/handlers.ts

@@ -1,4 +1,3 @@
-import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import express from 'express';
 
 import { TransportManager } from '../transport';
@@ -9,7 +8,6 @@ import { setupPingHandler } from './handlers/ping';
 
 export function setupRequestHandlers(
   app: express.Application,
-  server: Server,
   transportManager: TransportManager
 ): void {
   app.get('/', (_req, res) => {
@@ -18,8 +16,8 @@ export function setupRequestHandlers(
   });
 
   setupPingHandler(app);
-  setupMCPPostHandler(app, server, transportManager);
-  setupMCPDeleteHandler(app, server, transportManager);
+  setupMCPPostHandler(app, transportManager);
+  setupMCPDeleteHandler(app, transportManager);
 
   loggingContext.log('info', 'Request handlers setup complete');
 }

src/core/server/http/handlers/mcpDelete.ts

@@ -1,4 +1,3 @@
-import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import express from 'express';
 
 import type { TransportManager } from '@/core/server/transport';
@@ -7,7 +6,6 @@ import { loggingContext } from '../context';
 
 export function setupMCPDeleteHandler(
   app: express.Application,
-  _server: Server,
   transportManager: TransportManager
 ): void {
   // Handle MCP DELETE requests (session termination)
@@ -21,31 +19,26 @@ export function setupMCPDeleteHandler(
         if (
           sessionId === undefined ||
           sessionId.trim() === '' ||
-          !transportManager.hasTransport(sessionId)
+          !(await transportManager.hasTransport(sessionId))
         ) {
-          loggingContext.log('error', 'Session not found', {
-            data: { sessionId },
-          });
+          loggingContext.log('error', 'Session not found');
           res.status(200).json({ error: 'Session not found' }); // Return 200 to gracefully handle the request
           return;
         }
 
-        const transport = transportManager.getTransport(sessionId);
+        loggingContext.log('debug', 'Found session, getting transport');
+        const transport = await transportManager.getTransport(sessionId);
         if (!transport) {
-          loggingContext.log('error', 'Transport not found', {
-            data: { sessionId },
-          });
+          loggingContext.log('error', 'Transport not found');
           res.status(200).json({ error: 'Transport not found' }); // Return 200 to gracefully handle the request
           return;
         }
         await transport.handleRequest(req, res);
 
         // Clean up the transport
-        transportManager.deleteTransport(sessionId);
+        await transportManager.deleteTransport(sessionId);
 
-        loggingContext.log('info', 'Session terminated', {
-          data: { sessionId },
-        });
+        loggingContext.log('info', 'Session terminated');
       } catch (error: unknown) {
         loggingContext.log('error', 'Error handling DELETE request', {
           data: {

src/core/server/http/handlers/mcpPost.ts

@@ -1,4 +1,5 @@
-import { Server } from '@modelcontextprotocol/sdk/server/index.js';
+import { randomUUID } from 'node:crypto';
+
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
 import express from 'express';
@@ -18,7 +19,6 @@ function getSessionId(req: express.Request): string | undefined {
 
 export function setupMCPPostHandler(
   app: express.Application,
-  server: Server,
   transportManager: TransportManager
 ): void {
   // Handle MCP POST requests (Streamable HTTP)
@@ -29,47 +29,45 @@ export function setupMCPPostHandler(
         const sessionId = getSessionId(req);
 
         loggingContext.log('debug', 'POST /mcp request body', {
-          data: { requestBody, sessionId },
+          data: { requestBody },
         });
 
         let transport: StreamableHTTPServerTransport;
 
         if (
           sessionId !== undefined &&
           sessionId.trim() !== '' &&
-          transportManager.hasTransport(sessionId)
+          (await transportManager.hasTransport(sessionId))
         ) {
+          loggingContext.log('debug', 'Found session, getting transport');
           // Reuse existing transport
-          const existingTransport = transportManager.getTransport(sessionId);
+          const existingTransport =
+            await transportManager.getTransport(sessionId);
           if (!existingTransport) {
             loggingContext.log(
               'error',
-              'Transport not found despite has() check',
-              {
-                data: { sessionId },
-              }
+              'Transport not found despite has() check'
             );
             throw new Error('Transport not found despite has() check');
           }
+          loggingContext.log(
+            'debug',
+            'Transport found, using existing transport'
+          );
           transport = existingTransport;
         } else if (
           sessionId === undefined &&
           isInitializeRequest(requestBody)
         ) {
+          loggingContext.log('debug', 'No session found, creating new session');
           // New initialization request
-          transport = transportManager.createTransport();
+          const newSessionId = randomUUID();
 
-          // Connect the transport to the server
-          await server.connect(
-            transport as StreamableHTTPServerTransport & { sessionId: string }
-          );
+          transport = await transportManager.createTransport(newSessionId);
         } else {
           loggingContext.log(
             'error',
-            'Invalid request: missing session ID or not an initialization request',
-            {
-              data: { sessionId },
-            }
+            'Invalid request: missing session ID or not an initialization request'
           );
           res.status(400).json({
             error:
@@ -79,6 +77,7 @@ export function setupMCPPostHandler(
         }
 
         await transport.handleRequest(req, res, requestBody);
+        loggingContext.log('debug', 'Request handled');
         return;
       } catch (error) {
         loggingContext.log('error', 'Error handling HTTP request', {

src/core/server/http/server.ts

@@ -8,13 +8,13 @@ import { setupMiddleware } from './middleware';
 
 export function setupHttpServer(server: Server): express.Application {
   const app = express();
-  const transportManager = new TransportManager();
+  const transportManager = new TransportManager(server);
 
   // Setup middleware
   setupMiddleware(app);
 
   // Setup request handlers
-  setupRequestHandlers(app, server, transportManager);
+  setupRequestHandlers(app, transportManager);
 
   // Setup auth handlers
   setupAuthHandlers(app);

src/core/server/transport/manager.ts

@@ -1,38 +1,91 @@
-import { randomUUID } from 'node:crypto';
-
+import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 
+import { config } from '@/config/manager';
+import { createStorage } from '@/core/storage/storageFactory';
+import { Storage } from '@/core/storage/types';
+
 import { loggingContext } from '../http/context';
 
 export class TransportManager {
+  private server: Server;
+  private storage: Storage;
   private transports: Map<string, StreamableHTTPServerTransport> = new Map();
 
-  public getTransport(
+  private readonly CACHE_KEY_PREFIX = 'mcp-session';
+
+  constructor(server: Server) {
+    this.server = server;
+
+    try {
+      this.storage = createStorage(config.storage);
+    } catch (error: unknown) {
+      loggingContext.log('error', 'Failed to create storage', {
+        error: error instanceof Error ? error.message : 'Unknown error',
+        stack: error instanceof Error ? error.stack : undefined,
+      });
+      throw error;
+    }
+  }
+
+  public async getTransport(
     sessionId: string
-  ): StreamableHTTPServerTransport | undefined {
+  ): Promise<StreamableHTTPServerTransport | undefined> {
     loggingContext.log('debug', 'Getting transport', {
       data: { sessionId },
     });
-    return this.transports.get(sessionId);
+
+    // If storage contains sessionId, then check transports with sessionId. If not, then create a new transport with sessionId.
+    const session = await this.storage.get(
+      `${this.CACHE_KEY_PREFIX}:${sessionId}`
+    );
+    if (session !== null && session.trim() !== '') {
+      if (this.transports.has(sessionId)) {
+        loggingContext.log('debug', 'Transport found in transports', {
+          data: { sessionId },
+        });
+        return this.transports.get(sessionId);
+      }
+
+      // It exists in storage, but not in transports. Create a new transport with sessionId.
+      loggingContext.log(
+        'debug',
+        'Transport not found in transports, creating new transport'
+      );
+      const newTransport = await this.createTransport(sessionId);
+      this.transports.set(sessionId, newTransport);
+
+      return newTransport;
+    }
+
+    // If session is not found, then return undefined.
+    return undefined;
   }
 
-  public hasTransport(sessionId: string): boolean {
+  public async hasTransport(sessionId: string): Promise<boolean> {
     loggingContext.log('debug', 'Checking if transport exists', {
       data: { sessionId },
     });
-    return this.transports.has(sessionId);
+    const session = await this.storage.get(
+      `${this.CACHE_KEY_PREFIX}:${sessionId}`
+    );
+    return session !== null && session.trim() !== '';
   }
 
-  public createTransport(): StreamableHTTPServerTransport {
-    const newSessionId = randomUUID();
+  public async createTransport(
+    sessionId: string
+  ): Promise<StreamableHTTPServerTransport> {
     const transport = new StreamableHTTPServerTransport({
       /**
        * Function that generates a session ID for the transport.
        * The session ID SHOULD be globally unique and cryptographically secure (e.g., a securely generated UUID, a JWT, or a cryptographic hash)
        *
        * Return undefined to disable session management.
        */
-      sessionIdGenerator: (): string => newSessionId,
+      // This is disabled to make stateless mode.
+      sessionIdGenerator: undefined,
+      // Below is for stateful mode.
+      // sessionIdGenerator: (): string => sessionId,
       /**
        * If true, the server will return JSON responses instead of starting an SSE stream.
        * This can be useful for simple request/response scenarios without streaming.
@@ -42,45 +95,57 @@ export class TransportManager {
     });
 
     // Manually set the session ID to ensure it's available
-    transport.sessionId = newSessionId;
+    transport.sessionId = sessionId;
     loggingContext.log('debug', 'Creating transport', {
-      data: { sessionId: newSessionId },
+      data: { sessionId },
     });
-    this.transports.set(newSessionId, transport);
 
-    loggingContext.log('debug', 'Transport created', {
-      data: { sessionId: newSessionId },
-    });
+    this.transports.set(sessionId, transport);
+    await this.storage.set(
+      `${this.CACHE_KEY_PREFIX}:${sessionId}`,
+      JSON.stringify({
+        createdAt: new Date().toISOString(),
+      })
+    );
+
+    loggingContext.log('debug', 'Transport created');
 
     // Set up cleanup handler
     transport.onclose = (): void => {
       const currentSessionId = transport.sessionId;
       if (currentSessionId !== undefined && currentSessionId.trim() !== '') {
         this.transports.delete(currentSessionId);
+        void this.storage.delete(
+          `${this.CACHE_KEY_PREFIX}:${currentSessionId}`
+        );
         loggingContext.log('debug', 'Transport closed and cleaned up', {
           data: {
-            sessionId: currentSessionId,
             transportCount: this.transports.size,
           },
         });
       }
     };
 
+    // Connect the transport to the server
+    loggingContext.log('debug', 'Connecting transport to server');
+    await this.server.connect(
+      transport as StreamableHTTPServerTransport & { sessionId: string }
+    );
+
     return transport;
   }
 
-  public deleteTransport(sessionId: string): void {
+  public async deleteTransport(sessionId: string): Promise<void> {
     this.transports.delete(sessionId);
-    loggingContext.log('info', 'Transport deleted', {
-      data: { sessionId },
-    });
+    await this.storage.delete(`${this.CACHE_KEY_PREFIX}:${sessionId}`);
+    loggingContext.log('info', 'Transport deleted');
   }
 
   public getTransportCount(): number {
     return this.transports.size;
   }
 
-  public getAllSessionIds(): string[] {
-    return Array.from(this.transports.keys());
+  public async getAllSessionIds(): Promise<string[]> {
+    return this.storage.keys(`${this.CACHE_KEY_PREFIX}:*`);
   }
 }

src/core/storage/memory.ts

@@ -35,4 +35,8 @@ export class MemoryStorage implements Storage {
   async keys(pattern: string): Promise<string[]> {
     return Array.from(this.store.keys()).filter(key => key.startsWith(pattern));
   }
+
+  async length(): Promise<number> {
+    return this.store.size;
+  }
 }

src/core/storage/types.ts

@@ -4,4 +4,5 @@ export interface Storage {
   delete(key: string): Promise<boolean>;
   keys(pattern: string): Promise<string[]>;
   close(): Promise<void>;
+  length(): Promise<number>;
 }

src/core/storage/valkey.ts

@@ -47,4 +47,8 @@ export class ValkeyStorage implements Storage {
     const keys = await this.client.keys(pattern);
     return keys.map(key => key.replace(pattern, ''));
   }
+
+  async length(): Promise<number> {
+    return this.client.dbsize();
+  }
 }