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

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

@@ -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<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/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;
+}

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<typeof createMcpExpressApp>;
   config: ServerConfig;
   localAgentProviders: LocalAgentProviderAvailability[];
-  close(): void;
+  close(): Promise<void>;
 }
 
 type ToolContent =
@@ -1799,21 +1800,21 @@ export function createServer(config = loadConfig()): RunningServer {
     }
   });
 
-  let closed = false;
+  let closePromise: Promise<void> | 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);
 }