fix: close

Author: fmenezes

src/common/sessionStore.ts

@@ -2,25 +2,25 @@ import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/
 import logger, { LogId } from "./logger.js";
 
 export class SessionStore {
-    private sessions: Map<string, StreamableHTTPServerTransport> = new Map();
+    private sessions: { [sessionId: string]: StreamableHTTPServerTransport | undefined } = {};
 
     getSession(sessionId: string): StreamableHTTPServerTransport | undefined {
-        return this.sessions.get(sessionId);
+        return this.sessions[sessionId];
     }
 
     setSession(sessionId: string, transport: StreamableHTTPServerTransport): void {
-        if (this.sessions.has(sessionId)) {
+        if (this.sessions[sessionId]) {
             throw new Error(`Session ${sessionId} already exists`);
         }
-        this.sessions.set(sessionId, transport);
+        this.sessions[sessionId] = transport;
     }
 
     async closeSession(sessionId: string, closeTransport: boolean = true): Promise<void> {
-        if (!this.sessions.has(sessionId)) {
+        if (!this.sessions[sessionId]) {
             throw new Error(`Session ${sessionId} not found`);
         }
         if (closeTransport) {
-            const transport = this.sessions.get(sessionId);
+            const transport = this.sessions[sessionId];
             if (!transport) {
                 throw new Error(`Session ${sessionId} not found`);
             }
@@ -34,11 +34,15 @@ export class SessionStore {
                 );
             }
         }
-        this.sessions.delete(sessionId);
+        delete this.sessions[sessionId];
     }
 
     async closeAllSessions(): Promise<void> {
-        await Promise.all(Array.from(this.sessions.values()).map((transport) => transport.close()));
-        this.sessions.clear();
+        await Promise.all(
+            Object.values(this.sessions)
+                .filter((transport) => transport !== undefined)
+                .map((transport) => transport.close())
+        );
+        this.sessions = {};
     }
 }

src/index.ts

@@ -14,6 +14,7 @@ async function main() {
         transportRunner
             .close()
             .then(() => {
+                logger.info(LogId.serverClosed, "server", `Server closed`);
                 process.exit(0);
             })
             .catch((error: unknown) => {
@@ -22,10 +23,10 @@ async function main() {
             });
     };
 
-    process.once("SIGINT", shutdown);
-    process.once("SIGABRT", shutdown);
-    process.once("SIGTERM", shutdown);
-    process.once("SIGQUIT", shutdown);
+    process.on("SIGINT", shutdown);
+    process.on("SIGABRT", shutdown);
+    process.on("SIGTERM", shutdown);
+    process.on("SIGQUIT", shutdown);
 
     try {
         await transportRunner.start();

src/transports/streamableHttp.ts

@@ -155,5 +155,6 @@ export class StreamableHttpRunner extends TransportRunnerBase {
                 resolve();
             });
         });
+        await this.sessionStore.closeAllSessions();
     }
 }