Skip to content

Commit 3ba8a4a

Browse files
fmenezesfmenezes
authored
chore: address comments from #361 (#386)
1 parent 1297ed9 commit 3ba8a4a

File tree

5 files changed

+72
-112
lines changed

5 files changed

+72
-112
lines changed

‎src/common/logger.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,9 @@ export const LogId = {
4444
streamableHttpTransportStarted: mongoLogId(1_006_001),
4545
streamableHttpTransportSessionCloseFailure: mongoLogId(1_006_002),
4646
streamableHttpTransportSessionCloseNotification: mongoLogId(1_006_003),
47-
streamableHttpTransportRequestFailure: mongoLogId(1_006_004),
48-
streamableHttpTransportCloseFailure: mongoLogId(1_006_005),
47+
streamableHttpTransportSessionCloseNotificationFailure: mongoLogId(1_006_004),
48+
streamableHttpTransportRequestFailure: mongoLogId(1_006_005),
49+
streamableHttpTransportCloseFailure: mongoLogId(1_006_006),
4950
} as const;
5051

5152
export abstract class LoggerBase {

‎src/common/managedTimeout.ts

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
export interface ManagedTimeout {
2+
cancel: () => void;
3+
restart: () => void;
4+
}
5+
6+
export function setManagedTimeout(callback: () => Promise<void> | void, timeoutMS: number): ManagedTimeout {
7+
let timeoutId: NodeJS.Timeout | undefined = setTimeout(() => {
8+
void callback();
9+
}, timeoutMS);
10+
11+
function cancel() {
12+
clearTimeout(timeoutId);
13+
timeoutId = undefined;
14+
}
15+
16+
function restart() {
17+
cancel();
18+
timeoutId = setTimeout(() => {
19+
void callback();
20+
}, timeoutMS);
21+
}
22+
23+
return {
24+
cancel,
25+
restart,
26+
};
27+
}

‎src/common/sessionStore.ts

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
22
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
3-
import logger, { LogId, McpLogger } from "./logger.js";
4-
import { TimeoutManager } from "./timeoutManager.js";
3+
import logger, { LogId, LoggerBase, McpLogger } from "./logger.js";
4+
import { ManagedTimeout, setManagedTimeout } from "./managedTimeout.js";
55

66
export class SessionStore {
77
private sessions: {
88
[sessionId: string]: {
9-
mcpServer: McpServer;
9+
logger: LoggerBase;
1010
transport: StreamableHTTPServerTransport;
11-
abortTimeout: TimeoutManager;
12-
notificationTimeout: TimeoutManager;
11+
abortTimeout: ManagedTimeout;
12+
notificationTimeout: ManagedTimeout;
1313
};
1414
} = {};
1515

@@ -39,54 +39,61 @@ export class SessionStore {
3939
return;
4040
}
4141

42-
session.abortTimeout.reset();
42+
session.abortTimeout.restart();
4343

44-
session.notificationTimeout.reset();
44+
session.notificationTimeout.restart();
4545
}
4646

4747
private sendNotification(sessionId: string): void {
4848
const session = this.sessions[sessionId];
4949
if (!session) {
50+
logger.warning(
51+
LogId.streamableHttpTransportSessionCloseNotificationFailure,
52+
"sessionStore",
53+
`session ${sessionId} not found, no notification delivered`
54+
);
5055
return;
5156
}
52-
const logger = new McpLogger(session.mcpServer);
53-
logger.info(
57+
session.logger.info(
5458
LogId.streamableHttpTransportSessionCloseNotification,
5559
"sessionStore",
5660
"Session is about to be closed due to inactivity"
5761
);
5862
}
5963

6064
setSession(sessionId: string, transport: StreamableHTTPServerTransport, mcpServer: McpServer): void {
61-
if (this.sessions[sessionId]) {
65+
const session = this.sessions[sessionId];
66+
if (session) {
6267
throw new Error(`Session ${sessionId} already exists`);
6368
}
64-
const abortTimeout = new TimeoutManager(async () => {
65-
const logger = new McpLogger(mcpServer);
66-
logger.info(
67-
LogId.streamableHttpTransportSessionCloseNotification,
68-
"sessionStore",
69-
"Session closed due to inactivity"
70-
);
69+
const abortTimeout = setManagedTimeout(async () => {
70+
if (this.sessions[sessionId]) {
71+
this.sessions[sessionId].logger.info(
72+
LogId.streamableHttpTransportSessionCloseNotification,
73+
"sessionStore",
74+
"Session closed due to inactivity"
75+
);
7176

72-
await this.closeSession(sessionId);
77+
await this.closeSession(sessionId);
78+
}
7379
}, this.idleTimeoutMS);
74-
const notificationTimeout = new TimeoutManager(
80+
const notificationTimeout = setManagedTimeout(
7581
() => this.sendNotification(sessionId),
7682
this.notificationTimeoutMS
7783
);
78-
this.sessions[sessionId] = { mcpServer, transport, abortTimeout, notificationTimeout };
84+
this.sessions[sessionId] = { logger: new McpLogger(mcpServer), transport, abortTimeout, notificationTimeout };
7985
}
8086

8187
async closeSession(sessionId: string, closeTransport: boolean = true): Promise<void> {
82-
if (!this.sessions[sessionId]) {
88+
const session = this.sessions[sessionId];
89+
if (!session) {
8390
throw new Error(`Session ${sessionId} not found`);
8491
}
85-
this.sessions[sessionId].abortTimeout.clear();
86-
this.sessions[sessionId].notificationTimeout.clear();
92+
session.abortTimeout.cancel();
93+
session.notificationTimeout.cancel();
8794
if (closeTransport) {
8895
try {
89-
await this.sessions[sessionId].transport.close();
96+
await session.transport.close();
9097
} catch (error) {
9198
logger.error(
9299
LogId.streamableHttpTransportSessionCloseFailure,

‎src/common/timeoutManager.ts

Lines changed: 0 additions & 63 deletions
This file was deleted.

‎tests/unit/common/timeoutManager.test.ts renamed to ‎tests/unit/common/managedTimeout.test.ts

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
2-
import { TimeoutManager } from "../../../src/common/timeoutManager.js";
2+
import { setManagedTimeout } from "../../../src/common/managedTimeout.js";
33

4-
describe("TimeoutManager", () => {
4+
describe("setManagedTimeout", () => {
55
beforeAll(() => {
66
vi.useFakeTimers();
77
});
@@ -13,7 +13,7 @@ describe("TimeoutManager", () => {
1313
it("calls the timeout callback", () => {
1414
const callback = vi.fn();
1515

16-
new TimeoutManager(callback, 1000);
16+
setManagedTimeout(callback, 1000);
1717

1818
vi.advanceTimersByTime(1000);
1919
expect(callback).toHaveBeenCalled();
@@ -22,10 +22,10 @@ describe("TimeoutManager", () => {
2222
it("does not call the timeout callback if the timeout is cleared", () => {
2323
const callback = vi.fn();
2424

25-
const timeoutManager = new TimeoutManager(callback, 1000);
25+
const timeout = setManagedTimeout(callback, 1000);
2626

2727
vi.advanceTimersByTime(500);
28-
timeoutManager.clear();
28+
timeout.cancel();
2929
vi.advanceTimersByTime(500);
3030

3131
expect(callback).not.toHaveBeenCalled();
@@ -34,44 +34,32 @@ describe("TimeoutManager", () => {
3434
it("does not call the timeout callback if the timeout is reset", () => {
3535
const callback = vi.fn();
3636

37-
const timeoutManager = new TimeoutManager(callback, 1000);
37+
const timeout = setManagedTimeout(callback, 1000);
3838

3939
vi.advanceTimersByTime(500);
40-
timeoutManager.reset();
40+
timeout.restart();
4141
vi.advanceTimersByTime(500);
4242
expect(callback).not.toHaveBeenCalled();
4343
});
4444

45-
it("calls the onerror callback", () => {
46-
const onerrorCallback = vi.fn();
47-
48-
const tm = new TimeoutManager(() => {
49-
throw new Error("test");
50-
}, 1000);
51-
tm.onerror = onerrorCallback;
52-
53-
vi.advanceTimersByTime(1000);
54-
expect(onerrorCallback).toHaveBeenCalled();
55-
});
56-
5745
describe("if timeout is reset", () => {
5846
it("does not call the timeout callback within the timeout period", () => {
5947
const callback = vi.fn();
6048

61-
const timeoutManager = new TimeoutManager(callback, 1000);
49+
const timeout = setManagedTimeout(callback, 1000);
6250

6351
vi.advanceTimersByTime(500);
64-
timeoutManager.reset();
52+
timeout.restart();
6553
vi.advanceTimersByTime(500);
6654
expect(callback).not.toHaveBeenCalled();
6755
});
6856
it("calls the timeout callback after the timeout period", () => {
6957
const callback = vi.fn();
7058

71-
const timeoutManager = new TimeoutManager(callback, 1000);
59+
const timeout = setManagedTimeout(callback, 1000);
7260

7361
vi.advanceTimersByTime(500);
74-
timeoutManager.reset();
62+
timeout.restart();
7563
vi.advanceTimersByTime(1000);
7664
expect(callback).toHaveBeenCalled();
7765
});

0 commit comments

Comments
 (0)