Add idle-session TTL and cap to bound session memory growth

.env.example

@@ -23,3 +23,11 @@ BOOKSTACK_ENABLE_WRITE=false
 # REDIS_URL=redis://:password@redis:6379
 # REDIS_KEY_PREFIX=bookstack-mcp
 
+# HTTP transport only: idle timeout (ms) before an abandoned session is swept.
+# Default 1800000 (30m) — well above the ~5m SSE reconnect cycle seen in normal use.
+# BOOKSTACK_SESSION_IDLE_TTL_MS=1800000
+
+# HTTP transport only: hard cap on concurrent sessions, as a backstop independent
+# of the idle TTL sweep. Default 1000.
+# BOOKSTACK_MAX_SESSIONS=1000
+

src/index.ts

@@ -20,6 +20,7 @@ import {
   sendUnauthorized,
   OAuthConfig
 } from "./oauth/entra-proxy.js";
+import { SessionRegistry } from "./session-registry.js";
 
 // App-level config: the read-only credential is always present; the write credential and
 // OAuth proxy are optional. In OAuth mode the per-session credential is chosen by role.
@@ -1032,9 +1033,15 @@ async function startHttp(config: AppConfig): Promise<void> {
     return allowedHosts.includes(hostname) ? null : `Host '${hostname}' not allowed`;
   };
 
-  const transports: Record<string, StreamableHTTPServerTransport> = {};
-  // Per-session auth binding (OAuth mode only): subject + resolved write status, fixed at init.
-  const sessionAuth: Record<string, { sub?: string; isWriter: boolean }> = {};
+  const idleTtlRaw = process.env.BOOKSTACK_SESSION_IDLE_TTL_MS;
+  const idleTtlParsed = idleTtlRaw ? parseInt(idleTtlRaw, 10) : NaN;
+  const idleTtlMs = Number.isFinite(idleTtlParsed) && idleTtlParsed > 0 ? idleTtlParsed : undefined;
+
+  const maxSessionsRaw = process.env.BOOKSTACK_MAX_SESSIONS;
+  const maxSessionsParsed = maxSessionsRaw ? parseInt(maxSessionsRaw, 10) : NaN;
+  const maxSessions = Number.isFinite(maxSessionsParsed) && maxSessionsParsed > 0 ? maxSessionsParsed : undefined;
+
+  const sessions = new SessionRegistry({ idleTtlMs, maxSessions });
 
   const readJsonBody = (req: IncomingMessage): Promise<unknown> =>
     new Promise((resolve, reject) => {
@@ -1090,18 +1097,28 @@ async function startHttp(config: AppConfig): Promise<void> {
           parsedBody = await readJsonBody(req);
         }
 
-        if (sessionId && transports[sessionId]) {
+        if (sessionId && sessions.has(sessionId)) {
           // Reject a valid token belonging to a different subject than the session was bound to.
-          if (config.oauth && sessionAuth[sessionId] && auth?.sub !== sessionAuth[sessionId].sub) {
+          const boundAuth = sessions.getAuth(sessionId);
+          if (config.oauth && boundAuth && auth?.sub !== boundAuth.sub) {
             sendJson(res, 403, {
               jsonrpc: "2.0",
               error: { code: -32003, message: "Forbidden: token does not match session" },
               id: null
             });
             return;
           }
-          transport = transports[sessionId];
+          transport = sessions.get(sessionId);
         } else if (!sessionId && req.method === "POST" && isInitializeRequest(parsedBody)) {
+          if (sessions.atCapacity()) {
+            sendJson(res, 503, {
+              jsonrpc: "2.0",
+              error: { code: -32000, message: "Server busy: maximum concurrent sessions reached" },
+              id: null
+            });
+            return;
+          }
+
           // Choose the per-session credential: writers (with a write token configured) get the
           // write token and write tools; everyone else gets the read-only token.
           const useWrite = !!(auth?.isWriter && config.write);
@@ -1110,14 +1127,12 @@ async function startHttp(config: AppConfig): Promise<void> {
           transport = new StreamableHTTPServerTransport({
             sessionIdGenerator: () => randomUUID(),
             onsessioninitialized: (sid) => {
-              transports[sid] = transport!;
-              if (config.oauth && auth) sessionAuth[sid] = auth;
+              sessions.register(sid, transport!, config.oauth && auth ? auth : undefined);
             }
           });
           transport.onclose = () => {
             const sid = transport!.sessionId;
-            if (sid && transports[sid]) delete transports[sid];
-            if (sid && sessionAuth[sid]) delete sessionAuth[sid];
+            if (sid) sessions.delete(sid);
           };
           const server = buildServer(sessionConfig);
           await server.connect(transport);
@@ -1136,7 +1151,13 @@ async function startHttp(config: AppConfig): Promise<void> {
 
       // Health check
       if (pathname === "/health" || pathname === "/") {
-        sendJson(res, 200, { status: "ok", server: "bookstack-mcp", transport: "http" });
+        sendJson(res, 200, {
+          status: "ok",
+          server: "bookstack-mcp",
+          transport: "http",
+          sessions: sessions.size(),
+          memoryRssBytes: process.memoryUsage().rss
+        });
         return;
       }
 
@@ -1161,14 +1182,18 @@ async function startHttp(config: AppConfig): Promise<void> {
         ? `  Allowed Host headers: ${allowedHosts.join(", ")}`
         : `  Host header validation: DISABLED (set MCP_HTTP_ALLOWED_HOSTS to enable)`
     );
+    console.error(
+      `  Session limits: idle TTL ${Math.round(sessions.idleTtlMs / 60000)}m, max concurrent ${sessions.maxSessions}`
+    );
   });
 
   const shutdown = async () => {
     console.error("Shutting down HTTP server...");
-    for (const sid of Object.keys(transports)) {
-      try { await transports[sid].close(); } catch {}
-      delete transports[sid];
+    for (const sid of sessions.sessionIds()) {
+      try { await sessions.get(sid)?.close(); } catch {}
+      sessions.delete(sid);
     }
+    sessions.dispose();
     httpServer.close(() => process.exit(0));
   };
   process.on("SIGINT", shutdown);

src/session-registry.test.ts

@@ -0,0 +1,116 @@
+import test, { mock } from 'node:test';
+import assert from 'node:assert/strict';
+import type { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import { SessionRegistry } from './session-registry.js';
+
+// A stand-in for StreamableHTTPServerTransport: registry only ever reads/stores
+// the reference (plus, now, calls close() on eviction), so a minimal object with
+// a close() spy satisfies the type for this test.
+function fakeTransport() {
+  const close = mock.fn(async () => {});
+  const transport = { close } as unknown as StreamableHTTPServerTransport;
+  return { transport, close };
+}
+
+test('an abandoned session (onclose never fires) is swept after the idle TTL', (t) => {
+  t.mock.timers.enable({ apis: ['setInterval', 'Date'] });
+
+  const sessions = new SessionRegistry({ idleTtlMs: 30 * 60 * 1000, sweepIntervalMs: 5 * 60 * 1000 });
+  const { transport, close } = fakeTransport();
+  sessions.register('abandoned-session', transport, { sub: 'user-1', isWriter: false });
+  assert.equal(sessions.has('abandoned-session'), true, 'sanity check: session was registered');
+
+  // No `sessions.delete()` is ever called (that only happens from transport.onclose or
+  // process shutdown) — advance the clock past the idle TTL and let the sweep tick run.
+  t.mock.timers.tick(31 * 60 * 1000);
+
+  assert.equal(
+    sessions.has('abandoned-session'),
+    false,
+    'expected the abandoned session to be swept after the idle TTL'
+  );
+  assert.equal(
+    close.mock.callCount(),
+    1,
+    'expected the sweep to close() the transport before evicting it from bookkeeping'
+  );
+
+  sessions.dispose();
+});
+
+test('an expired session accessed between sweeps is closed by the lazy purge on access', (t) => {
+  t.mock.timers.enable({ apis: ['setInterval', 'Date'] });
+
+  // Sweep interval far longer than the TTL so only the lazy purgeIfExpired() check
+  // inside has() can evict this session before the next sweep tick would.
+  const sessions = new SessionRegistry({ idleTtlMs: 30 * 60 * 1000, sweepIntervalMs: 24 * 60 * 60 * 1000 });
+  const { transport, close } = fakeTransport();
+  sessions.register('stale-session', transport);
+
+  t.mock.timers.tick(31 * 60 * 1000);
+
+  assert.equal(
+    sessions.has('stale-session'),
+    false,
+    'expected the lazy purge to evict the expired session on access'
+  );
+  assert.equal(
+    close.mock.callCount(),
+    1,
+    'expected the lazy purge to close() the transport before evicting it from bookkeeping'
+  );
+
+  sessions.dispose();
+});
+
+test('a session touched within the idle TTL survives repeated sweep ticks', (t) => {
+  t.mock.timers.enable({ apis: ['setInterval', 'Date'] });
+
+  const sessions = new SessionRegistry({ idleTtlMs: 30 * 60 * 1000, sweepIntervalMs: 5 * 60 * 1000 });
+  const { transport, close } = fakeTransport();
+  sessions.register('active-session', transport);
+
+  // Simulate the ~5-minute SSE reconnect cycle seen in production: touch the session
+  // (get(), same as a real request) well inside the TTL window, across several cycles.
+  for (let i = 0; i < 5; i++) {
+    t.mock.timers.tick(5 * 60 * 1000);
+    assert.notEqual(
+      sessions.get('active-session'),
+      undefined,
+      `session should still be alive at minute ${(i + 1) * 5}`
+    );
+  }
+  assert.equal(close.mock.callCount(), 0, 'expected an active session to never be closed');
+
+  sessions.dispose();
+});
+
+test('registry enforces a hard cap on total sessions', () => {
+  const sessions = new SessionRegistry({ maxSessions: 2 });
+  try {
+    assert.equal(sessions.atCapacity(), false);
+    sessions.register('s1', fakeTransport().transport);
+    assert.equal(sessions.atCapacity(), false);
+    sessions.register('s2', fakeTransport().transport);
+    assert.equal(sessions.atCapacity(), true, 'expected at-capacity once maxSessions is reached');
+  } finally {
+    sessions.dispose();
+  }
+});
+
+test('constructor starts exactly one sweep timer, and dispose() clears it', (t) => {
+  const setIntervalSpy = t.mock.method(global, 'setInterval');
+  const clearIntervalSpy = t.mock.method(global, 'clearInterval');
+
+  const sessions = new SessionRegistry();
+  assert.equal(setIntervalSpy.mock.callCount(), 1, 'expected the constructor to start exactly one sweep timer');
+  const timerHandle = setIntervalSpy.mock.calls[0].result;
+
+  sessions.dispose();
+  assert.equal(clearIntervalSpy.mock.callCount(), 1, 'expected dispose() to clear the sweep timer');
+  assert.equal(
+    clearIntervalSpy.mock.calls[0].arguments[0],
+    timerHandle,
+    'expected dispose() to clear the same timer the constructor created'
+  );
+});

src/session-registry.ts

@@ -0,0 +1,119 @@
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+
+// Per-session auth binding (OAuth mode only): subject + resolved write status, fixed at init.
+export interface SessionAuthBinding {
+  sub?: string;
+  isWriter: boolean;
+}
+
+// Production traffic shows SSE streams recycling roughly every 5 minutes as normal,
+// healthy behavior (reconnect/backoff churn). 30 minutes gives 6x that margin, so a
+// normal reconnect cycle — even a slow or retried one — never gets mistaken for an
+// abandoned session, while still bounding how long a truly abandoned session's
+// transport (and the McpServer + tool closures it holds) can leak.
+const DEFAULT_IDLE_TTL_MS = 30 * 60 * 1000;
+
+// Sweep cadence mirrors oauth/kv-store.ts's InMemoryKvStore.
+const DEFAULT_SWEEP_INTERVAL_MS = 5 * 60 * 1000;
+
+// Hard cap as a second line of defense: bounds worst-case memory even if the TTL
+// logic has a bug, or a burst of sessions arrives faster than a sweep tick.
+const DEFAULT_MAX_SESSIONS = 1000;
+
+export interface SessionRegistryOptions {
+  idleTtlMs?: number;
+  sweepIntervalMs?: number;
+  maxSessions?: number;
+}
+
+/**
+ * Tracks live StreamableHTTPServerTransport sessions and their OAuth binding for the
+ * HTTP transport. Mirrors oauth/kv-store.ts's InMemoryKvStore: a periodic sweep
+ * (`setInterval`, `.unref()`'d) evicts sessions idle longer than `idleTtlMs`, with a
+ * lazy expiry check on access as a backstop between sweeps. `maxSessions` is a hard
+ * cap enforced independently of the TTL. `dispose()` stops the sweep timer.
+ */
+export class SessionRegistry {
+  private readonly transports: Record<string, StreamableHTTPServerTransport> = {};
+  private readonly sessionAuth: Record<string, SessionAuthBinding> = {};
+  private readonly lastSeen: Record<string, number> = {};
+  readonly idleTtlMs: number;
+  readonly maxSessions: number;
+  private sweep?: ReturnType<typeof setInterval>;
+
+  constructor(options: SessionRegistryOptions = {}) {
+    this.idleTtlMs = options.idleTtlMs ?? DEFAULT_IDLE_TTL_MS;
+    this.maxSessions = options.maxSessions ?? DEFAULT_MAX_SESSIONS;
+    const sweepIntervalMs = options.sweepIntervalMs ?? DEFAULT_SWEEP_INTERVAL_MS;
+    this.sweep = setInterval(() => this.sweepExpired(), sweepIntervalMs);
+    this.sweep.unref?.();
+  }
+
+  /** Closes the transport (best-effort) before dropping it from bookkeeping, mirroring shutdown(). */
+  private evict(sessionId: string): void {
+    this.transports[sessionId]?.close().catch(() => {});
+    this.delete(sessionId);
+  }
+
+  private sweepExpired(): void {
+    const now = Date.now();
+    for (const sid of Object.keys(this.transports)) {
+      if (now - (this.lastSeen[sid] ?? 0) >= this.idleTtlMs) this.evict(sid);
+    }
+  }
+
+  /** Evicts sessionId if idle-expired. Returns true if it was (or already is) absent. */
+  private purgeIfExpired(sessionId: string): boolean {
+    if (!(sessionId in this.transports)) return true;
+    const seen = this.lastSeen[sessionId];
+    if (seen === undefined || Date.now() - seen >= this.idleTtlMs) {
+      this.evict(sessionId);
+      return true;
+    }
+    return false;
+  }
+
+  has(sessionId: string): boolean {
+    return !this.purgeIfExpired(sessionId);
+  }
+
+  /** Also refreshes the session's idle timer — every real request is a keep-alive. */
+  get(sessionId: string): StreamableHTTPServerTransport | undefined {
+    if (this.purgeIfExpired(sessionId)) return undefined;
+    this.lastSeen[sessionId] = Date.now();
+    return this.transports[sessionId];
+  }
+
+  getAuth(sessionId: string): SessionAuthBinding | undefined {
+    if (this.purgeIfExpired(sessionId)) return undefined;
+    return this.sessionAuth[sessionId];
+  }
+
+  atCapacity(): boolean {
+    return this.size() >= this.maxSessions;
+  }
+
+  register(sessionId: string, transport: StreamableHTTPServerTransport, auth?: SessionAuthBinding): void {
+    this.transports[sessionId] = transport;
+    this.lastSeen[sessionId] = Date.now();
+    if (auth) this.sessionAuth[sessionId] = auth;
+  }
+
+  delete(sessionId: string): void {
+    delete this.transports[sessionId];
+    delete this.sessionAuth[sessionId];
+    delete this.lastSeen[sessionId];
+  }
+
+  sessionIds(): string[] {
+    return Object.keys(this.transports);
+  }
+
+  size(): number {
+    return this.sessionIds().length;
+  }
+
+  dispose(): void {
+    if (this.sweep) clearInterval(this.sweep);
+  }
+}