diff --git a/.vscode/launch.json b/.vscode/launch.json index 435c9f3..bbc2970 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -13,6 +13,67 @@ "env": { "NODE_ENV": "development" } + }, + { + "name": "Debug Jest Tests", + "type": "node", + "request": "launch", + "runtimeArgs": [ + "--experimental-vm-modules", + "--inspect-brk", + "${workspaceFolder}/node_modules/.bin/jest", + "--runInBand" + ], + "console": "integratedTerminal", + "internalConsoleOptions": "neverOpen", + "skipFiles": ["/**"], + "env": { + "NODE_OPTIONS": "--experimental-vm-modules" + } + }, + { + "name": "Debug Current Jest Test File", + "type": "node", + "request": "launch", + "runtimeArgs": [ + "--experimental-vm-modules", + "--inspect-brk", + "${workspaceFolder}/node_modules/.bin/jest", + "--runInBand", + "${relativeFile}" + ], + "console": "integratedTerminal", + "internalConsoleOptions": "neverOpen", + "skipFiles": ["/**"], + "env": { + "NODE_OPTIONS": "--experimental-vm-modules" + } + }, + { + "name": "Debug Jest Test by Pattern", + "type": "node", + "request": "launch", + "runtimeArgs": [ + "--experimental-vm-modules", + "--inspect-brk", + "${workspaceFolder}/node_modules/.bin/jest", + "--runInBand", + "--testNamePattern", + "${input:testNamePattern}" + ], + "console": "integratedTerminal", + "internalConsoleOptions": "neverOpen", + "skipFiles": ["/**"], + "env": { + "NODE_OPTIONS": "--experimental-vm-modules" + } + } + ], + "inputs": [ + { + "id": "testNamePattern", + "type": "promptString", + "description": "Test name pattern to match (e.g., 'should trigger onsessionclosed')" } ] } \ No newline at end of file diff --git a/src/handlers/shttp.integration.test.ts b/src/handlers/shttp.integration.test.ts index f9c70c0..a7c6ec3 100644 --- a/src/handlers/shttp.integration.test.ts +++ b/src/handlers/shttp.integration.test.ts @@ -5,6 +5,7 @@ import { MockRedisClient, setRedisClient } from '../redis.js'; import { handleStreamableHTTP } from './shttp.js'; import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js'; import { randomUUID } from 'crypto'; +import { shutdownSession } from '../services/redisTransport.js'; describe('Streamable HTTP Handler Integration Tests', () => { let mockRedis: MockRedisClient; @@ -21,6 +22,8 @@ describe('Streamable HTTP Handler Integration Tests', () => { status: jest.fn().mockReturnThis(), json: jest.fn().mockReturnThis(), on: jest.fn().mockReturnThis(), + once: jest.fn().mockReturnThis(), + emit: jest.fn().mockReturnThis(), headersSent: false, setHeader: jest.fn().mockReturnThis(), writeHead: jest.fn().mockReturnThis(), @@ -28,13 +31,17 @@ describe('Streamable HTTP Handler Integration Tests', () => { end: jest.fn().mockReturnThis(), getHeader: jest.fn(), removeHeader: jest.fn().mockReturnThis(), - } as Partial; + socket: { + setTimeout: jest.fn(), + }, + } as unknown as Partial; // Create mock request mockReq = { method: 'POST', headers: { 'content-type': 'application/json', + 'accept': 'application/json, text/event-stream', 'mcp-protocol-version': '2024-11-05', }, body: {}, @@ -57,12 +64,12 @@ describe('Streamable HTTP Handler Integration Tests', () => { }; // Helper to extract session ID from test context - const getSessionIdFromTest = () => { + const getSessionIdFromTest = (): string | undefined => { // Try to get from response headers first const setHeaderCalls = (mockRes.setHeader as jest.Mock).mock.calls; const sessionIdHeader = setHeaderCalls.find(([name]) => name === 'mcp-session-id'); if (sessionIdHeader?.[1]) { - return sessionIdHeader[1]; + return sessionIdHeader[1] as string; } // Fall back to extracting from Redis channels @@ -105,6 +112,10 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Wait longer for async initialization to complete await new Promise(resolve => setTimeout(resolve, 200)); + + // get the sessionId from the response + const sessionId = getSessionIdFromTest(); + expect(sessionId).toBeDefined(); // Check if any subscriptions were created on any channels // Since we don't know the exact sessionId generated, check all channels @@ -129,6 +140,10 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Verify cleanup handler was registered expect(mockRes.on).toHaveBeenCalledWith('finish', expect.any(Function)); + + if (sessionId) { + await shutdownSession(sessionId) + } }); it('should handle cleanup errors gracefully', async () => { @@ -175,7 +190,9 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Send DELETE request to clean up MCP server jest.clearAllMocks(); mockReq.method = 'DELETE'; - mockReq.headers['mcp-session-id'] = cleanupSessionId; + if (mockReq.headers) { + mockReq.headers['mcp-session-id'] = cleanupSessionId; + } mockReq.body = {}; await handleStreamableHTTP(mockReq as Request, mockRes as Response); @@ -223,7 +240,7 @@ describe('Streamable HTTP Handler Integration Tests', () => { let sessionId: string | undefined; if (jsonCalls.length > 0) { - const response = jsonCalls[0][0]; + const response = jsonCalls[0][0] as any; if (response?.result?._meta?.sessionId) { sessionId = response.result._meta.sessionId; } @@ -237,7 +254,7 @@ describe('Streamable HTTP Handler Integration Tests', () => { try { // SSE data format: "data: {...}\n\n" const jsonStr = data.replace(/^data: /, '').trim(); - const parsed = JSON.parse(jsonStr); + const parsed = JSON.parse(jsonStr) as any; if (parsed?.result?._meta?.sessionId) { sessionId = parsed.result._meta.sessionId; } @@ -327,7 +344,7 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Check JSON responses const jsonCalls = (mockRes.json as jest.Mock).mock.calls; if (jsonCalls.length > 0) { - const response = jsonCalls[0][0]; + const response = jsonCalls[0][0] as any; if (response?.result?._meta?.sessionId) { sessionId = response.result._meta.sessionId; } @@ -340,7 +357,7 @@ describe('Streamable HTTP Handler Integration Tests', () => { if (typeof data === 'string' && data.includes('sessionId')) { try { const jsonStr = data.replace(/^data: /, '').trim(); - const parsed = JSON.parse(jsonStr); + const parsed = JSON.parse(jsonStr) as any; if (parsed?.result?._meta?.sessionId) { sessionId = parsed.result._meta.sessionId; } @@ -374,13 +391,17 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Should return 401 for unauthorized access to another user's session expect(mockRes.status).toHaveBeenCalledWith(401); + + // shutdown the session + if (sessionId) { + await shutdownSession(sessionId) + } }); }); describe('User Session Isolation', () => { it('should prevent users from accessing sessions created by other users', async () => { // Create session for user 1 - const sessionId = randomUUID(); const user1Auth: AuthInfo = { clientId: 'user1-client', token: 'user1-token', @@ -421,7 +442,7 @@ describe('Streamable HTTP Handler Integration Tests', () => { // Check JSON responses const jsonCalls = (mockRes.json as jest.Mock).mock.calls; if (jsonCalls.length > 0) { - const response = jsonCalls[0][0]; + const response = jsonCalls[0][0] as any; if (response?.result?._meta?.sessionId) { actualSessionId = response.result._meta.sessionId; } diff --git a/src/handlers/shttp.ts b/src/handlers/shttp.ts index 45810ba..5ee1e10 100644 --- a/src/handlers/shttp.ts +++ b/src/handlers/shttp.ts @@ -27,6 +27,10 @@ export async function handleStreamableHTTP(req: Request, res: Response) { await shttpTransport?.close(); }); + const onsessionclosed = async (sessionId: string) => { + await shutdownSession(sessionId); + } + try { // Check for existing session ID const sessionId = req.headers['mcp-session-id'] as string | undefined; @@ -45,7 +49,7 @@ export async function handleStreamableHTTP(req: Request, res: Response) { return; } // Reuse existing transport for owned session - shttpTransport = await getShttpTransport(sessionId); + shttpTransport = await getShttpTransport(sessionId, onsessionclosed); } else if (isInitializeRequest(req.body)) { // New initialization request - use JSON response mode const onsessioninitialized = async (sessionId: string) => { @@ -59,16 +63,10 @@ export async function handleStreamableHTTP(req: Request, res: Response) { await setSessionOwner(sessionId, userId); } - const onsessionclosed = (sessionId: string) => { - console.log(`Session ${sessionId} closing`); - void shutdownSession(sessionId); - console.log(`Session ${sessionId} closed`); - } - const sessionId = randomUUID(); shttpTransport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => sessionId, - onsessionclosed, // dev build only + onsessionclosed, onsessioninitialized, }); shttpTransport.onclose = redisRelayToMcpServer(sessionId, shttpTransport); diff --git a/src/services/mcp.ts b/src/services/mcp.ts index 0f9d11d..04a5e65 100644 --- a/src/services/mcp.ts +++ b/src/services/mcp.ts @@ -674,7 +674,6 @@ export const createMcpServer = (): McpServerWrapper => { }); const cleanup = async () => { - console.log("Cleaning up MCP server"); if (subsUpdateInterval) clearInterval(subsUpdateInterval); if (logsUpdateInterval) clearInterval(logsUpdateInterval); if (stdErrUpdateInterval) clearInterval(stdErrUpdateInterval); diff --git a/src/services/redisTransport.ts b/src/services/redisTransport.ts index 955661b..8697f2b 100644 --- a/src/services/redisTransport.ts +++ b/src/services/redisTransport.ts @@ -202,11 +202,12 @@ export class ServerRedisTransport implements Transport { } } -export async function getShttpTransport(sessionId: string): Promise { +export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise): Promise { // Giving undefined here and setting the sessionId means the // transport wont try to create a new session. const shttpTransport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined, + onsessionclosed, }) shttpTransport.sessionId = sessionId;