From a5b6b67a9d072bb866413cd8b45e9e084e6a8474 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20K=2E?= Date: Mon, 31 Aug 2026 05:14:39 +0200 Subject: [PATCH] fix: bound retained MCP sessions --- docs/configuration.md | 6 + schema/v1/devspace.schema.json | 6 + src/config-schema.ts | 1 + src/config.test.ts | 3 + src/config.ts | 3 + src/mcp-sessions.test.ts | 274 +++++++++++++++++++++++++++++++-- src/mcp-sessions.ts | 235 ++++++++++++++++++++++++++-- src/server.test.ts | 190 ++++++++++++++++++++++- src/server.ts | 106 ++++++++++++- 9 files changed, 794 insertions(+), 30 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 6a6f607bd..f92d072ad 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -25,6 +25,7 @@ Run `devspace init` to create both files. `devspace config set publicBaseUrl "server": { "host": "127.0.0.1", "port": 7676, + "maxMcpSessions": 256, // Use the public origin only; do not append /mcp. "publicBaseUrl": "https://devspace.example.com", "allowedHosts": [], @@ -77,6 +78,11 @@ Omitted sections and keys use the defaults shown above. An empty `workspaces.allowedRoots` uses the current working directory. Unknown keys are rejected so spelling mistakes cannot silently alter behavior. +`server.maxMcpSessions` limits how many stateful MCP sessions DevSpace keeps in +memory. When the limit is reached, DevSpace closes the least-recently-used idle +session first. Active sessions are not evicted. If all retained sessions are +active, new initialize requests return `503` until a slot is free. + ## Tool modes and UI `tools.mode` accepts two values: diff --git a/schema/v1/devspace.schema.json b/schema/v1/devspace.schema.json index e7c18466e..61b1d428f 100644 --- a/schema/v1/devspace.schema.json +++ b/schema/v1/devspace.schema.json @@ -29,6 +29,12 @@ "minimum": 1, "maximum": 65535 }, + "maxMcpSessions": { + "default": 256, + "type": "integer", + "exclusiveMinimum": 0, + "maximum": 9007199254740991 + }, "publicBaseUrl": { "default": null, "anyOf": [ diff --git a/src/config-schema.ts b/src/config-schema.ts index c30bb3612..9e07c57fc 100644 --- a/src/config-schema.ts +++ b/src/config-schema.ts @@ -8,6 +8,7 @@ export const DEVSPACE_CONFIG_SCHEMA_URL = const serverConfigSchema = z.object({ host: z.string().trim().min(1).default("127.0.0.1"), port: z.number().int().min(1).max(65_535).default(7676), + maxMcpSessions: z.number().int().positive().default(256), publicBaseUrl: z.string().url().nullable().default(null), allowedHosts: z.array(z.string().trim().min(1)).default([]), trustProxy: z.boolean().default(false), diff --git a/src/config.test.ts b/src/config.test.ts index 47a39652a..8ab11d681 100644 --- a/src/config.test.ts +++ b/src/config.test.ts @@ -15,6 +15,7 @@ try { const defaults = loadConfig(env); assert.equal(defaults.host, "127.0.0.1"); assert.equal(defaults.port, 7676); + assert.equal(defaults.maxMcpSessions, 256); assert.equal(defaults.publicBaseUrl, "http://127.0.0.1:7676"); assert.deepEqual(defaults.allowedRoots, [process.cwd()]); assert.deepEqual(defaults.allowedHosts, ["localhost", "127.0.0.1", "::1"]); @@ -38,6 +39,7 @@ try { server: { host: "0.0.0.0", port: 8787, + maxMcpSessions: 32, publicBaseUrl: "https://devspace.example.com/", allowedHosts: ["example.internal"], trustProxy: true, @@ -76,6 +78,7 @@ try { assert.equal(configured.configDir, configDir); assert.equal(configured.host, "0.0.0.0"); assert.equal(configured.port, 8787); + assert.equal(configured.maxMcpSessions, 32); assert.equal(configured.publicBaseUrl, "https://devspace.example.com"); assert.deepEqual(configured.allowedRoots, [resolve("~/work".replace("~", process.env.HOME!))]); assert.deepEqual(configured.allowedHosts, [ diff --git a/src/config.ts b/src/config.ts index e53305268..c0eefc33b 100644 --- a/src/config.ts +++ b/src/config.ts @@ -12,6 +12,7 @@ export interface ServerConfig { configDir: string; host: string; port: number; + maxMcpSessions: number; oauth: OAuthConfig; allowedRoots: string[]; allowedHosts: string[]; @@ -31,6 +32,7 @@ export interface ServerConfig { logging: LoggingConfig; } +/** Load and normalize the persisted DevSpace server configuration. */ export function loadConfig(env: NodeJS.ProcessEnv = process.env): ServerConfig { const files = loadDevspaceFiles(env); const stored = files.config; @@ -52,6 +54,7 @@ export function loadConfig(env: NodeJS.ProcessEnv = process.env): ServerConfig { configDir: files.dir, host, port, + maxMcpSessions: stored.server.maxMcpSessions, oauth: { ownerToken: parseRequiredSecret( env.DEVSPACE_OAUTH_OWNER_TOKEN ?? files.auth.ownerToken, diff --git a/src/mcp-sessions.test.ts b/src/mcp-sessions.test.ts index 1b3ecf349..8ad83314d 100644 --- a/src/mcp-sessions.test.ts +++ b/src/mcp-sessions.test.ts @@ -16,16 +16,29 @@ function createTransport(closeError?: Error): FakeTransport { }; } +/** Reserve a slot and commit a fake transport for registry tests. */ +async function reserveAndCommit( + registry: McpSessionRegistry, + sessionId: string, + transport: FakeTransport, + options: { active?: boolean } = {}, +): Promise { + const result = await registry.reserve(); + assert.equal(result.ok, true); + if (!result.ok) return; + assert.equal(await result.reservation.commit(sessionId, transport, options), true); +} + let now = 0; -const registry = new McpSessionRegistry({ now: () => now }); +const registry = new McpSessionRegistry({ now: () => now, maxSessions: 4 }); const staleTransport = createTransport(); const activeTransport = createTransport(); -registry.register("stale", staleTransport); +await reserveAndCommit(registry, "stale", staleTransport); now = 1_000; -registry.register("active", activeTransport); +await reserveAndCommit(registry, "active", activeTransport); now = 1_500; -assert.equal(registry.get("active"), activeTransport); +assert.equal(registry.acquire("active"), activeTransport); now = 2_000; const idleResults = await registry.closeIdle(1_500); @@ -33,12 +46,14 @@ assert.deepEqual(idleResults, [{ sessionId: "stale" }]); assert.equal(staleTransport.closeCalls, 1); assert.equal(activeTransport.closeCalls, 0); assert.equal(registry.size, 1); -assert.equal(registry.get("stale"), undefined); -assert.equal(registry.get("active"), activeTransport); +assert.equal(registry.acquire("stale"), undefined); +assert.equal(registry.release("active"), true); +assert.equal(registry.release("active"), false); +assert.equal(registry.activeSessions, 0); const closeError = new Error("close failed"); const failingTransport = createTransport(closeError); -registry.register("failing", failingTransport); +await reserveAndCommit(registry, "failing", failingTransport); now = 10_000; const failingResults = await registry.closeIdle(1); @@ -46,12 +61,15 @@ assert.equal(failingResults.length, 2); assert.deepEqual(failingResults.map((result) => result.sessionId).sort(), ["active", "failing"]); assert.equal(failingResults.find((result) => result.sessionId === "failing")?.error, closeError); assert.equal(failingTransport.closeCalls, 1); -assert.equal(registry.size, 0); +assert.equal(registry.size, 1, "a failed close keeps its slot occupied"); +assert.equal(registry.acquire("failing"), failingTransport); +assert.equal(registry.release("failing"), true); +registry.remove("failing"); const first = createTransport(); const second = createTransport(); -registry.register("first", first); -registry.register("second", second); +await reserveAndCommit(registry, "first", first); +await reserveAndCommit(registry, "second", second); registry.remove("first"); const shutdownResults = await registry.closeAll(); @@ -60,6 +78,7 @@ assert.equal(first.closeCalls, 0); assert.equal(second.closeCalls, 1); assert.equal(registry.size, 0); +const delayedRegistry = new McpSessionRegistry(); let finishDelayedClose: (() => void) | undefined; let delayedCloseResolved = false; const delayedTransport: FakeTransport = { @@ -71,16 +90,243 @@ const delayedTransport: FakeTransport = { }); }, }; -registry.register("delayed", delayedTransport); -const delayedClose = registry.closeAll(); +await reserveAndCommit(delayedRegistry, "delayed", delayedTransport); +const delayedClose = delayedRegistry.closeAll(); void delayedClose.then(() => { delayedCloseResolved = true; }); -await Promise.resolve(); +await new Promise((resolve) => setImmediate(resolve)); assert.equal(delayedCloseResolved, false); assert.equal(delayedTransport.closeCalls, 1); finishDelayedClose?.(); await delayedClose; assert.equal(delayedCloseResolved, true); -assert.equal(registry.size, 0); +assert.equal(delayedRegistry.size, 0); + +const bounded = new McpSessionRegistry({ now: () => now, maxSessions: 2 }); +const oldest = createTransport(); +const protectedTransport = createTransport(); +const replacement = createTransport(); + +now = 20_000; +await reserveAndCommit(bounded, "oldest", oldest); +now = 21_000; +await reserveAndCommit(bounded, "protected", protectedTransport, { active: true }); +assert.equal(bounded.size, 2); +assert.equal(bounded.occupancy, 2); +assert.equal(bounded.activeSessions, 1); + +now = 22_000; +const replacementReservation = await bounded.reserve(); +assert.equal(replacementReservation.ok, true); +if (replacementReservation.ok) { + assert.equal(replacementReservation.reservation.evictedSessionId, "oldest"); + assert.equal(oldest.closeCalls, 1); + assert.equal(protectedTransport.closeCalls, 0); + assert.equal(await replacementReservation.reservation.commit("replacement", replacement), true); +} +assert.equal(bounded.size, 2); +assert.equal(bounded.occupancy, 2); +assert.equal(bounded.acquire("oldest"), undefined); + +assert.equal(bounded.acquire("replacement"), replacement); +assert.equal(bounded.activeSessions, 2); +const blocked = await bounded.reserve(); +assert.deepEqual(blocked, { ok: false, reason: "capacity_exhausted" }); +assert.equal(bounded.size, 2); +assert.equal(bounded.occupancy, 2); +assert.equal(protectedTransport.closeCalls, 0); +assert.equal(replacement.closeCalls, 0); +assert.equal(bounded.release("replacement"), true); +assert.equal(bounded.release("protected"), true); + +const closeFailure = new Error("cannot close candidate"); +const closeFailureRegistry = new McpSessionRegistry({ maxSessions: 1 }); +const stuck = createTransport(closeFailure); +await reserveAndCommit(closeFailureRegistry, "stuck", stuck); +const deniedByCloseFailure = await closeFailureRegistry.reserve(); +assert.equal(deniedByCloseFailure.ok, false); +if (!deniedByCloseFailure.ok) { + assert.equal(deniedByCloseFailure.reason, "close_failed"); + assert.equal(deniedByCloseFailure.sessionId, "stuck"); + assert.equal(deniedByCloseFailure.error, closeFailure); +} +assert.equal(closeFailureRegistry.size, 1); +assert.equal(closeFailureRegistry.occupancy, 1); +assert.equal(stuck.closeCalls, 1); + +const fallbackRegistry = new McpSessionRegistry({ now: () => now, maxSessions: 2 }); +const brokenCandidate = createTransport(new Error("broken candidate")); +const healthyCandidate = createTransport(); +now = 30_000; +await reserveAndCommit(fallbackRegistry, "broken", brokenCandidate); +now = 31_000; +await reserveAndCommit(fallbackRegistry, "healthy", healthyCandidate); +const fallbackReservation = await fallbackRegistry.reserve(); +assert.equal(fallbackReservation.ok, true); +if (fallbackReservation.ok) { + assert.equal(fallbackReservation.reservation.evictedSessionId, "healthy"); + assert.equal(await fallbackReservation.reservation.commit("replacement", createTransport()), true); +} +assert.equal(brokenCandidate.closeCalls, 1); +assert.equal(healthyCandidate.closeCalls, 1); +assert.equal(fallbackRegistry.size, 2); +assert.equal(fallbackRegistry.occupancy, 2); + +const concurrent = new McpSessionRegistry({ maxSessions: 2 }); +const base = createTransport(); +await reserveAndCommit(concurrent, "base", base); +const firstReservation = await concurrent.reserve(); +assert.equal(firstReservation.ok, true); +assert.equal(concurrent.occupancy, 2, "initialize reserves its slot before it runs"); +const secondReservation = await concurrent.reserve(); +assert.equal(secondReservation.ok, true); +assert.equal(base.closeCalls, 1, "a second initialize evicts the idle session to claim a slot"); +assert.equal(concurrent.occupancy, 2); +if (firstReservation.ok) { + assert.equal(await firstReservation.reservation.commit("new-1", createTransport(), { active: true }), true); +} +if (secondReservation.ok) { + assert.equal(await secondReservation.reservation.commit("new-2", createTransport(), { active: true }), true); +} +assert.equal(concurrent.size, 2); +assert.equal(concurrent.occupancy, 2); +assert.equal(concurrent.activeSessions, 2); +const allProtected = await concurrent.reserve(); +assert.deepEqual(allProtected, { ok: false, reason: "capacity_exhausted" }); + +const callbackRemovalRegistry = new McpSessionRegistry({ maxSessions: 2 }); +let finishCallbackClose: (() => void) | undefined; +let callbackCloseStarted!: () => void; +const callbackCloseStartedPromise = new Promise((resolve) => { + callbackCloseStarted = resolve; +}); +const callbackRemovingTransport: FakeTransport = { + closeCalls: 0, + close() { + this.closeCalls += 1; + callbackRemovalRegistry.remove("callback-removal"); + callbackCloseStarted(); + return new Promise((resolve) => { + finishCallbackClose = resolve; + }); + }, +}; +await reserveAndCommit(callbackRemovalRegistry, "callback-removal", callbackRemovingTransport); +await reserveAndCommit(callbackRemovalRegistry, "other", createTransport()); +const callbackFirstReservationPromise = callbackRemovalRegistry.reserve(); +await callbackCloseStartedPromise; +const callbackSecondReservationPromise = callbackRemovalRegistry.reserve(); +finishCallbackClose?.(); +const [callbackFirstReservation, callbackSecondReservation] = await Promise.all([ + callbackFirstReservationPromise, + callbackSecondReservationPromise, +]); +assert.equal(callbackFirstReservation.ok, true); +assert.equal(callbackSecondReservation.ok, true); +assert.equal( + callbackRemovalRegistry.occupancy, + 2, + "onclose removal must not make the slot look free during eviction", +); +if (callbackFirstReservation.ok) { + assert.equal( + await callbackFirstReservation.reservation.commit("callback-new-1", createTransport(), { active: true }), + true, + ); +} +if (callbackSecondReservation.ok) { + assert.equal( + await callbackSecondReservation.reservation.commit("callback-new-2", createTransport(), { active: true }), + true, + ); +} +assert.equal(callbackRemovalRegistry.size, 2); +assert.equal(callbackRemovalRegistry.occupancy, 2); + +const referenceCounted = new McpSessionRegistry({ maxSessions: 1 }); +const referenceTransport = createTransport(); +await reserveAndCommit(referenceCounted, "reference", referenceTransport); +assert.equal(referenceCounted.acquire("reference"), referenceTransport); +assert.equal(referenceCounted.acquire("reference"), referenceTransport); +assert.equal(referenceCounted.activeSessions, 1); +assert.equal(referenceCounted.release("reference"), true); +assert.equal(referenceCounted.activeSessions, 1); +assert.deepEqual(await referenceCounted.reserve(), { ok: false, reason: "capacity_exhausted" }); +assert.equal(referenceCounted.release("reference"), true); +assert.equal(referenceCounted.activeSessions, 0); + +const churnLimit = 256; +const churnRegistry = new McpSessionRegistry({ maxSessions: churnLimit }); +const churnTransports: FakeTransport[] = []; +for (let index = 0; index < 1_024; index += 1) { + const transport = createTransport(); + churnTransports.push(transport); + const reservation = await churnRegistry.reserve(); + assert.equal(reservation.ok, true); + if (!reservation.ok) continue; + assert.equal(await reservation.reservation.commit(`churn-${index}`, transport), true); + assert.ok(churnRegistry.size <= churnLimit); + assert.ok(churnRegistry.occupancy <= churnLimit); +} +assert.equal(churnRegistry.size, churnLimit); +assert.equal(churnRegistry.occupancy, churnLimit); +assert.equal( + churnTransports.slice(0, 1_024 - churnLimit).filter((transport) => transport.closeCalls === 1).length, + 1_024 - churnLimit, + "churn closes every evicted session", +); +assert.equal( + churnTransports.slice(1_024 - churnLimit).filter((transport) => transport.closeCalls !== 0).length, + 0, + "the newest sessions stay open", +); + +const releasedReservationRegistry = new McpSessionRegistry({ maxSessions: 1 }); +const pending = await releasedReservationRegistry.reserve(); +assert.equal(pending.ok, true); +assert.equal(releasedReservationRegistry.occupancy, 1); +if (pending.ok) pending.reservation.release(); +assert.equal(releasedReservationRegistry.occupancy, 0); +assert.equal(releasedReservationRegistry.size, 0); + +const shutdownRegistry = new McpSessionRegistry({ maxSessions: 1 }); +const lateReservation = await shutdownRegistry.reserve(); +assert.equal(lateReservation.ok, true); +assert.equal(shutdownRegistry.occupancy, 1); +let finishShutdownClose: (() => void) | undefined; +const lateTransport: FakeTransport = { + closeCalls: 0, + close() { + this.closeCalls += 1; + return new Promise((resolve) => { + finishShutdownClose = resolve; + }); + }, +}; +let shutdownResolved = false; +const shutdown = shutdownRegistry.closeAll().then((results) => { + shutdownResolved = true; + return results; +}); +await Promise.resolve(); +assert.equal(shutdownResolved, false); +if (lateReservation.ok) { + const lateCommit = lateReservation.reservation.commit("late", lateTransport); + await Promise.resolve(); + assert.equal(lateTransport.closeCalls, 1); + assert.equal(shutdownResolved, false); + finishShutdownClose?.(); + assert.equal(await lateCommit, false); +} +assert.deepEqual(await shutdown, []); +assert.equal(shutdownResolved, true); +assert.equal(shutdownRegistry.size, 0); +assert.equal(shutdownRegistry.occupancy, 0); +assert.deepEqual(await shutdownRegistry.reserve(), { ok: false, reason: "capacity_exhausted" }); + +assert.throws( + () => new McpSessionRegistry({ maxSessions: 0 }), + /positive integer/, +); diff --git a/src/mcp-sessions.ts b/src/mcp-sessions.ts index 3a844943a..1a0a17671 100644 --- a/src/mcp-sessions.ts +++ b/src/mcp-sessions.ts @@ -10,58 +10,239 @@ export interface McpSessionCloseResult { interface McpSessionEntry { transport: TTransport; lastActivityAt: number; + activeResponses: number; + closing: boolean; } export interface McpSessionRegistryOptions { now?: () => number; + maxSessions?: number; } +export interface McpSessionReservation { + readonly evictedSessionId?: string; + commit( + sessionId: string, + transport: TTransport, + options?: { active?: boolean }, + ): Promise; + release(): void; +} + +export type McpSessionReservationResult = + | { + ok: true; + reservation: McpSessionReservation; + } + | { + ok: false; + reason: "capacity_exhausted" | "close_failed"; + sessionId?: string; + error?: unknown; + }; + +const DEFAULT_MAX_SESSIONS = 256; + export class McpSessionRegistry { private readonly sessions = new Map>(); private readonly now: () => number; + private readonly maximumSessions: number; + private reservations = 0; + private closed = false; + private readonly reservationWaiters = new Set<() => void>(); + // Transport.close() can call onclose synchronously and remove its session. + // Serialize reservations so the callback cannot make a slot look free too early. + private reservationGate: Promise = Promise.resolve(); + /** Create a bounded session registry with an optional clock for tests. */ constructor(options: McpSessionRegistryOptions = {}) { this.now = options.now ?? Date.now; + this.maximumSessions = options.maxSessions ?? DEFAULT_MAX_SESSIONS; + if (!Number.isInteger(this.maximumSessions) || this.maximumSessions < 1) { + throw new TypeError("maxSessions must be a positive integer."); + } } + /** Number of retained sessions. */ get size(): number { return this.sessions.size; } - register(sessionId: string, transport: TTransport): void { - this.sessions.set(sessionId, { - transport, - lastActivityAt: this.now(), + /** Configured retained-session limit. */ + get maxSessions(): number { + return this.maximumSessions; + } + + /** Retained sessions plus pending initialization reservations. */ + get occupancy(): number { + return this.sessions.size + this.reservations; + } + + /** Number of sessions with at least one active HTTP response. */ + get activeSessions(): number { + let active = 0; + for (const entry of this.sessions.values()) { + if (entry.activeResponses > 0) active += 1; + } + return active; + } + + /** Reserve capacity for a session before its transport is initialized. */ + async reserve(): Promise> { + let unlock = (): void => undefined; + const previousReservation = this.reservationGate; + this.reservationGate = new Promise((resolve) => { + unlock = resolve; }); + + await previousReservation; + try { + if (this.closed) return { ok: false, reason: "capacity_exhausted" }; + return await this.reserveSerially(); + } finally { + unlock(); + } + } + + /** Evict an idle session when needed, then create one pending reservation. */ + private async reserveSerially(): Promise> { + let evictedSessionId: string | undefined; + + if (this.occupancy >= this.maximumSessions) { + const failedCandidates = new Set(); + let lastCloseFailure: { sessionId: string; error: unknown } | undefined; + + while (this.occupancy >= this.maximumSessions) { + const candidate = this.oldestEvictableSession(failedCandidates); + if (!candidate) { + return lastCloseFailure + ? { ok: false, reason: "close_failed", ...lastCloseFailure } + : { ok: false, reason: "capacity_exhausted" }; + } + + const [sessionId, entry] = candidate; + entry.closing = true; + try { + await entry.transport.close(); + } catch (error) { + if (this.sessions.get(sessionId) === entry) entry.closing = false; + failedCandidates.add(sessionId); + lastCloseFailure = { sessionId, error }; + continue; + } + this.sessions.delete(sessionId); + evictedSessionId = sessionId; + } + } + + if (this.closed) return { ok: false, reason: "capacity_exhausted" }; + + this.reservations += 1; + let settled = false; + return { + ok: true, + reservation: { + evictedSessionId, + commit: async (sessionId, transport, options = {}) => { + if (settled) throw new Error("MCP session reservation has already been settled."); + settled = true; + + if (this.closed) { + try { + await transport.close(); + } finally { + this.settleReservation(); + } + return false; + } + + this.settleReservation(); + if (this.sessions.size >= this.maximumSessions) { + throw new Error("MCP session reservation exceeded the configured limit."); + } + this.sessions.set(sessionId, { + transport, + lastActivityAt: this.now(), + activeResponses: options.active ? 1 : 0, + closing: false, + }); + return true; + }, + release: () => { + if (settled) return; + settled = true; + this.settleReservation(); + }, + }, + }; } - get(sessionId: string): TTransport | undefined { + /** Mark one HTTP response as active for a retained session. */ + acquire(sessionId: string): TTransport | undefined { const entry = this.sessions.get(sessionId); - if (!entry) return undefined; + if (!entry || entry.closing) return undefined; + entry.activeResponses += 1; entry.lastActivityAt = this.now(); return entry.transport; } + /** Release one active HTTP response from a retained session. */ + release(sessionId: string): boolean { + const entry = this.sessions.get(sessionId); + if (!entry || entry.activeResponses < 1) return false; + + entry.activeResponses -= 1; + entry.lastActivityAt = this.now(); + return true; + } + + /** Remove a session from the registry without closing its transport. */ remove(sessionId: string): boolean { return this.sessions.delete(sessionId); } + /** Close sessions that are idle and have no active HTTP responses. */ async closeIdle(idleTimeoutMs: number): Promise { const cutoff = this.now() - idleTimeoutMs; - const idleSessions: Array<{ sessionId: string; transport: TTransport }> = []; + const idleSessions: Array<{ + sessionId: string; + entry: McpSessionEntry; + }> = []; for (const [sessionId, entry] of this.sessions) { - if (entry.lastActivityAt > cutoff) continue; + if ( + entry.closing + || entry.activeResponses > 0 + || entry.lastActivityAt > cutoff + ) { + continue; + } - this.sessions.delete(sessionId); - idleSessions.push({ sessionId, transport: entry.transport }); + entry.closing = true; + idleSessions.push({ sessionId, entry }); } - return closeSessions(idleSessions); + return Promise.all( + idleSessions.map(async ({ sessionId, entry }) => { + try { + await entry.transport.close(); + this.sessions.delete(sessionId); + return { sessionId }; + } catch (error) { + if (this.sessions.get(sessionId) === entry) entry.closing = false; + return { sessionId, error }; + } + }), + ); } + /** Stop new reservations and close every retained or initializing session. */ async closeAll(): Promise { + this.closed = true; + await this.reservationGate; + await this.waitForReservations(); + const sessions = Array.from(this.sessions, ([sessionId, entry]) => ({ sessionId, transport: entry.transport, @@ -69,8 +250,40 @@ export class McpSessionRegistry { this.sessions.clear(); return closeSessions(sessions); } + + /** Find the oldest idle session that has not failed during this eviction attempt. */ + private oldestEvictableSession( + excluded: ReadonlySet = new Set(), + ): [string, McpSessionEntry] | undefined { + let candidate: [string, McpSessionEntry] | undefined; + for (const current of this.sessions) { + const [sessionId, entry] = current; + if (excluded.has(sessionId) || entry.closing || entry.activeResponses > 0) continue; + if (!candidate || entry.lastActivityAt < candidate[1].lastActivityAt) { + candidate = current; + } + } + return candidate; + } + + /** Decrement pending reservation count and wake shutdown when the last one settles. */ + private settleReservation(): void { + this.reservations -= 1; + if (this.reservations !== 0) return; + for (const resolve of this.reservationWaiters) resolve(); + this.reservationWaiters.clear(); + } + + /** Wait until all reservations have committed or been released. */ + private waitForReservations(): Promise { + if (this.reservations === 0) return Promise.resolve(); + return new Promise((resolve) => { + this.reservationWaiters.add(resolve); + }); + } } +/** Close a fixed set of transports and preserve per-session close errors. */ async function closeSessions( sessions: Array<{ sessionId: string; transport: TTransport }>, ): Promise { diff --git a/src/server.test.ts b/src/server.test.ts index 79c21a66c..741a6c126 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -1,19 +1,24 @@ import assert from "node:assert/strict"; import { execFile } from "node:child_process"; +import { once } from "node:events"; import { mkdtemp, mkdir, rm, writeFile } from "node:fs/promises"; +import { createServer as createHttpServer } from "node:http"; import { tmpdir } from "node:os"; import { join } from "node:path"; import test, { type TestContext } from "node:test"; import { promisify } from "node:util"; import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js"; +import type { Response as ExpressResponse } from "express"; import { loadConfig, type ServerConfig, type ToolMode } from "./config.js"; import type { LocalAgentProviderAvailability } from "./local-agent-availability.js"; import { buildLocalAgentProviderStatuses } from "./local-agent-catalog.js"; import type { SubagentsConfig } from "./local-agent-config.js"; import { createReviewCheckpointManager } from "./review-checkpoints.js"; +import { SingleUserOAuthProvider } from "./oauth-provider.js"; import { ProcessSessionManager } from "./process-sessions.js"; -import { createMcpServer } from "./server.js"; +import { createMcpServer, createServer as createDevspaceServer } from "./server.js"; import { SqliteWorkspaceStore } from "./workspace-store.js"; import { WorkspaceRegistry } from "./workspaces.js"; import { writeTestDevspaceConfig } from "./test-support/config.test.js"; @@ -287,6 +292,189 @@ test("open_workspace scopes checkout reuse to OpenAI session metadata", async (t assert.ok(Array.isArray(structuredContent(unscoped).agentsFiles)); }); +test("HTTP returns 503 when all MCP sessions are active", async (t) => { + const root = await mkdtemp(join(tmpdir(), "devspace-http-capacity-test-")); + const configDir = join(root, ".config"); + const stateDir = join(root, ".state"); + const httpServer = createHttpServer(); + httpServer.listen(0, "127.0.0.1"); + await once(httpServer, "listening"); + const address = httpServer.address(); + assert.ok(address && typeof address !== "string"); + const port = address.port; + const publicBaseUrl = `http://127.0.0.1:${port}`; + const config = loadConfig(writeTestDevspaceConfig(configDir, { + server: { + host: "127.0.0.1", + port, + publicBaseUrl, + maxMcpSessions: 2, + }, + storage: { stateDir }, + workspaces: { allowedRoots: [root] }, + logging: { level: "silent", requests: false, toolCalls: false }, + })); + let running: ReturnType | undefined; + t.after(async () => { + await running?.close(); + await new Promise((resolve, reject) => { + httpServer.close((error) => error ? reject(error) : resolve()); + }); + await rm(root, { recursive: true, force: true }); + }); + + const accessToken = await issueTestAccessToken(config); + running = createDevspaceServer(config, { incomingArtifactAdapters: [] }); + httpServer.on("request", running.app); + + const firstSession = await initializeHttpSession(publicBaseUrl, accessToken, 1); + const secondSession = await initializeHttpSession(publicBaseUrl, accessToken, 2); + const firstSse = await openSessionSse(publicBaseUrl, accessToken, firstSession); + const secondSse = await openSessionSse(publicBaseUrl, accessToken, secondSession); + + try { + const blocked = await initializeHttpResponse(publicBaseUrl, accessToken, 3); + assert.equal(blocked.status, 503); + assert.equal(blocked.headers.get("retry-after"), "5"); + const payload = await blocked.json() as { + error?: { message?: string }; + }; + assert.match(payload.error?.message ?? "", /session limit reached/); + } finally { + firstSse.abort(); + secondSse.abort(); + } +}); + +/** Issue an OAuth access token through the real provider used by the HTTP test. */ +async function issueTestAccessToken(config: ServerConfig): Promise { + await mkdir(config.stateDir, { recursive: true }); + const mcpUrl = new URL("/mcp", config.publicBaseUrl); + const provider = new SingleUserOAuthProvider(config.oauth, mcpUrl, config.stateDir); + const redirectUri = "http://127.0.0.1/callback"; + assert.ok(provider.clientsStore.registerClient); + const client: OAuthClientInformationFull = await provider.clientsStore.registerClient({ + redirect_uris: [redirectUri], + client_name: "HTTP capacity test", + }); + let redirectLocation: string | undefined; + const response = { + req: { + method: "POST", + body: { owner_token: config.oauth.ownerToken }, + }, + redirect(status: number, location: string) { + assert.equal(status, 302); + redirectLocation = location; + return response; + }, + } as unknown as ExpressResponse; + + try { + await provider.authorize(client, { + codeChallenge: "http-capacity-test-challenge", + redirectUri, + resource: mcpUrl, + scopes: config.oauth.scopes, + }, response); + assert.ok(redirectLocation); + const code = new URL(redirectLocation).searchParams.get("code"); + assert.ok(code); + const tokens = await provider.exchangeAuthorizationCode( + client, + code, + undefined, + redirectUri, + mcpUrl, + ); + return tokens.access_token; + } finally { + provider.close(); + } +} + +/** Send one MCP initialize request to the HTTP endpoint. */ +async function initializeHttpResponse( + publicBaseUrl: string, + accessToken: string, + id: number, +): Promise { + return fetch(`${publicBaseUrl}/mcp`, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + Accept: "application/json, text/event-stream", + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "http-capacity-test", version: "1.0.0" }, + }, + }), + }); +} + +/** Initialize one MCP session and send its initialized notification. */ +async function initializeHttpSession( + publicBaseUrl: string, + accessToken: string, + id: number, +): Promise { + const response = await initializeHttpResponse(publicBaseUrl, accessToken, id); + assert.equal(response.status, 200); + const sessionId = response.headers.get("mcp-session-id"); + assert.ok(sessionId); + await response.text(); + + const initialized = await fetch(`${publicBaseUrl}/mcp`, { + method: "POST", + headers: { + Authorization: `Bearer ${accessToken}`, + Accept: "application/json, text/event-stream", + "Content-Type": "application/json", + "mcp-session-id": sessionId, + "mcp-protocol-version": "2025-11-25", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + method: "notifications/initialized", + }), + }); + assert.equal(initialized.status, 202); + await initialized.text(); + return sessionId; +} + +/** Open a live SSE response so the session remains protected from eviction. */ +async function openSessionSse( + publicBaseUrl: string, + accessToken: string, + sessionId: string, +): Promise<{ response: Response; abort(): void }> { + const controller = new AbortController(); + const response = await fetch(`${publicBaseUrl}/mcp`, { + method: "GET", + headers: { + Authorization: `Bearer ${accessToken}`, + Accept: "text/event-stream", + "mcp-session-id": sessionId, + "mcp-protocol-version": "2025-11-25", + }, + signal: controller.signal, + }); + assert.equal(response.status, 200); + assert.match(response.headers.get("content-type") ?? "", /text\/event-stream/); + return { + response, + abort: () => controller.abort(), + }; +} + interface ServerFixture { client: Client; project: string; diff --git a/src/server.ts b/src/server.ts index 9e7ded7fd..862e66a46 100644 --- a/src/server.ts +++ b/src/server.ts @@ -37,6 +37,7 @@ import { SingleUserOAuthProvider } from "./oauth-provider.js"; import { McpSessionRegistry, type McpSessionCloseResult, + type McpSessionReservation, } from "./mcp-sessions.js"; import { ProcessSessionManager } from "./process-sessions.js"; import { createReviewCheckpointManager } from "./review-checkpoints.js"; @@ -706,6 +707,26 @@ export interface CreateServerOptions { incomingArtifactAdapters?: readonly IncomingArtifactAdapter[]; } +/** Release a session's active-response reference when its HTTP response ends. */ +function releaseMcpSessionWhenResponseEnds( + res: Response, + transports: McpSessionRegistry, + sessionId: string, +): () => void { + let released = false; + const release = () => { + if (released) return; + released = true; + res.off("finish", release); + res.off("close", release); + transports.release(sessionId); + }; + res.once("finish", release); + res.once("close", release); + return release; +} + +/** Create the authenticated HTTP server and its shared MCP session registry. */ export function createServer( config = loadConfig(), options: CreateServerOptions = {}, @@ -719,7 +740,9 @@ export function createServer( host: config.host, ...(allowedHosts ? { allowedHosts } : {}), }); - const transports = new McpSessionRegistry(); + const transports = new McpSessionRegistry({ + maxSessions: config.maxMcpSessions, + }); const mcpUrl = new URL("/mcp", config.publicBaseUrl); const resourceServerUrl = resourceUrlFromServerUrl(mcpUrl); const oauthProvider = new SingleUserOAuthProvider(config.oauth, mcpUrl, config.stateDir); @@ -862,23 +885,91 @@ export function createServer( isInitialize: initializeRequest, }); + let pendingReservation: McpSessionReservation | undefined; + let releaseActiveSession: (() => void) | undefined; try { let transport: Transport | undefined; if (sessionId) { - transport = transports.get(sessionId); + transport = transports.acquire(sessionId); if (!transport) { sendJsonRpcError(res, 404, -32000, "Unknown MCP session"); return; } + releaseActiveSession = releaseMcpSessionWhenResponseEnds( + res, + transports, + sessionId, + ); } else if (initializeRequest) { + const reservationResult = await transports.reserve(); + if (!reservationResult.ok) { + const closeError = reservationResult.error; + logEvent( + config.logging, + reservationResult.reason === "close_failed" ? "error" : "warn", + reservationResult.reason === "close_failed" + ? "mcp_session_capacity_close_failed" + : "mcp_session_capacity_exhausted", + { + requestId, + sessionIdPrefix: sessionIdPrefix(reservationResult.sessionId), + sessionCount: transports.size, + activeSessions: transports.activeSessions, + capacity: transports.maxSessions, + ...(closeError + ? { + error: + closeError instanceof Error + ? closeError.message + : String(closeError), + } + : {}), + ...requestLogFields(req, config), + }, + ); + res.setHeader("Retry-After", "5"); + sendJsonRpcError( + res, + 503, + -32000, + "MCP session limit reached; retry shortly.", + ); + return; + } + + pendingReservation = reservationResult.reservation; + if (pendingReservation.evictedSessionId) { + logEvent(config.logging, "info", "mcp_session_evicted", { + requestId, + reason: "capacity", + sessionIdPrefix: sessionIdPrefix( + pendingReservation.evictedSessionId, + ), + sessionCount: transports.size, + capacity: transports.maxSessions, + }); + } transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), - onsessioninitialized: (newSessionId) => { - if (transport) transports.register(newSessionId, transport); + onsessioninitialized: async (newSessionId) => { + if (!transport || !pendingReservation) return; + const reservation = pendingReservation; + pendingReservation = undefined; + const committed = await reservation.commit(newSessionId, transport, { active: true }); + if (!committed) return; + + releaseActiveSession = releaseMcpSessionWhenResponseEnds( + res, + transports, + newSessionId, + ); logEvent(config.logging, "info", "mcp_session_created", { requestId, sessionIdPrefix: sessionIdPrefix(newSessionId), + sessionCount: transports.size, + activeSessions: transports.activeSessions, + capacity: transports.maxSessions, ...requestLogFields(req, config), }); }, @@ -890,6 +981,8 @@ export function createServer( logEvent(config.logging, "info", "mcp_session_closed", { reason: "transport_close", sessionIdPrefix: sessionIdPrefix(closedSessionId), + sessionCount: transports.size, + capacity: transports.maxSessions, }); } }; @@ -917,6 +1010,11 @@ export function createServer( if (!res.headersSent) { sendJsonRpcError(res, 500, -32603, "Internal server error"); } + } finally { + pendingReservation?.release(); + if (releaseActiveSession && (res.writableEnded || res.destroyed)) { + releaseActiveSession(); + } } });