diff --git a/server/src/mcp/server.ts b/server/src/mcp/server.ts index 2a3514d..0b331f9 100644 --- a/server/src/mcp/server.ts +++ b/server/src/mcp/server.ts @@ -13,7 +13,7 @@ * Protocol versions: 2025-03-26, 2024-11-05 */ -import { createServer, IncomingMessage, ServerResponse } from "http"; +import { createServer, IncomingMessage, Server, ServerResponse } from "http"; import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"; import { config } from "./config"; @@ -96,6 +96,22 @@ function sendJson(res: ServerResponse, status: number, data: unknown): void { res.end(JSON.stringify(data)); } +/** + * Send 404 for an unknown or expired session. + * Per the MCP spec, clients must re-initialize when they receive this. + * Matches the SDK's own session-not-found response. + */ +function sendSessionNotFound(res: ServerResponse): void { + sendJson(res, 404, { + jsonrpc: "2.0", + error: { + code: -32001, + message: "Session not found", + }, + id: null, + }); +} + /** * Parse request body as JSON */ @@ -122,11 +138,32 @@ interface SessionContext { } const sessions = new Map(); +// sessionStore owns TTL and per-key limits. Close the transport whenever +// it drops a session, so the two maps never drift apart. +sessionStore.onDelete((sessionId) => { + const sessionCtx = sessions.get(sessionId); + if (!sessionCtx) return; + + sessions.delete(sessionId); + sessionCtx.transport.close().catch((error) => { + console.error("[MCP] Error closing transport:", error); + }); +}); + /** - * Start the MCP server + * Look up a live session. Goes through sessionStore so each request + * refreshes the TTL and an expired session is dropped. */ -export function startMcpServer(): void { - const server = createServer( +function getActiveSession(sessionId: string): SessionContext | undefined { + if (!sessionStore.validate(sessionId)) return undefined; + return sessions.get(sessionId); +} + +/** + * Create the MCP HTTP server without starting it + */ +export function createMcpHttpServer(): Server { + return createServer( async (req: IncomingMessage, res: ServerResponse) => { // CORS headers res.setHeader("Access-Control-Allow-Origin", "*"); @@ -198,11 +235,7 @@ export function startMcpServer(): void { // Handle DELETE - session termination if (req.method === "DELETE") { if (sessionId) { - const sessionCtx = sessions.get(sessionId); - if (sessionCtx) { - await sessionCtx.transport.close(); - sessions.delete(sessionId); - } + // The onDelete listener closes the transport sessionStore.delete(sessionId); } res.writeHead(204); @@ -216,24 +249,15 @@ export function startMcpServer(): void { // Get or create transport for this session let transport: StreamableHTTPServerTransport; - if (sessionId && sessions.has(sessionId)) { + if (sessionId) { // Existing session - reuse transport - const sessionCtx = sessions.get(sessionId); + const sessionCtx = getActiveSession(sessionId); if (!sessionCtx) { - throw new Error("Session not found despite has() check"); + // Unknown, expired, or evicted session + sendSessionNotFound(res); + return; } transport = sessionCtx.transport; - } else if (sessionId) { - // Session ID provided but transport not found - session expired - sendJson(res, 400, { - jsonrpc: "2.0", - error: { - code: -32000, - message: "Bad Request: Session not found or expired", - }, - id: null, - }); - return; } else { // No session ID - this should be an initialize request // Create new MCP server instance for this session @@ -283,16 +307,19 @@ export function startMcpServer(): void { // Handle GET - SSE streaming (for server-initiated messages) if (req.method === "GET") { - if (!sessionId || !sessions.has(sessionId)) { - sendJson(res, 400, { - error: "Session ID required for SSE connection", + // No session: we offer no standalone SSE stream, so per the MCP + // spec answer 405 rather than 400 + if (!sessionId) { + res.setHeader("Allow", "POST, DELETE"); + sendJson(res, 405, { + error: "SSE stream requires an initialized session", }); return; } - const sessionCtx = sessions.get(sessionId); + const sessionCtx = getActiveSession(sessionId); if (!sessionCtx) { - sendJson(res, 400, { error: "Session not found" }); + sendSessionNotFound(res); return; } await sessionCtx.transport.handleRequest(req, res); @@ -308,6 +335,13 @@ export function startMcpServer(): void { sendJson(res, 404, { error: "Not found" }); }, ); +} + +/** + * Start the MCP server + */ +export function startMcpServer(): void { + const server = createMcpHttpServer(); server.listen(config.port, config.host, () => { console.log( diff --git a/server/src/mcp/sessions.ts b/server/src/mcp/sessions.ts index 26847e3..71c5ebb 100644 --- a/server/src/mcp/sessions.ts +++ b/server/src/mcp/sessions.ts @@ -14,11 +14,14 @@ interface Session { initialized: boolean; } +type SessionDeleteListener = (sessionId: string) => void; + /** * Simple in-memory session store with TTL and limits */ export class SessionStore { private sessions = new Map(); + private deleteListeners: SessionDeleteListener[] = []; private cleanupInterval: ReturnType | null = null; constructor() { @@ -92,18 +95,32 @@ export class SessionStore { return true; } + /** + * Register a callback that runs whenever a session is removed + * (explicit delete, TTL expiry, or per-key eviction) + */ + onDelete(listener: SessionDeleteListener): void { + this.deleteListeners.push(listener); + } + /** * Delete a session */ delete(sessionId: string): boolean { const deleted = this.sessions.delete(sessionId); - if (deleted && config.debug) { + if (!deleted) return false; + + if (config.debug) { // Sanitize sessionId before logging to prevent log injection (CWE-117) // eslint-disable-next-line no-control-regex const safeId = sessionId.replace(/[\r\n\t\x00-\x1f\x7f]/g, "_"); console.log(`[MCP] Session deleted: ${safeId}`); } - return deleted; + + for (const listener of this.deleteListeners) { + listener(sessionId); + } + return true; } /** @@ -131,7 +148,7 @@ export class SessionStore { for (const [id, session] of this.sessions) { if (now - session.lastAccessedAt > config.sessionTtlMs) { - this.sessions.delete(id); + this.delete(id); cleaned++; } } diff --git a/server/tests/mcp.test.ts b/server/tests/mcp.test.ts index 1c55798..eb5f946 100644 --- a/server/tests/mcp.test.ts +++ b/server/tests/mcp.test.ts @@ -12,8 +12,10 @@ import path from 'path'; // Import MCP components for unit testing import { SessionStore, sessionStore as globalSessionStore } from '../src/mcp/sessions'; -import { loadConfig } from '../src/mcp/config'; +import { config, loadConfig } from '../src/mcp/config'; import { tools, getTool, executeTool } from '../src/mcp/tools/registry'; +import { createMcpHttpServer } from '../src/mcp/server'; +import { AddressInfo } from 'net'; const TEST_DATA_DIR = path.join(__dirname, '../test-data-mcp'); @@ -215,6 +217,70 @@ describe('MCP Session Store', () => { expect(store.count()).toBe(6); }); }); + + describe('Delete Listeners', () => { + afterEach(() => { + jest.useRealTimers(); + }); + + it('should notify listeners on explicit delete', () => { + const listener = jest.fn(); + store.onDelete(listener); + const session = store.create('test-api-key'); + + store.delete(session.id); + + expect(listener).toHaveBeenCalledTimes(1); + expect(listener).toHaveBeenCalledWith(session.id); + }); + + it('should not notify listeners when deleting a non-existent session', () => { + const listener = jest.fn(); + store.onDelete(listener); + + store.delete('non-existent'); + + expect(listener).not.toHaveBeenCalled(); + }); + + it('should notify listeners when a session is evicted by the per-key limit', () => { + const listener = jest.fn(); + store.onDelete(listener); + const oldest = store.create('same-api-key'); + for (let i = 0; i < config.maxSessionsPerKey; i++) { + store.create('same-api-key'); + } + + expect(listener).toHaveBeenCalledWith(oldest.id); + }); + + it('should notify listeners when get() finds an expired session', () => { + jest.useFakeTimers({ now: Date.now() }); + const listener = jest.fn(); + store.onDelete(listener); + const session = store.create('test-api-key'); + + jest.setSystemTime(Date.now() + config.sessionTtlMs + 1); + + expect(store.get(session.id)).toBeUndefined(); + expect(listener).toHaveBeenCalledWith(session.id); + }); + + it('should notify listeners when the cleanup interval removes an expired session', () => { + jest.useFakeTimers({ now: Date.now() }); + // Create after faking timers so the cleanup interval is faked too + const timedStore = new SessionStore(); + const listener = jest.fn(); + timedStore.onDelete(listener); + const session = timedStore.create('test-api-key'); + + jest.advanceTimersByTime(config.sessionTtlMs + 60000); + + expect(timedStore.count()).toBe(0); + expect(listener).toHaveBeenCalledWith(session.id); + timedStore.stopCleanup(); + }); + }); }); // ============================================================================ @@ -607,6 +673,171 @@ describe('MCP HTTP Server Integration', () => { }); }); +// ============================================================================ +// Integration Tests: MCP Session Status Codes (real server) +// ============================================================================ + +describe('MCP Session Status Codes', () => { + let server: Server; + let baseUrl: string; + + // Only needed when MCP_API_KEY is set in the environment + const authHeaders: Record = process.env.MCP_API_KEY + ? { Authorization: `Bearer ${process.env.MCP_API_KEY}` } + : {}; + + beforeAll(async () => { + server = createMcpHttpServer(); + await new Promise((resolve) => server.listen(0, '127.0.0.1', () => resolve())); + const { port } = server.address() as AddressInfo; + baseUrl = `http://127.0.0.1:${port}/mcp`; + }); + + afterAll(async () => { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + }); + + async function initializeSession(): Promise { + const response = await fetch(baseUrl, { + method: 'POST', + headers: { + ...authHeaders, + 'Content-Type': 'application/json', + Accept: 'application/json, text/event-stream', + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: 1, + method: 'initialize', + params: { + protocolVersion: '2025-06-18', + capabilities: {}, + clientInfo: { name: 'codex-test', version: '1.0.0' }, + }, + }), + }); + + expect(response.status).toBe(200); + const sessionId = response.headers.get('mcp-session-id'); + if (!sessionId) throw new Error('No session ID returned'); + return sessionId; + } + + it('should return 405 for GET without a session ID', async () => { + const response = await fetch(baseUrl, { + headers: { ...authHeaders, Accept: 'text/event-stream' }, + }); + + expect(response.status).toBe(405); + expect(response.headers.get('allow')).toBe('POST, DELETE'); + }); + + it('should return 404 for GET with an unknown session ID', async () => { + const response = await fetch(baseUrl, { + headers: { + ...authHeaders, + Accept: 'text/event-stream', + 'Mcp-Session-Id': 'unknown-session', + }, + }); + const data = await response.json() as { error: { code: number } }; + + expect(response.status).toBe(404); + expect(data.error.code).toBe(-32001); + }); + + it('should return 404 for POST with an unknown session ID', async () => { + const response = await fetch(baseUrl, { + method: 'POST', + headers: { + ...authHeaders, + 'Content-Type': 'application/json', + Accept: 'application/json, text/event-stream', + 'Mcp-Session-Id': 'unknown-session', + }, + body: JSON.stringify({ jsonrpc: '2.0', id: 2, method: 'tools/list' }), + }); + const data = await response.json() as { error: { code: number } }; + + expect(response.status).toBe(404); + expect(data.error.code).toBe(-32001); + }); + + it('should open an SSE stream for GET with a valid session ID', async () => { + const sessionId = await initializeSession(); + const controller = new AbortController(); + + const response = await fetch(baseUrl, { + headers: { + ...authHeaders, + Accept: 'text/event-stream', + 'Mcp-Session-Id': sessionId, + 'Mcp-Protocol-Version': '2025-06-18', + }, + signal: controller.signal, + }); + + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toContain('text/event-stream'); + controller.abort(); + }); + + async function getWithSession(sessionId: string): Promise { + return fetch(baseUrl, { + headers: { + ...authHeaders, + Accept: 'text/event-stream', + 'Mcp-Session-Id': sessionId, + }, + }); + } + + it('should return 404 after the session is terminated with DELETE', async () => { + const sessionId = await initializeSession(); + + const deleteResponse = await fetch(baseUrl, { + method: 'DELETE', + headers: { ...authHeaders, 'Mcp-Session-Id': sessionId }, + }); + expect(deleteResponse.status).toBe(204); + expect(globalSessionStore.validate(sessionId)).toBe(false); + + const response = await getWithSession(sessionId); + expect(response.status).toBe(404); + }); + + it('should return 404 after the session store drops the session', async () => { + const sessionId = await initializeSession(); + + // Same path as TTL expiry and the cleanup interval + globalSessionStore.delete(sessionId); + + const response = await getWithSession(sessionId); + expect(response.status).toBe(404); + }); + + it('should return 404 for a session evicted by the per-key limit', async () => { + const oldestId = await initializeSession(); + for (let i = 0; i < config.maxSessionsPerKey; i++) { + await initializeSession(); + } + + const response = await fetch(baseUrl, { + method: 'POST', + headers: { + ...authHeaders, + 'Content-Type': 'application/json', + Accept: 'application/json, text/event-stream', + 'Mcp-Session-Id': oldestId, + }, + body: JSON.stringify({ jsonrpc: '2.0', id: 3, method: 'tools/list' }), + }); + + expect(response.status).toBe(404); + }); +}); + // ============================================================================ // Integration Tests: Tool Execution with FileSystem // ============================================================================