diff --git a/CLAUDE.md b/CLAUDE.md index 082c791..7bc5b9e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -53,6 +53,7 @@ Digital Samba Embedded API MCP Server - a Model Context Protocol server for Digi - `DIGITAL_SAMBA_DEVELOPER_KEY`, `DIGITAL_SAMBA_API_URL`, `OAUTH_API_URL` - `OAUTH_CLIENT_ID`, `OAUTH_CLIENT_SECRET`, `OAUTH_AUTHORIZE_URL`, `OAUTH_TOKEN_URL`, `OAUTH_REDIRECT_URI`, `OAUTH_ISSUER` - `REDIS_URL` +- `SESSION_IDLE_TIMEOUT_MS` (default 30min; 0 disables), `SESSION_SWEEP_INTERVAL_MS` (default 5min) — idle HTTP session eviction ## Architecture @@ -64,6 +65,7 @@ src/ ├── oauth.ts # OAuth 2.0 / PKCE / DCR implementation ├── session-store.ts # Redis-backed session store (memory fallback) ├── auth.ts # AsyncLocalStorage API-key context +├── session-registry.ts # HTTP session tracking + idle sweep (clients rarely send DELETE) ├── cache.ts # Simple memory cache ├── logger.ts # Console logger (writes to stderr in stdio mode) ├── errors.ts # Error type definitions diff --git a/src/session-registry.ts b/src/session-registry.ts new file mode 100644 index 0000000..ef24595 --- /dev/null +++ b/src/session-registry.ts @@ -0,0 +1,135 @@ +/** + * Tracking of active HTTP transport sessions. + * + * Most MCP clients never send `DELETE /mcp`, and a transport's `onclose` only + * fires on an explicit close, so a client that simply goes away would otherwise + * leave its entry — and the Server instance behind it — in memory forever. + * This registry adds last-seen tracking and an idle sweep on top of the map. + * + * Sessions holding an open SSE stream are never swept, however long they have + * been quiet: streaming is activity. + * + * @module session-registry + */ + +import logger from "./logger.js"; + +/** The subset of a transport this registry needs. */ +export interface ClosableTransport { + close(): Promise | void; +} + +export interface SessionRegistryOptions { + /** Evict sessions idle for longer than this. 0 disables sweeping. */ + idleTimeoutMs?: number; +} + +export class SessionRegistry { + private readonly transports = new Map(); + private readonly lastActivity = new Map(); + private readonly openStreams = new Map(); + + readonly idleTimeoutMs: number; + + /** Cumulative count of swept sessions, surfaced via /health. */ + sweptCount = 0; + + constructor(options: SessionRegistryOptions = {}) { + this.idleTimeoutMs = options.idleTimeoutMs ?? 30 * 60 * 1000; + } + + get size(): number { + return this.transports.size; + } + + /** Number of sessions currently holding at least one open stream. */ + get streamingCount(): number { + return this.openStreams.size; + } + + has(sessionId: string): boolean { + return this.transports.has(sessionId); + } + + get(sessionId: string): T | undefined { + return this.transports.get(sessionId); + } + + values(): IterableIterator { + return this.transports.values(); + } + + entries(): IterableIterator<[string, T]> { + return this.transports.entries(); + } + + add(sessionId: string, transport: T, now: number = Date.now()): void { + this.transports.set(sessionId, transport); + this.lastActivity.set(sessionId, now); + } + + /** Record traffic on a session, resetting its idle clock. */ + touch(sessionId: string, now: number = Date.now()): void { + if (this.transports.has(sessionId)) { + this.lastActivity.set(sessionId, now); + } + } + + /** Drop a session and all its bookkeeping. Does not close the transport. */ + remove(sessionId: string): void { + this.transports.delete(sessionId); + this.lastActivity.delete(sessionId); + this.openStreams.delete(sessionId); + } + + /** Mark an SSE stream as opened; the session can't be swept while open. */ + openStream(sessionId: string, now: number = Date.now()): void { + this.openStreams.set(sessionId, (this.openStreams.get(sessionId) ?? 0) + 1); + this.touch(sessionId, now); + } + + /** Mark an SSE stream as closed; the idle clock restarts from now. */ + closeStream(sessionId: string, now: number = Date.now()): void { + const remaining = (this.openStreams.get(sessionId) ?? 1) - 1; + if (remaining > 0) { + this.openStreams.set(sessionId, remaining); + } else { + this.openStreams.delete(sessionId); + } + this.touch(sessionId, now); + } + + /** + * Close and evict every session idle beyond the timeout. + * + * @returns the number of sessions swept + */ + async sweep(now: number = Date.now()): Promise { + if (this.idleTimeoutMs <= 0) return 0; + + let swept = 0; + for (const [id, transport] of [...this.transports]) { + if ((this.openStreams.get(id) ?? 0) > 0) continue; // streaming = active + + const seen = this.lastActivity.get(id) ?? 0; + if (now - seen < this.idleTimeoutMs) continue; + + // Remove first: close() triggers onclose, and a throwing close() must not + // leave the entry behind — that would defeat the whole point of the sweep. + this.remove(id); + try { + await transport.close(); + } catch (err: any) { + logger.warn(`Error closing idle session ${id}: ${err?.message}`); + } + swept++; + logger.info(`Swept idle session: ${id}`); + } + + if (swept > 0) { + this.sweptCount += swept; + logger.info(`Swept ${swept} idle session(s), ${this.size} remaining`); + } + return swept; + } +} diff --git a/src/transports/http.ts b/src/transports/http.ts index 4fbf576..eeadccb 100644 --- a/src/transports/http.ts +++ b/src/transports/http.ts @@ -24,6 +24,7 @@ import { } from "../server.js"; import logger from "../logger.js"; import apiKeyContext from "../auth.js"; +import { SessionRegistry } from "../session-registry.js"; import { loadOAuthConfig, generateState, @@ -55,8 +56,18 @@ export interface HttpTransportConfig { requireAuth?: boolean; } -// Active transport sessions -const transports: Map = new Map(); +// 0 disables the idle sweep entirely. +const SESSION_IDLE_TIMEOUT_MS = Number( + process.env.SESSION_IDLE_TIMEOUT_MS ?? 30 * 60 * 1000, +); +const SESSION_SWEEP_INTERVAL_MS = Number( + process.env.SESSION_SWEEP_INTERVAL_MS ?? 5 * 60 * 1000, +); + +// Active transport sessions, with idle tracking (see session-registry.ts). +const transports = new SessionRegistry({ + idleTimeoutMs: SESSION_IDLE_TIMEOUT_MS, +}); /** * Authentication middleware @@ -212,6 +223,8 @@ export async function startHttpServer( commit: GIT_COMMIT, transport: "http", activeSessions: transports.size, + streamingSessions: transports.streamingCount, + sweptSessions: transports.sweptCount, oauthSessions: await getActiveSessionCount(), registeredClients: await getRegisteredClientCount(), }); @@ -563,6 +576,7 @@ export async function startHttpServer( if (mcpSessionId && transports.has(mcpSessionId)) { // Reuse existing session transport = transports.get(mcpSessionId)!; + transports.touch(mcpSessionId); logger.debug(`Reusing session: ${mcpSessionId}`); } else if (!mcpSessionId && isInitializeRequest(req.body)) { // New session initialization @@ -571,18 +585,18 @@ export async function startHttpServer( transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), onsessioninitialized: (id) => { - transports.set(id, transport); + transports.add(id, transport); logger.info(`Session initialized: ${id}`); }, onsessionclosed: (id) => { - transports.delete(id); + transports.remove(id); logger.info(`Session closed: ${id}`); }, }); transport.onclose = () => { if (transport.sessionId) { - transports.delete(transport.sessionId); + transports.remove(transport.sessionId); logger.debug( `Transport closed, session cleaned up: ${transport.sessionId}`, ); @@ -667,6 +681,14 @@ export async function startHttpServer( } const transport = transports.get(sessionId)!; + + // An SSE stream can stay open for hours with no other traffic. Count it as + // in-flight for the whole time it is open so the idle sweep can't evict a + // session that is actively streaming. + transports.openStream(sessionId); + // Start the idle clock from when the stream ends, not when it began. + res.on("close", () => transports.closeStream(sessionId)); + const apiKey = (req as any).apiKey; if (apiKey) { await apiKeyContext.run(apiKey, async () => { @@ -694,6 +716,7 @@ export async function startHttpServer( } const transport = transports.get(sessionId)!; + transports.touch(sessionId); const apiKey = (req as any).apiKey; if (apiKey) { await apiKeyContext.run(apiKey, async () => { @@ -1333,6 +1356,20 @@ export async function startHttpServer( }); app.delete("/", handleMcpDelete); + // Evict sessions abandoned without a DELETE (most clients never send one). + const sweepTimer = + SESSION_IDLE_TIMEOUT_MS > 0 + ? setInterval(() => { + transports + .sweep() + .catch((err) => + logger.error(`Idle session sweep failed: ${err?.message}`), + ); + }, SESSION_SWEEP_INTERVAL_MS) + : null; + // Don't hold the event loop open on account of the sweep. + sweepTimer?.unref(); + // Start listening app.listen(port, host, () => { logger.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"); @@ -1342,13 +1379,21 @@ export async function startHttpServer( logger.info(`Listening: http://${host}:${port}`); logger.info(`MCP Endpoint: http://${host}:${port}/mcp`); logger.info(`Auth: ${requireAuth ? "Required" : "Optional"}`); + logger.info( + `Idle session sweep: ${ + SESSION_IDLE_TIMEOUT_MS > 0 + ? `every ${SESSION_SWEEP_INTERVAL_MS / 1000}s, timeout ${SESSION_IDLE_TIMEOUT_MS / 60000}min` + : "disabled" + }`, + ); logger.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"); }); // Graceful shutdown process.on("SIGINT", async () => { logger.info("Shutting down HTTP server..."); - for (const [id, transport] of transports) { + if (sweepTimer) clearInterval(sweepTimer); + for (const [id, transport] of transports.entries()) { logger.debug(`Closing session: ${id}`); await transport.close(); } @@ -1357,6 +1402,7 @@ export async function startHttpServer( process.on("SIGTERM", async () => { logger.info("Received SIGTERM, shutting down..."); + if (sweepTimer) clearInterval(sweepTimer); for (const transport of transports.values()) { await transport.close(); } diff --git a/tests/unit/session-registry.test.ts b/tests/unit/session-registry.test.ts new file mode 100644 index 0000000..e559e1e --- /dev/null +++ b/tests/unit/session-registry.test.ts @@ -0,0 +1,236 @@ +/** + * Unit tests for session-registry.ts + * + * The registry exists because most MCP clients never send DELETE /mcp, so + * abandoned sessions used to accumulate in memory forever. These tests pin the + * eviction behaviour, and in particular the cases where a session must NOT be + * evicted (recent traffic, open SSE stream). + * + * Time is passed in explicitly rather than faked, so nothing here depends on + * timers or wall-clock. + * + * @module tests/unit/session-registry + */ + +import { SessionRegistry } from "../../src/session-registry.js"; + +const MINUTE = 60 * 1000; + +/** Minimal transport double; records whether close() was called. */ +function makeTransport(): { close: jest.Mock; closed: () => boolean } { + const close = jest.fn().mockResolvedValue(undefined); + return { close, closed: () => close.mock.calls.length > 0 }; +} + +describe("SessionRegistry", () => { + describe("basic bookkeeping", () => { + it("adds, retrieves and removes sessions", () => { + const registry = new SessionRegistry(); + const transport = makeTransport(); + + registry.add("s1", transport); + expect(registry.has("s1")).toBe(true); + expect(registry.get("s1")).toBe(transport); + expect(registry.size).toBe(1); + + registry.remove("s1"); + expect(registry.has("s1")).toBe(false); + expect(registry.size).toBe(0); + }); + + it("does not close the transport on remove()", () => { + const registry = new SessionRegistry(); + const transport = makeTransport(); + + registry.add("s1", transport); + registry.remove("s1"); + + // remove() is bookkeeping only - it is called *from* onclose handlers, + // so closing here would recurse. + expect(transport.closed()).toBe(false); + }); + + it("ignores touch() for unknown sessions", () => { + const registry = new SessionRegistry(); + registry.touch("never-existed"); + expect(registry.size).toBe(0); + }); + }); + + describe("idle sweeping", () => { + it("evicts and closes a session idle beyond the timeout", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = makeTransport(); + registry.add("stale", transport, 0); + + const swept = await registry.sweep(31 * MINUTE); + + expect(swept).toBe(1); + expect(registry.has("stale")).toBe(false); + expect(registry.size).toBe(0); + expect(transport.closed()).toBe(true); + }); + + it("keeps a session that is idle but within the timeout", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = makeTransport(); + registry.add("fresh", transport, 0); + + const swept = await registry.sweep(29 * MINUTE); + + expect(swept).toBe(0); + expect(registry.has("fresh")).toBe(true); + expect(transport.closed()).toBe(false); + }); + + it("keeps a session whose idle clock was reset by traffic", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = makeTransport(); + registry.add("busy", transport, 0); + + // A request arrives at t=29min, well before the session goes stale. + registry.touch("busy", 29 * MINUTE); + + // At t=31min it is only 2 minutes idle, so it must survive. + expect(await registry.sweep(31 * MINUTE)).toBe(0); + expect(registry.has("busy")).toBe(true); + }); + + it("sweeps only the stale sessions, leaving active ones alone", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const stale = makeTransport(); + const active = makeTransport(); + + registry.add("stale", stale, 0); + registry.add("active", active, 0); + registry.touch("active", 30 * MINUTE); + + const swept = await registry.sweep(31 * MINUTE); + + expect(swept).toBe(1); + expect(registry.has("stale")).toBe(false); + expect(registry.has("active")).toBe(true); + expect(stale.closed()).toBe(true); + expect(active.closed()).toBe(false); + }); + + it("accumulates sweptCount across sweeps", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + registry.add("a", makeTransport(), 0); + registry.add("b", makeTransport(), 0); + + await registry.sweep(31 * MINUTE); + expect(registry.sweptCount).toBe(2); + + registry.add("c", makeTransport(), 31 * MINUTE); + await registry.sweep(62 * MINUTE); + expect(registry.sweptCount).toBe(3); + }); + + it("is disabled when idleTimeoutMs is 0", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 0 }); + const transport = makeTransport(); + registry.add("ancient", transport, 0); + + const swept = await registry.sweep(1000 * MINUTE); + + expect(swept).toBe(0); + expect(registry.has("ancient")).toBe(true); + expect(transport.closed()).toBe(false); + }); + + it("evicts the session even if close() rejects", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = { + close: jest.fn().mockRejectedValue(new Error("socket already gone")), + }; + registry.add("broken", transport, 0); + + const swept = await registry.sweep(31 * MINUTE); + + // A failing close() must not leave the entry behind - that is the exact + // leak this registry exists to prevent. + expect(swept).toBe(1); + expect(registry.has("broken")).toBe(false); + expect(registry.size).toBe(0); + }); + }); + + describe("open SSE streams", () => { + it("never sweeps a session with an open stream, however old", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = makeTransport(); + registry.add("streaming", transport, 0); + registry.openStream("streaming", 0); + + const swept = await registry.sweep(600 * MINUTE); + + expect(swept).toBe(0); + expect(registry.has("streaming")).toBe(true); + expect(transport.closed()).toBe(false); + }); + + it("restarts the idle clock when the stream closes", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + const transport = makeTransport(); + registry.add("streamed", transport, 0); + + registry.openStream("streamed", 0); + registry.closeStream("streamed", 100 * MINUTE); + + // Idle is measured from stream close (100min), not session start. + expect(await registry.sweep(120 * MINUTE)).toBe(0); + expect(registry.has("streamed")).toBe(true); + + expect(await registry.sweep(131 * MINUTE)).toBe(1); + expect(registry.has("streamed")).toBe(false); + }); + + it("stays protected until the last of several streams closes", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + registry.add("multi", makeTransport(), 0); + + registry.openStream("multi", 0); + registry.openStream("multi", 0); + expect(registry.streamingCount).toBe(1); + + // One stream ends; the other is still open, so it must survive. + registry.closeStream("multi", 0); + expect(await registry.sweep(600 * MINUTE)).toBe(0); + expect(registry.has("multi")).toBe(true); + + // Now the last one ends and the idle clock applies again. + registry.closeStream("multi", 600 * MINUTE); + expect(registry.streamingCount).toBe(0); + expect(await registry.sweep(631 * MINUTE)).toBe(1); + }); + + it("clears stream state when the session is removed", () => { + const registry = new SessionRegistry(); + registry.add("s1", makeTransport()); + registry.openStream("s1"); + expect(registry.streamingCount).toBe(1); + + registry.remove("s1"); + expect(registry.streamingCount).toBe(0); + }); + }); + + describe("the leak this prevents", () => { + it("drains sessions abandoned without DELETE", async () => { + const registry = new SessionRegistry({ idleTimeoutMs: 30 * MINUTE }); + + // 449 sessions initialized and never closed - the state observed on the + // dev host, where every session leaked because clients never send DELETE. + for (let i = 0; i < 449; i++) { + registry.add(`abandoned-${i}`, makeTransport(), 0); + } + expect(registry.size).toBe(449); + + await registry.sweep(31 * MINUTE); + + expect(registry.size).toBe(0); + expect(registry.sweptCount).toBe(449); + }); + }); +});