chore: address comments from #361

src/common/logger.ts

@@ -44,8 +44,9 @@ export const LogId = {
     streamableHttpTransportStarted: mongoLogId(1_006_001),
     streamableHttpTransportSessionCloseFailure: mongoLogId(1_006_002),
     streamableHttpTransportSessionCloseNotification: mongoLogId(1_006_003),
-    streamableHttpTransportRequestFailure: mongoLogId(1_006_004),
-    streamableHttpTransportCloseFailure: mongoLogId(1_006_005),
+    streamableHttpTransportSessionCloseNotificationFailure: mongoLogId(1_006_004),
+    streamableHttpTransportRequestFailure: mongoLogId(1_006_005),
+    streamableHttpTransportCloseFailure: mongoLogId(1_006_006),
 } as const;
 
 export abstract class LoggerBase {

src/common/managedTimeout.ts

@@ -0,0 +1,27 @@
+export interface ManagedTimeout {
+    cancel: () => void;
+    restart: () => void;
+}
+
+export function setManagedTimeout(callback: () => Promise<void> | void, timeoutMS: number): ManagedTimeout {
+    let timeoutId: NodeJS.Timeout | undefined = setTimeout(() => {
+        void callback();
+    }, timeoutMS);
+
+    function cancel() {
+        clearTimeout(timeoutId);
+        timeoutId = undefined;
+    }
+
+    function restart() {
+        cancel();
+        timeoutId = setTimeout(() => {
+            void callback();
+        }, timeoutMS);
+    }
+
+    return {
+        cancel,
+        restart,
+    };
+}

src/common/sessionStore.ts

@@ -1,15 +1,15 @@
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
-import logger, { LogId, McpLogger } from "./logger.js";
-import { TimeoutManager } from "./timeoutManager.js";
+import logger, { LogId, LoggerBase, McpLogger } from "./logger.js";
+import { ManagedTimeout, setManagedTimeout } from "./managedTimeout.js";
 
 export class SessionStore {
     private sessions: {
         [sessionId: string]: {
-            mcpServer: McpServer;
+            logger: LoggerBase;
             transport: StreamableHTTPServerTransport;
-            abortTimeout: TimeoutManager;
-            notificationTimeout: TimeoutManager;
+            abortTimeout: ManagedTimeout;
+            notificationTimeout: ManagedTimeout;
         };
     } = {};
 
@@ -39,54 +39,61 @@ export class SessionStore {
             return;
         }
 
-        session.abortTimeout.reset();
+        session.abortTimeout.restart();
 
-        session.notificationTimeout.reset();
+        session.notificationTimeout.restart();
     }
 
     private sendNotification(sessionId: string): void {
         const session = this.sessions[sessionId];
         if (!session) {
+            logger.warning(
+                LogId.streamableHttpTransportSessionCloseNotificationFailure,
+                "sessionStore",
+                `session ${sessionId} not found, no notification delivered`
+            );
             return;
         }
-        const logger = new McpLogger(session.mcpServer);
-        logger.info(
+        session.logger.info(
             LogId.streamableHttpTransportSessionCloseNotification,
             "sessionStore",
             "Session is about to be closed due to inactivity"
         );
     }
 
     setSession(sessionId: string, transport: StreamableHTTPServerTransport, mcpServer: McpServer): void {
-        if (this.sessions[sessionId]) {
+        const session = this.sessions[sessionId];
+        if (session) {
             throw new Error(`Session ${sessionId} already exists`);
         }
-        const abortTimeout = new TimeoutManager(async () => {
-            const logger = new McpLogger(mcpServer);
-            logger.info(
-                LogId.streamableHttpTransportSessionCloseNotification,
-                "sessionStore",
-                "Session closed due to inactivity"
-            );
+        const abortTimeout = setManagedTimeout(async () => {
+            if (this.sessions[sessionId]) {
+                this.sessions[sessionId].logger.info(
+                    LogId.streamableHttpTransportSessionCloseNotification,
+                    "sessionStore",
+                    "Session closed due to inactivity"
+                );
 
-            await this.closeSession(sessionId);
+                await this.closeSession(sessionId);
+            }
         }, this.idleTimeoutMS);
-        const notificationTimeout = new TimeoutManager(
+        const notificationTimeout = setManagedTimeout(
             () => this.sendNotification(sessionId),
             this.notificationTimeoutMS
         );
-        this.sessions[sessionId] = { mcpServer, transport, abortTimeout, notificationTimeout };
+        this.sessions[sessionId] = { logger: new McpLogger(mcpServer), transport, abortTimeout, notificationTimeout };
     }
 
     async closeSession(sessionId: string, closeTransport: boolean = true): Promise<void> {
-        if (!this.sessions[sessionId]) {
+        const session = this.sessions[sessionId];
+        if (!session) {
             throw new Error(`Session ${sessionId} not found`);
         }
-        this.sessions[sessionId].abortTimeout.clear();
-        this.sessions[sessionId].notificationTimeout.clear();
+        session.abortTimeout.cancel();
+        session.notificationTimeout.cancel();
         if (closeTransport) {
             try {
-                await this.sessions[sessionId].transport.close();
+                await session.transport.close();
             } catch (error) {
                 logger.error(
                     LogId.streamableHttpTransportSessionCloseFailure,

tests/unit/common/timeoutManager.test.ts → tests/unit/common/managedTimeout.test.ts

@@ -1,7 +1,7 @@
 import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
-import { TimeoutManager } from "../../../src/common/timeoutManager.js";
+import { setManagedTimeout } from "../../../src/common/managedTimeout.js";
 
-describe("TimeoutManager", () => {
+describe("setManagedTimeout", () => {
     beforeAll(() => {
         vi.useFakeTimers();
     });
@@ -13,7 +13,7 @@ describe("TimeoutManager", () => {
     it("calls the timeout callback", () => {
         const callback = vi.fn();
 
-        new TimeoutManager(callback, 1000);
+        setManagedTimeout(callback, 1000);
 
         vi.advanceTimersByTime(1000);
         expect(callback).toHaveBeenCalled();
@@ -22,10 +22,10 @@ describe("TimeoutManager", () => {
     it("does not call the timeout callback if the timeout is cleared", () => {
         const callback = vi.fn();
 
-        const timeoutManager = new TimeoutManager(callback, 1000);
+        const timeout = setManagedTimeout(callback, 1000);
 
         vi.advanceTimersByTime(500);
-        timeoutManager.clear();
+        timeout.cancel();
         vi.advanceTimersByTime(500);
 
         expect(callback).not.toHaveBeenCalled();
@@ -34,44 +34,32 @@ describe("TimeoutManager", () => {
     it("does not call the timeout callback if the timeout is reset", () => {
         const callback = vi.fn();
 
-        const timeoutManager = new TimeoutManager(callback, 1000);
+        const timeout = setManagedTimeout(callback, 1000);
 
         vi.advanceTimersByTime(500);
-        timeoutManager.reset();
+        timeout.restart();
         vi.advanceTimersByTime(500);
         expect(callback).not.toHaveBeenCalled();
     });
 
-    it("calls the onerror callback", () => {
-        const onerrorCallback = vi.fn();
-
-        const tm = new TimeoutManager(() => {
-            throw new Error("test");
-        }, 1000);
-        tm.onerror = onerrorCallback;
-
-        vi.advanceTimersByTime(1000);
-        expect(onerrorCallback).toHaveBeenCalled();
-    });
-
     describe("if timeout is reset", () => {
         it("does not call the timeout callback within the timeout period", () => {
             const callback = vi.fn();
 
-            const timeoutManager = new TimeoutManager(callback, 1000);
+            const timeout = setManagedTimeout(callback, 1000);
 
             vi.advanceTimersByTime(500);
-            timeoutManager.reset();
+            timeout.restart();
             vi.advanceTimersByTime(500);
             expect(callback).not.toHaveBeenCalled();
         });
         it("calls the timeout callback after the timeout period", () => {
             const callback = vi.fn();
 
-            const timeoutManager = new TimeoutManager(callback, 1000);
+            const timeout = setManagedTimeout(callback, 1000);
 
             vi.advanceTimersByTime(500);
-            timeoutManager.reset();
+            timeout.restart();
             vi.advanceTimersByTime(1000);
             expect(callback).toHaveBeenCalled();
         });