Fixed the issue where transport was overwritten when multiple clients connected to the same server

Author: JasmineLCY

src/cli.ts

@@ -91,7 +91,8 @@ async function runServer(port: number | null) {
             console.log('Received message');
 
             const sessionId = req.query.sessionId as string;
-            const transport = servers.map(s => s.transport as SSEServerTransport).find(t => t.sessionId === sessionId);
+ 
+            const transport = servers.reduce((pre, val) => pre.concat(Array.from(val.transportMap.values()) as SSEServerTransport[]), [] as SSEServerTransport[]).find(t => t.sessionId === sessionId);
             if (!transport) {
                 res.status(404).send('Session not found');
                 return;

src/client/index.ts

@@ -103,7 +103,7 @@ export class Client<
      * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization).
      */
     public registerCapabilities(capabilities: ClientCapabilities): void {
-        if (this.transport) {
+        if (this.transportMap.size > 0) {
             throw new Error('Cannot register capabilities after connecting to transport');
         }
 

src/inMemory.ts

@@ -18,6 +18,8 @@ export class InMemoryTransport implements Transport {
     onerror?: (error: Error) => void;
     onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void;
     sessionId?: string;
+    onsessionclosed?: ((sessionId: string) => void | Promise<void>) | undefined;
+    onsessioninitialized?: ((sessionId: string) => void | Promise<void>) | undefined;
 
     /**
      * Creates a pair of linked in-memory transports that can communicate with each other. One should be passed to a Client and one to a Server.
@@ -31,6 +33,7 @@ export class InMemoryTransport implements Transport {
     }
 
     async start(): Promise<void> {
+        this.onsessioninitialized?.(this.sessionId!);
         // Process any messages that were queued before start was called
         while (this._messageQueue.length > 0) {
             const queuedMessage = this._messageQueue.shift()!;
@@ -39,6 +42,7 @@ export class InMemoryTransport implements Transport {
     }
 
     async close(): Promise<void> {
+        this.onsessionclosed?.(this.sessionId!);
         const other = this._otherTransport;
         this._otherTransport = undefined;
         await other?.close();

src/server/index.ts

@@ -131,7 +131,7 @@ export class Server<
      * The new capabilities will be merged with any existing capabilities previously given (e.g., at initialization).
      */
     public registerCapabilities(capabilities: ServerCapabilities): void {
-        if (this.transport) {
+        if (this.transportMap.size > 0) {
             throw new Error('Cannot register capabilities after connecting to transport');
         }
         this._capabilities = mergeCapabilities(this._capabilities, capabilities);
@@ -277,8 +277,8 @@ export class Server<
         return this._capabilities;
     }
 
-    async ping() {
-        return this.request({ method: 'ping' }, EmptyResultSchema);
+    async ping(sessionId?: string) {
+        return this.request({ method: 'ping' }, EmptyResultSchema, { sessionId });
     }
 
     async createMessage(params: CreateMessageRequest['params'], options?: RequestOptions) {
@@ -327,7 +327,7 @@ export class Server<
     async sendLoggingMessage(params: LoggingMessageNotification['params'], sessionId?: string) {
         if (this._capabilities.logging) {
             if (!this.isMessageIgnored(params.level, sessionId)) {
-                return this.notification({ method: 'notifications/message', params });
+                return this.notification({ method: 'notifications/message', params }, { sessionId});
             }
         }
     }

src/server/mcp.ts

@@ -902,7 +902,7 @@ export class McpServer {
      * @returns True if the server is connected
      */
     isConnected() {
-        return this.server.transport !== undefined;
+        return this.server.transportMap.size > 0;
     }
 
     /**