chore: POC work for Sticky sessions

src/common/sessionStore.ts

@@ -2,20 +2,29 @@ import type { LoggerBase } from "./logging/index.js";
 import { LogId } from "./logging/index.js";
 import type { ManagedTimeout } from "./managedTimeout.js";
 import { setManagedTimeout } from "./managedTimeout.js";
+import { type Server } from "../server.js";
+import { type StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 
-/**
- * Minimal interface for a transport that can be stored in a SessionStore.
- * The transport must have a close method for cleanup.
- */
-export type CloseableTransport = {
-    close(): Promise<void>;
-};
+export interface ISessionStore {
+    getSession(sessionId: string): Promise<StreamableHTTPServerTransport | undefined>;
+    setSession(
+        sessionId: string,
+        sessionState: {
+            transport: StreamableHTTPServerTransport;
+            server: Server;
+            logger: LoggerBase;
+        }
+    ): Promise<void>;
+    closeSession(sessionId: string, closeTransport: boolean): Promise<void>;
+    closeAllSessions(): Promise<void>;
+}
 
-export class SessionStore<T extends CloseableTransport = CloseableTransport> {
+export class SessionStore implements ISessionStore {
     private sessions: {
         [sessionId: string]: {
             logger: LoggerBase;
-            transport: T;
+            transport: StreamableHTTPServerTransport;
+            server: Server;
             abortTimeout: ManagedTimeout;
             notificationTimeout: ManagedTimeout;
         };
@@ -37,7 +46,8 @@ export class SessionStore<T extends CloseableTransport = CloseableTransport> {
         }
     }
 
-    getSession(sessionId: string): T | undefined {
+    // eslint-disable-next-line @typescript-eslint/require-await
+    async getSession(sessionId: string): Promise<StreamableHTTPServerTransport | undefined> {
         this.resetTimeout(sessionId);
         return this.sessions[sessionId]?.transport;
     }
@@ -70,7 +80,15 @@ export class SessionStore<T extends CloseableTransport = CloseableTransport> {
         });
     }
 
-    setSession(sessionId: string, transport: T, logger: LoggerBase): void {
+    // eslint-disable-next-line @typescript-eslint/require-await
+    async setSession(
+        sessionId: string,
+        sessionState: {
+            transport: StreamableHTTPServerTransport;
+            server: Server;
+            logger: LoggerBase;
+        }
+    ): Promise<void> {
         const session = this.sessions[sessionId];
         if (session) {
             throw new Error(`Session ${sessionId} already exists`);
@@ -91,10 +109,9 @@ export class SessionStore<T extends CloseableTransport = CloseableTransport> {
             this.notificationTimeoutMS
         );
         this.sessions[sessionId] = {
-            transport,
+            ...sessionState,
             abortTimeout,
             notificationTimeout,
-            logger,
         };
     }
 

src/index.ts

@@ -42,7 +42,7 @@ import { parseUserConfig } from "./common/config/parseUserConfig.js";
 import { type UserConfig } from "./common/config/userConfig.js";
 import { packageInfo } from "./common/packageInfo.js";
 import { StdioRunner } from "./transports/stdio.js";
-import { StreamableHttpRunner } from "./transports/streamableHttp.js";
+import { StreamableHttpRunner } from "./transports/streamableHttp/streamableHttpRunner.js";
 import { systemCA } from "@mongodb-js/devtools-proxy-support";
 import { Keychain } from "./common/keychain.js";
 import { DryRunModeRunner } from "./transports/dryModeRunner.js";

src/lib.ts

@@ -32,7 +32,7 @@ export {
     type LogLevel,
     CompositeLogger,
 } from "./common/logging/index.js";
-export { StreamableHttpRunner } from "./transports/streamableHttp.js";
+export { StreamableHttpRunner } from "./transports/streamableHttp/streamableHttpRunner.js";
 export { StdioRunner } from "./transports/stdio.js";
 export {
     TransportRunnerBase,
@@ -62,7 +62,7 @@ export { Keychain, registerGlobalSecretToRedact } from "./common/keychain.js";
 export type { Secret } from "./common/keychain.js";
 export { Elicitation } from "./elicitation.js";
 export { applyConfigOverrides } from "./common/config/configOverrides.js";
-export { SessionStore, type CloseableTransport } from "./common/sessionStore.js";
+export { SessionStore, type ISessionStore } from "./common/sessionStore.js";
 export { ApiClient, type ApiClientOptions } from "./common/atlas/apiClient.js";
 export type { AuthProvider } from "./common/atlas/auth/authProvider.js";
 export { type UIRegistryOptions } from "./ui/registry/registry.js";