remove clients from session id map when disconnecting

Author: JReinhold

src/mcp-handler.ts

@@ -10,7 +10,7 @@ import type { Options, CoreConfig } from "storybook/internal/types";
 import type { IncomingMessage, ServerResponse } from "node:http";
 import {
   collectTelemetry,
-  setClientForSession,
+  mcpSessionIdToClientMap,
   setDisableTelemetry,
 } from "./telemetry";
 
@@ -26,7 +26,7 @@ async function createMcpServer(options: Options, client: string) {
         {},
       );
       setDisableTelemetry(disableTelemetry);
-      setClientForSession(sessionId, client);
+      mcpSessionIdToClientMap[sessionId] = client;
 
       await collectTelemetry({
         event: "session:initialized",
@@ -36,8 +36,13 @@ async function createMcpServer(options: Options, client: string) {
   });
 
   transport.onclose = () => {
-    if (transport.sessionId) {
-      delete transports[transport.sessionId];
+    if (!transport.sessionId) {
+      return;
+    }
+
+    delete transports[transport.sessionId];
+    if (mcpSessionIdToClientMap[transport.sessionId]) {
+      delete mcpSessionIdToClientMap[transport.sessionId];
     }
   };
 

src/telemetry.ts

@@ -26,11 +26,7 @@ export async function collectTelemetry({
   }
 }
 
-const mcpSessionIdToClientMap: Record<string, string> = {};
-
-export const setClientForSession = (sessionId: string, client: string) => {
-  mcpSessionIdToClientMap[sessionId] = client;
-};
+export const mcpSessionIdToClientMap: Record<string, string> = {};
 
 let disableTelemetry = false;
 export const setDisableTelemetry = (value = false) => {