From 40b081af046fcdbc158ec0d4fdebe09f75d8fe04 Mon Sep 17 00:00:00 2001 From: Zhenyu Wu <62387006+Zhenyu98@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:24:20 +0800 Subject: [PATCH 1/2] fix: clean up stale MCP sessions --- package.json | 2 +- src/mcp-sessions.test.ts | 61 ++++++++++++++++++++++++++++ src/mcp-sessions.ts | 87 ++++++++++++++++++++++++++++++++++++++++ src/server.ts | 51 +++++++++++++++++++++-- 4 files changed, 196 insertions(+), 5 deletions(-) create mode 100644 src/mcp-sessions.test.ts create mode 100644 src/mcp-sessions.ts diff --git a/package.json b/package.json index 0b8c5eb..5436101 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "dev": "node scripts/dev-server.mjs", "postinstall": "node scripts/fix-node-pty-permissions.mjs", "start": "node dist/cli.js serve", - "test": "tsx src/config.test.ts && tsx src/ui/card-types.test.ts && tsx src/ui/patch-display.test.ts && tsx src/ui/tool-display.test.ts && tsx src/apply-patch.test.ts && tsx src/process-platform.test.ts && tsx src/process-sessions.test.ts && tsx src/local-agent-runtime.test.ts && tsx src/local-agent-adapters.test.ts && tsx src/local-agent-availability.test.ts && tsx src/local-agent-profiles.test.ts && tsx src/local-agent-targets.test.ts && tsx src/local-agent-store.test.ts && tsx src/roots.test.ts && tsx src/skills.test.ts && tsx src/workspaces.test.ts && tsx src/review-checkpoints.test.ts && tsx src/oauth-store.test.ts && tsx src/cli.test.ts", + "test": "tsx src/config.test.ts && tsx src/ui/card-types.test.ts && tsx src/ui/patch-display.test.ts && tsx src/ui/tool-display.test.ts && tsx src/apply-patch.test.ts && tsx src/process-platform.test.ts && tsx src/process-sessions.test.ts && tsx src/mcp-sessions.test.ts && tsx src/local-agent-runtime.test.ts && tsx src/local-agent-adapters.test.ts && tsx src/local-agent-availability.test.ts && tsx src/local-agent-profiles.test.ts && tsx src/local-agent-targets.test.ts && tsx src/local-agent-store.test.ts && tsx src/roots.test.ts && tsx src/skills.test.ts && tsx src/workspaces.test.ts && tsx src/review-checkpoints.test.ts && tsx src/oauth-store.test.ts && tsx src/cli.test.ts", "typecheck": "tsc -p tsconfig.json --noEmit" }, "keywords": [], diff --git a/src/mcp-sessions.test.ts b/src/mcp-sessions.test.ts new file mode 100644 index 0000000..d3d6663 --- /dev/null +++ b/src/mcp-sessions.test.ts @@ -0,0 +1,61 @@ +import assert from "node:assert/strict"; +import { McpSessionRegistry } from "./mcp-sessions.js"; + +interface FakeTransport { + closeCalls: number; + close(): Promise; +} + +function createTransport(closeError?: Error): FakeTransport { + return { + closeCalls: 0, + async close() { + this.closeCalls += 1; + if (closeError) throw closeError; + }, + }; +} + +let now = 0; +const registry = new McpSessionRegistry({ now: () => now }); +const staleTransport = createTransport(); +const activeTransport = createTransport(); + +registry.register("stale", staleTransport); +now = 1_000; +registry.register("active", activeTransport); +now = 1_500; +assert.equal(registry.get("active"), activeTransport); +now = 2_000; + +const idleResults = await registry.closeIdle(1_500); +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); + +const closeError = new Error("close failed"); +const failingTransport = createTransport(closeError); +registry.register("failing", failingTransport); +now = 10_000; + +const failingResults = await registry.closeIdle(1); +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); + +const first = createTransport(); +const second = createTransport(); +registry.register("first", first); +registry.register("second", second); +registry.remove("first"); + +const shutdownResults = await registry.closeAll(); +assert.deepEqual(shutdownResults, [{ sessionId: "second" }]); +assert.equal(first.closeCalls, 0); +assert.equal(second.closeCalls, 1); +assert.equal(registry.size, 0); diff --git a/src/mcp-sessions.ts b/src/mcp-sessions.ts new file mode 100644 index 0000000..3a84494 --- /dev/null +++ b/src/mcp-sessions.ts @@ -0,0 +1,87 @@ +export interface ClosableMcpTransport { + close(): Promise; +} + +export interface McpSessionCloseResult { + sessionId: string; + error?: unknown; +} + +interface McpSessionEntry { + transport: TTransport; + lastActivityAt: number; +} + +export interface McpSessionRegistryOptions { + now?: () => number; +} + +export class McpSessionRegistry { + private readonly sessions = new Map>(); + private readonly now: () => number; + + constructor(options: McpSessionRegistryOptions = {}) { + this.now = options.now ?? Date.now; + } + + get size(): number { + return this.sessions.size; + } + + register(sessionId: string, transport: TTransport): void { + this.sessions.set(sessionId, { + transport, + lastActivityAt: this.now(), + }); + } + + get(sessionId: string): TTransport | undefined { + const entry = this.sessions.get(sessionId); + if (!entry) return undefined; + + entry.lastActivityAt = this.now(); + return entry.transport; + } + + remove(sessionId: string): boolean { + return this.sessions.delete(sessionId); + } + + async closeIdle(idleTimeoutMs: number): Promise { + const cutoff = this.now() - idleTimeoutMs; + const idleSessions: Array<{ sessionId: string; transport: TTransport }> = []; + + for (const [sessionId, entry] of this.sessions) { + if (entry.lastActivityAt > cutoff) continue; + + this.sessions.delete(sessionId); + idleSessions.push({ sessionId, transport: entry.transport }); + } + + return closeSessions(idleSessions); + } + + async closeAll(): Promise { + const sessions = Array.from(this.sessions, ([sessionId, entry]) => ({ + sessionId, + transport: entry.transport, + })); + this.sessions.clear(); + return closeSessions(sessions); + } +} + +async function closeSessions( + sessions: Array<{ sessionId: string; transport: TTransport }>, +): Promise { + return Promise.all( + sessions.map(async ({ sessionId, transport }) => { + try { + await transport.close(); + return { sessionId }; + } catch (error) { + return { sessionId, error }; + } + }), + ); +} diff --git a/src/server.ts b/src/server.ts index bbf5dc3..dba3a0d 100644 --- a/src/server.ts +++ b/src/server.ts @@ -36,6 +36,10 @@ import { writeFileTool, } from "./pi-tools.js"; import { SingleUserOAuthProvider } from "./oauth-provider.js"; +import { + McpSessionRegistry, + type McpSessionCloseResult, +} from "./mcp-sessions.js"; import { ProcessSessionManager, type ProcessSnapshot } from "./process-sessions.js"; import { createReviewCheckpointManager } from "./review-checkpoints.js"; import { formatPathForPrompt } from "./skills.js"; @@ -49,6 +53,10 @@ import { } from "./local-agent-availability.js"; type Transport = StreamableHTTPServerTransport; +// MCP clients can reconnect without closing the previous transport. Bound stale +// session retention so abandoned MCP servers do not accumulate for the life of the process. +const MCP_SESSION_IDLE_TIMEOUT_MS = 24 * 60 * 60 * 1_000; +const MCP_SESSION_CLEANUP_INTERVAL_MS = 5 * 60 * 1_000; const WORKSPACE_APP_URI = "ui://devspace/workspace-app.html"; const WORKSPACE_APP_MANIFEST_ENTRY = "workspace-app.html"; const WRITE_TOOL_ANNOTATIONS = { @@ -1596,7 +1604,7 @@ export function createServer(config = loadConfig()): RunningServer { host: config.host, ...(allowedHosts ? { allowedHosts } : {}), }); - const transports = new Map(); + const transports = new McpSessionRegistry(); const mcpUrl = new URL("/mcp", config.publicBaseUrl); const resourceServerUrl = resourceUrlFromServerUrl(mcpUrl); const oauthProvider = new SingleUserOAuthProvider(config.oauth, mcpUrl, config.stateDir); @@ -1613,6 +1621,37 @@ export function createServer(config = loadConfig()): RunningServer { ? getLocalAgentProviderAvailabilitySnapshot() : []; + const logSessionCloseResults = ( + reason: "idle_timeout" | "server_shutdown", + results: McpSessionCloseResult[], + ) => { + for (const result of results) { + if (result.error) { + logEvent(config.logging, "warn", "mcp_session_close_failed", { + reason, + sessionIdPrefix: sessionIdPrefix(result.sessionId), + error: + result.error instanceof Error + ? result.error.message + : String(result.error), + }); + continue; + } + + logEvent(config.logging, "info", "mcp_session_closed", { + reason, + sessionIdPrefix: sessionIdPrefix(result.sessionId), + }); + } + }; + + const sessionCleanupTimer = setInterval(() => { + void transports + .closeIdle(MCP_SESSION_IDLE_TIMEOUT_MS) + .then((results) => logSessionCloseResults("idle_timeout", results)); + }, MCP_SESSION_CLEANUP_INTERVAL_MS); + sessionCleanupTimer.unref(); + if (config.logging.trustProxy) { app.set("trust proxy", true); } @@ -1716,7 +1755,7 @@ export function createServer(config = loadConfig()): RunningServer { transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID(), onsessioninitialized: (newSessionId) => { - if (transport) transports.set(newSessionId, transport); + if (transport) transports.register(newSessionId, transport); logEvent(config.logging, "info", "mcp_session_created", { requestId, sessionIdPrefix: sessionIdPrefix(newSessionId), @@ -1727,9 +1766,9 @@ export function createServer(config = loadConfig()): RunningServer { transport.onclose = () => { const closedSessionId = transport?.sessionId; - if (closedSessionId) { - transports.delete(closedSessionId); + if (closedSessionId && transports.remove(closedSessionId)) { logEvent(config.logging, "info", "mcp_session_closed", { + reason: "transport_close", sessionIdPrefix: sessionIdPrefix(closedSessionId), }); } @@ -1768,6 +1807,10 @@ export function createServer(config = loadConfig()): RunningServer { close: () => { if (closed) return; closed = true; + clearInterval(sessionCleanupTimer); + void transports + .closeAll() + .then((results) => logSessionCloseResults("server_shutdown", results)); processSessions.shutdown(); oauthProvider.close(); workspaceStore.close?.(); From ee372e7e4c2a5f16862a6cc5fbfb208baf47c5e5 Mon Sep 17 00:00:00 2001 From: Zhenyu Wu <62387006+Zhenyu98@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:15:46 +0800 Subject: [PATCH 2/2] fix: await MCP cleanup during shutdown --- package.json | 2 +- src/cli.ts | 20 +++++--- src/mcp-sessions.test.ts | 25 ++++++++++ src/server-shutdown.test.ts | 97 +++++++++++++++++++++++++++++++++++++ src/server-shutdown.ts | 18 +++++++ src/server.ts | 42 +++++++++------- 6 files changed, 180 insertions(+), 24 deletions(-) create mode 100644 src/server-shutdown.test.ts create mode 100644 src/server-shutdown.ts diff --git a/package.json b/package.json index 5436101..25f2105 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "dev": "node scripts/dev-server.mjs", "postinstall": "node scripts/fix-node-pty-permissions.mjs", "start": "node dist/cli.js serve", - "test": "tsx src/config.test.ts && tsx src/ui/card-types.test.ts && tsx src/ui/patch-display.test.ts && tsx src/ui/tool-display.test.ts && tsx src/apply-patch.test.ts && tsx src/process-platform.test.ts && tsx src/process-sessions.test.ts && tsx src/mcp-sessions.test.ts && tsx src/local-agent-runtime.test.ts && tsx src/local-agent-adapters.test.ts && tsx src/local-agent-availability.test.ts && tsx src/local-agent-profiles.test.ts && tsx src/local-agent-targets.test.ts && tsx src/local-agent-store.test.ts && tsx src/roots.test.ts && tsx src/skills.test.ts && tsx src/workspaces.test.ts && tsx src/review-checkpoints.test.ts && tsx src/oauth-store.test.ts && tsx src/cli.test.ts", + "test": "tsx src/config.test.ts && tsx src/ui/card-types.test.ts && tsx src/ui/patch-display.test.ts && tsx src/ui/tool-display.test.ts && tsx src/apply-patch.test.ts && tsx src/process-platform.test.ts && tsx src/process-sessions.test.ts && tsx src/mcp-sessions.test.ts && tsx src/server-shutdown.test.ts && tsx src/local-agent-runtime.test.ts && tsx src/local-agent-adapters.test.ts && tsx src/local-agent-availability.test.ts && tsx src/local-agent-profiles.test.ts && tsx src/local-agent-targets.test.ts && tsx src/local-agent-store.test.ts && tsx src/roots.test.ts && tsx src/skills.test.ts && tsx src/workspaces.test.ts && tsx src/review-checkpoints.test.ts && tsx src/oauth-store.test.ts && tsx src/cli.test.ts", "typecheck": "tsc -p tsconfig.json --noEmit" }, "keywords": [], diff --git a/src/cli.ts b/src/cli.ts index bbae1ae..7a1ac63 100644 --- a/src/cli.ts +++ b/src/cli.ts @@ -38,6 +38,7 @@ import { type DevspaceUserConfig, } from "./user-config.js"; import { expandHomePath } from "./roots.js"; +import { shutdownHttpServer } from "./server-shutdown.js"; type Command = "serve" | "init" | "doctor" | "config" | "agents" | "help" | "version"; const require = createRequire(import.meta.url); @@ -228,14 +229,21 @@ async function serve(): Promise { } }); - const shutdown = () => { - httpServer.close(() => { - close(); - process.exit(0); + let shuttingDown = false; + const shutdown = async () => { + if (shuttingDown) return; + shuttingDown = true; + await shutdownHttpServer(httpServer, close); + process.exit(0); + }; + const handleShutdown = () => { + void shutdown().catch((error) => { + console.error("devspace shutdown failed", error); + process.exit(1); }); }; - process.once("SIGINT", shutdown); - process.once("SIGTERM", shutdown); + process.once("SIGINT", handleShutdown); + process.once("SIGTERM", handleShutdown); } async function runDoctor(): Promise { diff --git a/src/mcp-sessions.test.ts b/src/mcp-sessions.test.ts index d3d6663..1b3ecf3 100644 --- a/src/mcp-sessions.test.ts +++ b/src/mcp-sessions.test.ts @@ -59,3 +59,28 @@ assert.deepEqual(shutdownResults, [{ sessionId: "second" }]); assert.equal(first.closeCalls, 0); assert.equal(second.closeCalls, 1); assert.equal(registry.size, 0); + +let finishDelayedClose: (() => void) | undefined; +let delayedCloseResolved = false; +const delayedTransport: FakeTransport = { + closeCalls: 0, + close() { + this.closeCalls += 1; + return new Promise((resolve) => { + finishDelayedClose = resolve; + }); + }, +}; +registry.register("delayed", delayedTransport); +const delayedClose = registry.closeAll(); +void delayedClose.then(() => { + delayedCloseResolved = true; +}); + +await Promise.resolve(); +assert.equal(delayedCloseResolved, false); +assert.equal(delayedTransport.closeCalls, 1); +finishDelayedClose?.(); +await delayedClose; +assert.equal(delayedCloseResolved, true); +assert.equal(registry.size, 0); diff --git a/src/server-shutdown.test.ts b/src/server-shutdown.test.ts new file mode 100644 index 0000000..7dca2ee --- /dev/null +++ b/src/server-shutdown.test.ts @@ -0,0 +1,97 @@ +import assert from "node:assert/strict"; +import { shutdownHttpServer } from "./server-shutdown.js"; + +let finishHttpClose: (() => void) | undefined; +let applicationCloseStarted = false; + +const drainingHttpServer = { + close(callback: (error?: Error) => void) { + finishHttpClose = () => callback(); + }, +}; + +const drainingShutdown = shutdownHttpServer(drainingHttpServer, async () => { + applicationCloseStarted = true; + assert.ok( + finishHttpClose, + "HTTP draining must start before application cleanup", + ); + finishHttpClose(); +}); + +await Promise.resolve(); +assert.equal( + applicationCloseStarted, + true, + "application cleanup must start while the HTTP server is draining", +); +await drainingShutdown; + +let finishApplicationClose: (() => void) | undefined; +let shutdownResolved = false; + +const immediatelyClosedHttpServer = { + close(callback: (error?: Error) => void) { + callback(); + }, +}; + +const delayedApplicationClose = () => + new Promise((resolve) => { + finishApplicationClose = resolve; + }); + +const delayedShutdown = shutdownHttpServer( + immediatelyClosedHttpServer, + delayedApplicationClose, +); +void delayedShutdown.then(() => { + shutdownResolved = true; +}); + +await Promise.resolve(); +assert.equal( + shutdownResolved, + false, + "shutdown must wait for asynchronous application cleanup", +); +finishApplicationClose?.(); +await delayedShutdown; +assert.equal(shutdownResolved, true); + +let finishDelayedHttpClose: (() => void) | undefined; +let httpDrainResolved = false; +const delayedHttpDrain = shutdownHttpServer( + { + close(callback: (error?: Error) => void) { + finishDelayedHttpClose = () => callback(); + }, + }, + async () => {}, +); +void delayedHttpDrain.then(() => { + httpDrainResolved = true; +}); + +await Promise.resolve(); +assert.equal( + httpDrainResolved, + false, + "shutdown must wait for active HTTP responses to drain", +); +finishDelayedHttpClose?.(); +await delayedHttpDrain; +assert.equal(httpDrainResolved, true); + +const httpCloseError = new Error("http close failed"); +await assert.rejects( + shutdownHttpServer( + { + close(callback: (error?: Error) => void) { + callback(httpCloseError); + }, + }, + async () => {}, + ), + httpCloseError, +); diff --git a/src/server-shutdown.ts b/src/server-shutdown.ts new file mode 100644 index 0000000..b9e6f4f --- /dev/null +++ b/src/server-shutdown.ts @@ -0,0 +1,18 @@ +export interface ClosableHttpServer { + close(callback: (error?: Error) => void): void; +} + +export async function shutdownHttpServer( + httpServer: ClosableHttpServer, + closeApplication: () => Promise, +): Promise { + const httpClosed = new Promise((resolve, reject) => { + httpServer.close((error) => { + if (error) reject(error); + else resolve(); + }); + }); + + await closeApplication(); + await httpClosed; +} diff --git a/src/server.ts b/src/server.ts index dba3a0d..7dd9629 100644 --- a/src/server.ts +++ b/src/server.ts @@ -42,6 +42,7 @@ import { } from "./mcp-sessions.js"; import { ProcessSessionManager, type ProcessSnapshot } from "./process-sessions.js"; import { createReviewCheckpointManager } from "./review-checkpoints.js"; +import { shutdownHttpServer } from "./server-shutdown.js"; import { formatPathForPrompt } from "./skills.js"; import { createWorkspaceStore } from "./workspace-store.js"; import { formatAgentsPath, WorkspaceRegistry } from "./workspaces.js"; @@ -82,7 +83,7 @@ interface RunningServer { app: ReturnType; config: ServerConfig; localAgentProviders: LocalAgentProviderAvailability[]; - close(): void; + close(): Promise; } type ToolContent = @@ -1799,21 +1800,21 @@ export function createServer(config = loadConfig()): RunningServer { } }); - let closed = false; + let closePromise: Promise | undefined; return { app, config, localAgentProviders, close: () => { - if (closed) return; - closed = true; - clearInterval(sessionCleanupTimer); - void transports - .closeAll() - .then((results) => logSessionCloseResults("server_shutdown", results)); - processSessions.shutdown(); - oauthProvider.close(); - workspaceStore.close?.(); + closePromise ??= (async () => { + clearInterval(sessionCleanupTimer); + const results = await transports.closeAll(); + logSessionCloseResults("server_shutdown", results); + processSessions.shutdown(); + oauthProvider.close(); + workspaceStore.close?.(); + })(); + return closePromise; }, }; } @@ -1843,12 +1844,19 @@ if (await isMainModule()) { } }); - const shutdown = () => { - httpServer.close(() => { - close(); - process.exit(0); + let shuttingDown = false; + const shutdown = async () => { + if (shuttingDown) return; + shuttingDown = true; + await shutdownHttpServer(httpServer, close); + process.exit(0); + }; + const handleShutdown = () => { + void shutdown().catch((error) => { + console.error("devspace shutdown failed", error); + process.exit(1); }); }; - process.once("SIGINT", shutdown); - process.once("SIGTERM", shutdown); + process.once("SIGINT", handleShutdown); + process.once("SIGTERM", handleShutdown); }