more comments

Author: fmenezes

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, LoggerBase, McpLogger } from "./logger.js";
-import { TimeoutManager } from "./timeoutManager.js";
+import { ManagedTimeout, setManagedTimeout } from "./managedTimeout.js";
 
 export class SessionStore {
     private sessions: {
         [sessionId: string]: {
             logger: LoggerBase;
             transport: StreamableHTTPServerTransport;
-            abortTimeout: TimeoutManager;
-            notificationTimeout: TimeoutManager;
+            abortTimeout: ManagedTimeout;
+            notificationTimeout: ManagedTimeout;
         };
     } = {};
 
@@ -39,9 +39,9 @@ export class SessionStore {
             return;
         }
 
-        session.abortTimeout.reset();
+        session.abortTimeout.restart();
 
-        session.notificationTimeout.reset();
+        session.notificationTimeout.restart();
     }
 
     private sendNotification(sessionId: string): void {
@@ -66,7 +66,7 @@ export class SessionStore {
         if (session) {
             throw new Error(`Session ${sessionId} already exists`);
         }
-        const abortTimeout = new TimeoutManager(async () => {
+        const abortTimeout = setManagedTimeout(async () => {
             if (this.sessions[sessionId]) {
                 this.sessions[sessionId].logger.info(
                     LogId.streamableHttpTransportSessionCloseNotification,
@@ -77,7 +77,7 @@ export class SessionStore {
                 await this.closeSession(sessionId);
             }
         }, this.idleTimeoutMS);
-        const notificationTimeout = new TimeoutManager(
+        const notificationTimeout = setManagedTimeout(
             () => this.sendNotification(sessionId),
             this.notificationTimeoutMS
         );
@@ -89,8 +89,8 @@ export class SessionStore {
         if (!session) {
             throw new Error(`Session ${sessionId} not found`);
         }
-        session.abortTimeout.clear();
-        session.notificationTimeout.clear();
+        session.abortTimeout.cancel();
+        session.notificationTimeout.cancel();
         if (closeTransport) {
             try {
                 await session.transport.close();

src/common/timeoutManager.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();
         });