fix: clean up stale MCP sessions to bound memory growth (#71)

* fix: clean up stale MCP sessions

* fix: await MCP cleanup during shutdown

Author: Zhenyu98

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/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": [],

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<void> {
     }
   });
 
-  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<void> {

src/mcp-sessions.test.ts

@@ -0,0 +1,86 @@
+import assert from "node:assert/strict";
+import { McpSessionRegistry } from "./mcp-sessions.js";
+
+interface FakeTransport {
+  closeCalls: number;
+  close(): Promise<void>;
+}
+
+function createTransport(closeError?: Error): FakeTransport {
+  return {
+    closeCalls: 0,
+    async close() {
+      this.closeCalls += 1;
+      if (closeError) throw closeError;
+    },
+  };
+}
+
+let now = 0;
+const registry = new McpSessionRegistry<FakeTransport>({ 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);
+
+let finishDelayedClose: (() => void) | undefined;
+let delayedCloseResolved = false;
+const delayedTransport: FakeTransport = {
+  closeCalls: 0,
+  close() {
+    this.closeCalls += 1;
+    return new Promise<void>((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);

src/mcp-sessions.ts

@@ -0,0 +1,87 @@
+export interface ClosableMcpTransport {
+  close(): Promise<void>;
+}
+
+export interface McpSessionCloseResult {
+  sessionId: string;
+  error?: unknown;
+}
+
+interface McpSessionEntry<TTransport> {
+  transport: TTransport;
+  lastActivityAt: number;
+}
+
+export interface McpSessionRegistryOptions {
+  now?: () => number;
+}
+
+export class McpSessionRegistry<TTransport extends ClosableMcpTransport> {
+  private readonly sessions = new Map<string, McpSessionEntry<TTransport>>();
+  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<McpSessionCloseResult[]> {
+    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<McpSessionCloseResult[]> {
+    const sessions = Array.from(this.sessions, ([sessionId, entry]) => ({
+      sessionId,
+      transport: entry.transport,
+    }));
+    this.sessions.clear();
+    return closeSessions(sessions);
+  }
+}
+
+async function closeSessions<TTransport extends ClosableMcpTransport>(
+  sessions: Array<{ sessionId: string; transport: TTransport }>,
+): Promise<McpSessionCloseResult[]> {
+  return Promise.all(
+    sessions.map(async ({ sessionId, transport }) => {
+      try {
+        await transport.close();
+        return { sessionId };
+      } catch (error) {
+        return { sessionId, error };
+      }
+    }),
+  );
+}

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<void>((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,
+);

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<void>,
+): Promise<void> {
+  const httpClosed = new Promise<void>((resolve, reject) => {
+    httpServer.close((error) => {
+      if (error) reject(error);
+      else resolve();
+    });
+  });
+
+  await closeApplication();
+  await httpClosed;
+}