diff --git a/src/http/session-registry.test.ts b/src/http/session-registry.test.ts new file mode 100644 index 0000000..b6d0895 --- /dev/null +++ b/src/http/session-registry.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it, vi } from 'vitest'; + +import SessionRegistry from './session-registry.js'; + +function transport() { + return { close: vi.fn(async () => undefined) }; +} + +describe('SessionRegistry', () => { + it('expires idle sessions but preserves active requests', async () => { + let now = 1_000; + const registry = new SessionRegistry({ idleTtlMs: 100, maxSessions: 4, now: () => now }); + const idle = transport(); + const active = transport(); + + registry.add('idle', idle); + registry.add('active', active); + expect(registry.acquire('active')).toBe(active); + + now = 1_101; + expect(await registry.pruneIdle()).toBe(1); + expect(idle.close).toHaveBeenCalledOnce(); + expect(active.close).not.toHaveBeenCalled(); + expect(registry.size).toBe(1); + + registry.release('active'); + now = 1_202; + expect(await registry.pruneIdle()).toBe(1); + expect(active.close).toHaveBeenCalledOnce(); + expect(registry.size).toBe(0); + }); + + it('evicts the least-recently used inactive session at capacity', async () => { + let now = 1_000; + const registry = new SessionRegistry({ idleTtlMs: 10_000, maxSessions: 2, now: () => now }); + const oldest = transport(); + const newest = transport(); + + registry.add('oldest', oldest); + now = 1_010; + registry.add('newest', newest); + + expect(await registry.ensureCapacity()).toBe(true); + expect(oldest.close).toHaveBeenCalledOnce(); + expect(newest.close).not.toHaveBeenCalled(); + expect(registry.size).toBe(1); + }); + + it('refuses new capacity rather than closing active sessions', async () => { + const registry = new SessionRegistry({ idleTtlMs: 10_000, maxSessions: 1 }); + const active = transport(); + + registry.add('active', active); + registry.acquire('active'); + + expect(await registry.ensureCapacity()).toBe(false); + expect(active.close).not.toHaveBeenCalled(); + expect(registry.size).toBe(1); + }); + + it('does not let a late close callback remove a replacement session', () => { + const registry = new SessionRegistry({ idleTtlMs: 10_000, maxSessions: 2 }); + const first = transport(); + const replacement = transport(); + + registry.add('same-id', first); + registry.add('same-id', replacement); + + expect(registry.remove('same-id', first)).toBe(false); + expect(registry.size).toBe(1); + expect(registry.acquire('same-id')).toBe(replacement); + }); +}); diff --git a/src/http/session-registry.ts b/src/http/session-registry.ts new file mode 100644 index 0000000..aae2c0b --- /dev/null +++ b/src/http/session-registry.ts @@ -0,0 +1,120 @@ +export interface ClosableSession { + close(): Promise | void; +} + +interface SessionEntry { + transport: T; + lastAccess: number; + activeRequests: number; +} + +interface SessionRegistryOptions { + idleTtlMs: number; + maxSessions: number; + now?: () => number; +} + +/** + * Bounds long-lived Streamable HTTP sessions without interrupting active requests. + * Idle sessions are expired by TTL; when capacity is reached, the least-recently + * used inactive session is evicted before accepting a new one. + */ +export default class SessionRegistry { + private readonly sessions = new Map>(); + + private readonly idleTtlMs: number; + + private readonly maxSessions: number; + + private readonly now: () => number; + + constructor(options: SessionRegistryOptions) { + if (!Number.isFinite(options.idleTtlMs) || options.idleTtlMs <= 0) { + throw new Error('idleTtlMs must be a positive number'); + } + if (!Number.isInteger(options.maxSessions) || options.maxSessions <= 0) { + throw new Error('maxSessions must be a positive integer'); + } + + this.idleTtlMs = options.idleTtlMs; + this.maxSessions = options.maxSessions; + this.now = options.now ?? Date.now; + } + + get size(): number { + return this.sessions.size; + } + + acquire(sessionId: string): T | undefined { + const entry = this.sessions.get(sessionId); + if (!entry) return undefined; + + entry.activeRequests += 1; + entry.lastAccess = this.now(); + return entry.transport; + } + + add(sessionId: string, transport: T, activeRequests = 0): void { + this.sessions.set(sessionId, { + transport, + lastAccess: this.now(), + activeRequests, + }); + } + + release(sessionId: string): void { + const entry = this.sessions.get(sessionId); + if (!entry) return; + + entry.activeRequests = Math.max(0, entry.activeRequests - 1); + entry.lastAccess = this.now(); + } + + remove(sessionId: string, transport?: T): boolean { + const entry = this.sessions.get(sessionId); + if (!entry || (transport && entry.transport !== transport)) return false; + return this.sessions.delete(sessionId); + } + + async pruneIdle(): Promise { + const cutoff = this.now() - this.idleTtlMs; + const expired = [...this.sessions.entries()] + .filter(([, entry]) => entry.activeRequests === 0 && entry.lastAccess <= cutoff) + .map(([sessionId]) => sessionId); + + await Promise.allSettled(expired.map(async (sessionId) => this.closeSession(sessionId))); + return expired.length; + } + + async ensureCapacity(): Promise { + await this.pruneIdle(); + if (this.sessions.size < this.maxSessions) return true; + + const candidates = [...this.sessions.entries()] + .filter(([, entry]) => entry.activeRequests === 0) + .sort((a, b) => a[1].lastAccess - b[1].lastAccess); + + for (const [sessionId] of candidates) { + if (this.sessions.size < this.maxSessions) break; + await this.closeSession(sessionId); + } + + return this.sessions.size < this.maxSessions; + } + + async closeAll(): Promise { + const sessionIds = [...this.sessions.keys()]; + await Promise.allSettled( + sessionIds.map(async (sessionId) => this.closeSession(sessionId, true)), + ); + } + + private async closeSession(sessionId: string, includeActive = false): Promise { + const entry = this.sessions.get(sessionId); + if (!entry || (!includeActive && entry.activeRequests > 0)) return false; + + this.sessions.delete(sessionId); + await entry.transport.close(); + return true; + } +} diff --git a/src/main.ts b/src/main.ts index b57ece2..a1b8970 100644 --- a/src/main.ts +++ b/src/main.ts @@ -21,6 +21,7 @@ import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js'; import { loadConfig } from './config/loader.js'; import ConnectionManager from './connections/manager.js'; +import SessionRegistry from './http/session-registry.js'; import { bindServer, markInitialized, mcpLog } from './logging.js'; import registerAllPrompts from './prompts/register.js'; import registerAllResources from './resources/register.js'; @@ -234,7 +235,12 @@ async function runHttpServer(port: number): Promise { return server; } - const transports = new Map(); + const sessionIdleTtlMs = 30 * 60_000; + const maxHttpSessions = 32; + const sessionRegistry = new SessionRegistry({ + idleTtlMs: sessionIdleTtlMs, + maxSessions: maxHttpSessions, + }); const httpServer = createHttpServer(async (req: IncomingMessage, res: ServerResponse) => { if (req.url === '/health') { @@ -273,19 +279,32 @@ async function runHttpServer(port: number): Promise { const sessionId = Array.isArray(sessionIdHeader) ? sessionIdHeader[0] : sessionIdHeader; let transport: StreamableHTTPServerTransport; - const existing = sessionId ? transports.get(sessionId) : undefined; + const existing = sessionId ? sessionRegistry.acquire(sessionId) : undefined; + let acquiredSessionId = existing ? sessionId : undefined; if (existing) { transport = existing; } else if (!sessionId && req.method === 'POST' && isInitializeRequest(body)) { + if (!(await sessionRegistry.ensureCapacity())) { + res.writeHead(503, { 'Content-Type': 'application/json' }); + res.end( + JSON.stringify({ + jsonrpc: '2.0', + error: { code: -32000, message: 'Server busy: too many active MCP sessions' }, + id: null, + }), + ); + return; + } const newTransport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), onsessioninitialized: (sid) => { - transports.set(sid, newTransport); + sessionRegistry.add(sid, newTransport, 1); + acquiredSessionId = sid; }, }); newTransport.onclose = () => { const sid = newTransport.sessionId; - if (sid) transports.delete(sid); + if (sid) sessionRegistry.remove(sid, newTransport); }; const mcpServer = buildMcpSession(); await mcpServer.connect(newTransport); @@ -326,11 +345,20 @@ async function runHttpServer(port: number): Promise { return; } - await transport.handleRequest(req, res, body); + try { + await transport.handleRequest(req, res, body); + } finally { + if (acquiredSessionId) sessionRegistry.release(acquiredSessionId); + } }); await watcherService.start(); + const sessionCleanupInterval = setInterval(() => { + // eslint-disable-next-line no-void + void sessionRegistry.pruneIdle(); + }, 60_000); + const checkInterval = setInterval(async () => { try { await schedulerService.checkAndSend(); @@ -360,10 +388,10 @@ async function runHttpServer(port: number): Promise { const shutdown = async () => { clearInterval(checkInterval); + clearInterval(sessionCleanupInterval); hooksService.stop(); await watcherService.stop(); - // Close all transports concurrently; errors are ignored individually. - await Promise.allSettled(Array.from(transports.values(), async (t) => t.close())); + await sessionRegistry.closeAll(); await connections.closeAll(); httpServer.close(); };