Add inactivity timeout to ServerRedisTransport
- Add 5-minute inactivity timeout to prevent memory leaks - Reset timer on each message received from client - Clear timer on transport close - Add comprehensive tests for timeout behavior - Fix async/await in redisRelayToMcpServer calls Sessions will now automatically shutdown after 5 minutes of no client activity. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
728c3f66ea
commit
d779401b5f
4 changed files with 3995 additions and 4802 deletions
File diff suppressed because one or more lines are too long
|
|
@ -44,7 +44,7 @@ describe('Redis Transport Integration', () => {
|
||||||
};
|
};
|
||||||
|
|
||||||
// 3. Set up the Redis relay (this is what happens in the HTTP handler)
|
// 3. Set up the Redis relay (this is what happens in the HTTP handler)
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||||
|
|
||||||
// Track messages received by server
|
// Track messages received by server
|
||||||
const serverReceivedMessages: JSONRPCMessage[] = [];
|
const serverReceivedMessages: JSONRPCMessage[] = [];
|
||||||
|
|
@ -148,7 +148,7 @@ describe('Redis Transport Integration', () => {
|
||||||
start: jest.fn(() => Promise.resolve())
|
start: jest.fn(() => Promise.resolve())
|
||||||
};
|
};
|
||||||
|
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||||
|
|
||||||
// Send tools/list request
|
// Send tools/list request
|
||||||
const toolsListMessage: JSONRPCMessage = {
|
const toolsListMessage: JSONRPCMessage = {
|
||||||
|
|
@ -202,7 +202,7 @@ describe('Redis Transport Integration', () => {
|
||||||
start: jest.fn(() => Promise.resolve())
|
start: jest.fn(() => Promise.resolve())
|
||||||
};
|
};
|
||||||
|
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||||
|
|
||||||
// Set up notification subscription manually since notifications don't have an id
|
// Set up notification subscription manually since notifications don't have an id
|
||||||
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
|
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
|
||||||
|
|
@ -268,7 +268,7 @@ describe('Redis Transport Integration', () => {
|
||||||
start: jest.fn(() => Promise.resolve())
|
start: jest.fn(() => Promise.resolve())
|
||||||
};
|
};
|
||||||
|
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockClientTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||||
|
|
||||||
// Send a notification (no id field)
|
// Send a notification (no id field)
|
||||||
const notification: JSONRPCMessage = {
|
const notification: JSONRPCMessage = {
|
||||||
|
|
|
||||||
|
|
@ -172,7 +172,7 @@ describe('Redis Transport', () => {
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should set up message relay from transport to server', async () => {
|
it('should set up message relay from transport to server', async () => {
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||||
|
|
||||||
// Simulate a message from the transport
|
// Simulate a message from the transport
|
||||||
const requestMessage: JSONRPCMessage = {
|
const requestMessage: JSONRPCMessage = {
|
||||||
|
|
@ -203,7 +203,7 @@ describe('Redis Transport', () => {
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should subscribe to response channel for request messages', async () => {
|
it('should subscribe to response channel for request messages', async () => {
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||||
|
|
||||||
const requestMessage: JSONRPCMessage = {
|
const requestMessage: JSONRPCMessage = {
|
||||||
jsonrpc: '2.0',
|
jsonrpc: '2.0',
|
||||||
|
|
@ -241,7 +241,7 @@ describe('Redis Transport', () => {
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should not subscribe for notification messages (no id)', async () => {
|
it('should not subscribe for notification messages (no id)', async () => {
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, mockTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||||
|
|
||||||
const notificationMessage: JSONRPCMessage = {
|
const notificationMessage: JSONRPCMessage = {
|
||||||
jsonrpc: '2.0',
|
jsonrpc: '2.0',
|
||||||
|
|
@ -337,7 +337,7 @@ describe('Redis Transport', () => {
|
||||||
start: jest.fn(() => Promise.resolve())
|
start: jest.fn(() => Promise.resolve())
|
||||||
};
|
};
|
||||||
|
|
||||||
const cleanup = redisRelayToMcpServer(sessionId, clientTransport);
|
const cleanup = await redisRelayToMcpServer(sessionId, clientTransport);
|
||||||
|
|
||||||
// Client sends a request
|
// Client sends a request
|
||||||
const listToolsRequest: JSONRPCMessage = {
|
const listToolsRequest: JSONRPCMessage = {
|
||||||
|
|
@ -434,4 +434,104 @@ describe('Redis Transport', () => {
|
||||||
expect(onCloseMock).toHaveBeenCalled();
|
expect(onCloseMock).toHaveBeenCalled();
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('Inactivity Timeout', () => {
|
||||||
|
const sessionId = 'test-inactivity-session';
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
jest.useFakeTimers({ doNotFake: ['setImmediate', 'nextTick'] });
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
jest.useRealTimers();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should shutdown session after 5 minutes of inactivity', async () => {
|
||||||
|
const transport = new ServerRedisTransport(sessionId);
|
||||||
|
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||||
|
|
||||||
|
await transport.start();
|
||||||
|
|
||||||
|
// Fast-forward time by 5 minutes
|
||||||
|
jest.advanceTimersByTime(5 * 60 * 1000);
|
||||||
|
|
||||||
|
// Should have published shutdown control message
|
||||||
|
expect(shutdownSpy).toHaveBeenCalledWith(
|
||||||
|
`mcp:control:${sessionId}`,
|
||||||
|
expect.stringContaining('"action":"SHUTDOWN"')
|
||||||
|
);
|
||||||
|
|
||||||
|
await transport.close();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should reset timeout when message is received', async () => {
|
||||||
|
const transport = new ServerRedisTransport(sessionId);
|
||||||
|
const onMessageMock = jest.fn();
|
||||||
|
transport.onmessage = onMessageMock;
|
||||||
|
|
||||||
|
await transport.start();
|
||||||
|
|
||||||
|
// Fast-forward 4 minutes
|
||||||
|
jest.advanceTimersByTime(4 * 60 * 1000);
|
||||||
|
|
||||||
|
// Manually publish a message to trigger the subscription handler
|
||||||
|
const testMessage = { jsonrpc: '2.0', method: 'ping' };
|
||||||
|
await mockRedis.publish(
|
||||||
|
`mcp:shttp:toserver:${sessionId}`,
|
||||||
|
JSON.stringify({
|
||||||
|
type: 'mcp',
|
||||||
|
message: testMessage
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
// Wait for message to be processed
|
||||||
|
await new Promise(resolve => setImmediate(resolve));
|
||||||
|
|
||||||
|
// Verify message was received
|
||||||
|
expect(onMessageMock).toHaveBeenCalledWith(testMessage, undefined);
|
||||||
|
|
||||||
|
// Clear the publish spy to check only future calls
|
||||||
|
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||||
|
shutdownSpy.mockClear();
|
||||||
|
|
||||||
|
// Fast-forward 4 more minutes (total 8, but only 4 since last message)
|
||||||
|
jest.advanceTimersByTime(4 * 60 * 1000);
|
||||||
|
|
||||||
|
// Should not have shutdown yet
|
||||||
|
expect(shutdownSpy).not.toHaveBeenCalledWith(
|
||||||
|
`mcp:control:${sessionId}`,
|
||||||
|
expect.stringContaining('"action":"SHUTDOWN"')
|
||||||
|
);
|
||||||
|
|
||||||
|
// Fast-forward 2 more minutes to exceed timeout
|
||||||
|
jest.advanceTimersByTime(2 * 60 * 1000);
|
||||||
|
|
||||||
|
// Now should have shutdown
|
||||||
|
expect(shutdownSpy).toHaveBeenCalledWith(
|
||||||
|
`mcp:control:${sessionId}`,
|
||||||
|
expect.stringContaining('"action":"SHUTDOWN"')
|
||||||
|
);
|
||||||
|
|
||||||
|
await transport.close();
|
||||||
|
}, 10000);
|
||||||
|
|
||||||
|
it('should clear timeout on close', async () => {
|
||||||
|
const transport = new ServerRedisTransport(sessionId);
|
||||||
|
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||||
|
|
||||||
|
await transport.start();
|
||||||
|
|
||||||
|
// Close transport before timeout
|
||||||
|
await transport.close();
|
||||||
|
|
||||||
|
// Fast-forward past timeout
|
||||||
|
jest.advanceTimersByTime(10 * 60 * 1000);
|
||||||
|
|
||||||
|
// Should not have triggered shutdown
|
||||||
|
expect(shutdownSpy).not.toHaveBeenCalledWith(
|
||||||
|
`mcp:control:${sessionId}`,
|
||||||
|
expect.stringContaining('"action":"SHUTDOWN"')
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|
@ -134,6 +134,8 @@ export class ServerRedisTransport implements Transport {
|
||||||
private controlCleanup?: (() => Promise<void>);
|
private controlCleanup?: (() => Promise<void>);
|
||||||
private serverCleanup?: (() => Promise<void>);
|
private serverCleanup?: (() => Promise<void>);
|
||||||
private shouldShutdown = false;
|
private shouldShutdown = false;
|
||||||
|
private inactivityTimeout?: NodeJS.Timeout;
|
||||||
|
private readonly INACTIVITY_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
|
||||||
|
|
||||||
onclose?: (() => void) | undefined;
|
onclose?: (() => void) | undefined;
|
||||||
onerror?: ((error: Error) => void) | undefined;
|
onerror?: ((error: Error) => void) | undefined;
|
||||||
|
|
@ -144,7 +146,29 @@ export class ServerRedisTransport implements Transport {
|
||||||
this._sessionId = sessionId;
|
this._sessionId = sessionId;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private resetInactivityTimer(): void {
|
||||||
|
// Clear existing timeout if any
|
||||||
|
if (this.inactivityTimeout) {
|
||||||
|
clearTimeout(this.inactivityTimeout);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set new timeout
|
||||||
|
this.inactivityTimeout = setTimeout(() => {
|
||||||
|
console.log(`Session ${this._sessionId} timed out due to inactivity`);
|
||||||
|
void shutdownSession(this._sessionId);
|
||||||
|
}, this.INACTIVITY_TIMEOUT_MS);
|
||||||
|
}
|
||||||
|
|
||||||
|
private clearInactivityTimer(): void {
|
||||||
|
if (this.inactivityTimeout) {
|
||||||
|
clearTimeout(this.inactivityTimeout);
|
||||||
|
this.inactivityTimeout = undefined;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async start(): Promise<void> {
|
async start(): Promise<void> {
|
||||||
|
// Start inactivity timer
|
||||||
|
this.resetInactivityTimer();
|
||||||
|
|
||||||
// Subscribe to MCP messages from clients
|
// Subscribe to MCP messages from clients
|
||||||
const serverChannel = getToServerChannel(this._sessionId);
|
const serverChannel = getToServerChannel(this._sessionId);
|
||||||
|
|
@ -153,6 +177,8 @@ export class ServerRedisTransport implements Transport {
|
||||||
(messageJson) => {
|
(messageJson) => {
|
||||||
const redisMessage = JSON.parse(messageJson) as RedisMessage;
|
const redisMessage = JSON.parse(messageJson) as RedisMessage;
|
||||||
if (redisMessage.type === 'mcp') {
|
if (redisMessage.type === 'mcp') {
|
||||||
|
// Reset inactivity timer on each message from client
|
||||||
|
this.resetInactivityTimer();
|
||||||
this.onmessage?.(redisMessage.message, redisMessage.extra);
|
this.onmessage?.(redisMessage.message, redisMessage.extra);
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|
@ -191,6 +217,8 @@ export class ServerRedisTransport implements Transport {
|
||||||
}
|
}
|
||||||
|
|
||||||
async close(): Promise<void> {
|
async close(): Promise<void> {
|
||||||
|
// Clear inactivity timer
|
||||||
|
this.clearInactivityTimer();
|
||||||
|
|
||||||
// Clean up server message subscription
|
// Clean up server message subscription
|
||||||
if (this.serverCleanup) {
|
if (this.serverCleanup) {
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue