fix: address comments

Author: fmenezes

src/common/sessionStore.ts

@@ -1,47 +1,7 @@
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import logger, { LogId, McpLogger } from "./logger.js";
-
-class TimeoutManager {
-    private timeoutId?: NodeJS.Timeout;
-    public onerror?: (error: unknown) => void;
-
-    constructor(
-        private readonly callback: () => Promise<void> | void,
-        private readonly timeoutMS: number
-    ) {
-        if (timeoutMS <= 0) {
-            throw new Error("timeoutMS must be greater than 0");
-        }
-        this.reset();
-    }
-
-    clear() {
-        if (this.timeoutId) {
-            clearTimeout(this.timeoutId);
-            this.timeoutId = undefined;
-        }
-    }
-
-    private async runCallback() {
-        if (this.callback) {
-            try {
-                await this.callback();
-            } catch (error: unknown) {
-                this.onerror?.(error);
-            }
-        }
-    }
-
-    reset() {
-        this.clear();
-        this.timeoutId = setTimeout(() => {
-            void this.runCallback().finally(() => {
-                this.timeoutId = undefined;
-            });
-        }, this.timeoutMS);
-    }
-}
+import { TimeoutManager } from "./timeoutManager.js";
 
 export class SessionStore {
     private sessions: {

src/common/timeoutManager.ts

@@ -0,0 +1,63 @@
+/**
+ * A class that manages timeouts for a callback function.
+ * It is used to ensure that a callback function is called after a certain amount of time.
+ * If the callback function is not called after the timeout, it will be called with an error.
+ */
+export class TimeoutManager {
+    private timeoutId?: NodeJS.Timeout;
+
+    /**
+     * A callback function that is called when the timeout is reached.
+     */
+    public onerror?: (error: unknown) => void;
+
+    /**
+     * Creates a new TimeoutManager.
+     * @param callback - A callback function that is called when the timeout is reached.
+     * @param timeoutMS - The timeout in milliseconds.
+     */
+    constructor(
+        private readonly callback: () => Promise<void> | void,
+        private readonly timeoutMS: number
+    ) {
+        if (timeoutMS <= 0) {
+            throw new Error("timeoutMS must be greater than 0");
+        }
+        this.reset();
+    }
+
+    /**
+     * Clears the timeout.
+     */
+    clear() {
+        if (this.timeoutId) {
+            clearTimeout(this.timeoutId);
+            this.timeoutId = undefined;
+        }
+    }
+
+    /**
+     * Runs the callback function.
+     */
+    private async runCallback() {
+        if (this.callback) {
+            try {
+                await this.callback();
+            } catch (error: unknown) {
+                this.onerror?.(error);
+            }
+        }
+    }
+
+    /**
+     * Resets the timeout.
+     */
+    reset() {
+        this.clear();
+        this.timeoutId = setTimeout(() => {
+            void this.runCallback().finally(() => {
+                this.timeoutId = undefined;
+            });
+        }, this.timeoutMS);
+    }
+}

tests/unit/common/timeoutManager.test.ts

@@ -0,0 +1,79 @@
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
+import { TimeoutManager } from "../../../src/common/timeoutManager.js";
+
+describe("TimeoutManager", () => {
+    beforeAll(() => {
+        vi.useFakeTimers();
+    });
+
+    afterAll(() => {
+        vi.useRealTimers();
+    });
+
+    it("calls the timeout callback", () => {
+        const callback = vi.fn();
+
+        new TimeoutManager(callback, 1000);
+
+        vi.advanceTimersByTime(1000);
+        expect(callback).toHaveBeenCalled();
+    });
+
+    it("does not call the timeout callback if the timeout is cleared", () => {
+        const callback = vi.fn();
+
+        const timeoutManager = new TimeoutManager(callback, 1000);
+
+        vi.advanceTimersByTime(500);
+        timeoutManager.clear();
+        vi.advanceTimersByTime(500);
+
+        expect(callback).not.toHaveBeenCalled();
+    });
+
+    it("does not call the timeout callback if the timeout is reset", () => {
+        const callback = vi.fn();
+
+        const timeoutManager = new TimeoutManager(callback, 1000);
+
+        vi.advanceTimersByTime(500);
+        timeoutManager.reset();
+        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);
+
+            vi.advanceTimersByTime(500);
+            timeoutManager.reset();
+            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);
+
+            vi.advanceTimersByTime(500);
+            timeoutManager.reset();
+            vi.advanceTimersByTime(1000);
+            expect(callback).toHaveBeenCalled();
+        });
+    });
+});