fix(cli): stabilize pm2 session handling

Author: pallyoung

packages/cli/src/pm2-control.test.ts

@@ -45,7 +45,10 @@ describe("pm2-control", () => {
     delete process.env.CODER_STUDIO_RUNTIME_JSON_PATH;
 
     connect.mockImplementation((callback: (error: Error | null) => void) => callback(null));
-    disconnect.mockImplementation(() => undefined);
+    disconnect.mockImplementation((callback?: (error: Error | null) => void) => {
+      callback?.(null);
+      return undefined;
+    });
     start.mockImplementation(
       (_config: unknown, callback: (error: Error | null, apps: unknown[]) => void) => {
         writeRuntimeConfig({
@@ -173,6 +176,67 @@ describe("pm2-control", () => {
     ).resolves.toBe("waiting");
 
     expect(start).not.toHaveBeenCalled();
+    await pendingStart;
+  });
+
+  it("reuses one pm2 session while polling deletion during startup", async () => {
+    describeProcess
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [{ pid: 111, pm2_env: { status: "online", restart_time: 0 } }])
+      )
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [{ pid: 111, pm2_env: { status: "stopping", restart_time: 0 } }])
+      )
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [])
+      );
+
+    await startManagedServer({
+      script: "/cli/dist/esm/server-runner.js",
+      cwd: "/repo",
+      waitMs: 10,
+    });
+
+    expect(connect).toHaveBeenCalledTimes(1);
+    expect(disconnect).toHaveBeenCalledTimes(1);
+  });
+
+  it("keeps waiting during startup when delete reports missing but the old app still lingers", async () => {
+    describeProcess
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [{ pid: 111, pm2_env: { status: "online", restart_time: 0 } }])
+      )
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [{ pid: 111, pm2_env: { status: "stopping", restart_time: 0 } }])
+      )
+      .mockImplementationOnce(
+        (_name: string, callback: (error: Error | null, result: unknown[]) => void) =>
+          callback(null, [])
+      );
+    deleteProcess.mockImplementationOnce((_name: string, callback: (error: Error | null) => void) =>
+      callback(new Error("process or namespace not found"))
+    );
+
+    const pendingStart = startManagedServer({
+      script: "/cli/dist/esm/server-runner.js",
+      cwd: "/repo",
+      waitMs: 10,
+    });
+
+    await expect(
+      Promise.race([
+        pendingStart.then(() => "started"),
+        new Promise((resolve) => setTimeout(() => resolve("waiting"), 20)),
+      ])
+    ).resolves.toBe("waiting");
+
+    expect(start).not.toHaveBeenCalled();
+    await pendingStart;
   });
 
   it("ignores delete-time missing errors when requested", async () => {
@@ -185,6 +249,8 @@ describe("pm2-control", () => {
     );
 
     await expect(deleteManagedServer({ ignoreMissing: true })).resolves.toBe(false);
+    expect(connect).toHaveBeenCalledTimes(1);
+    expect(disconnect).toHaveBeenCalledTimes(1);
   });
 
   it("fails background startup when runtime readiness times out", async () => {
@@ -263,6 +329,8 @@ describe("pm2-control", () => {
       pm2Pid: null,
       restartCount: 0,
     });
+    expect(connect).toHaveBeenCalledTimes(1);
+    expect(disconnect).toHaveBeenCalledTimes(1);
   });
 
   it("maps an online PM2 app to running status", async () => {

packages/cli/src/pm2-control.ts

@@ -9,6 +9,7 @@ const PM2_RESTART_DELAY_MS = 2000;
 const PM2_MIN_UPTIME = "5s";
 const PM2_MAX_RESTARTS = 10;
 const PM2_DELETE_WAIT_MS = 5000;
+const PM2_DISCONNECT_WAIT_MS = 1000;
 const STARTUP_POLL_INTERVAL_MS = 100;
 const STARTUP_FAILURE_GUIDANCE =
   "Run `coder-studio logs` for details or `coder-studio serve --foreground` for interactive debugging.";
@@ -62,7 +63,7 @@ const isPm2BrokenStateError = (error: unknown): boolean => {
 
 type Pm2Module = {
   connect: (cb: (err: Error | null) => void) => void;
-  disconnect: () => void;
+  disconnect: (cb?: (err: Error | null, data?: unknown) => void) => void;
   describe: (name: string, cb: (err: Error | null, result: unknown[]) => void) => void;
   delete: (name: string, cb: (err: Error | null) => void) => void;
   start: (opts: unknown, cb: (err: Error | null) => void) => void;
@@ -111,36 +112,51 @@ const sleep = async (ms: number): Promise<void> =>
 
 const disconnectPm2 = async (): Promise<void> => {
   const pm2 = await loadPm2();
-  pm2.disconnect();
+  await new Promise<void>((resolve) => {
+    let settled = false;
+    const finish = () => {
+      if (settled) {
+        return;
+      }
+
+      settled = true;
+      resolve();
+    };
+
+    const timer = setTimeout(finish, PM2_DISCONNECT_WAIT_MS);
+    try {
+      pm2.disconnect(() => {
+        clearTimeout(timer);
+        finish();
+      });
+    } catch {
+      clearTimeout(timer);
+      finish();
+    }
+  });
 };
 
-const describeManagedServer = async (): Promise<Pm2ProcessDescription[]> => {
-  const pm2 = await loadPm2();
-  return new Promise((resolve, reject) => {
+const describeManagedServer = async (pm2: Pm2Module): Promise<Pm2ProcessDescription[]> =>
+  new Promise((resolve, reject) => {
     pm2.describe(MANAGED_SERVER_NAME, (error, result) => {
       if (error) {
         reject(error);
         return;
       }
-
       resolve((result ?? []) as Pm2ProcessDescription[]);
     });
   });
-};
 
-const removeManagedServer = async (): Promise<void> => {
-  const pm2 = await loadPm2();
-  return new Promise((resolve, reject) => {
+const removeManagedServer = async (pm2: Pm2Module): Promise<void> =>
+  new Promise((resolve, reject) => {
     pm2.delete(MANAGED_SERVER_NAME, (error) => {
       if (error) {
         reject(error);
         return;
       }
-
       resolve();
     });
   });
-};
 
 /**
  * Kill the PM2 daemon to clear stale paths/caches.
@@ -159,9 +175,10 @@ const killPm2Daemon = async (): Promise<void> => {
  * Try to connect to PM2, and if it's in a broken state (stale worktree path),
  * kill the daemon and reconnect fresh.
  */
-const connectWithRecovery = async (): Promise<void> => {
+const connectWithRecovery = async (): Promise<Pm2Module> => {
   try {
     await connectPm2();
+    return loadPm2();
   } catch (error) {
     if (isPm2BrokenStateError(error)) {
       console.warn("PM2 daemon is in a stale state. Killing and reconnecting...");
@@ -174,23 +191,25 @@ const connectWithRecovery = async (): Promise<void> => {
       // Clear cached module so next loadPm2 gets a fresh instance
       cachedPm2 = null;
       await connectPm2();
+      return loadPm2();
     } else {
       throw error;
     }
   }
 };
 
-const withPm2Connection = async <T>(operation: () => Promise<T>): Promise<T> => {
-  await connectWithRecovery();
+const withPm2Connection = async <T>(operation: (pm2: Pm2Module) => Promise<T>): Promise<T> => {
+  const pm2 = await connectWithRecovery();
 
   try {
-    return await operation();
+    return await operation(pm2);
   } finally {
     await disconnectPm2();
   }
 };
 
 const waitForRuntimeReady = async (
+  pm2: Pm2Module,
   waitMs: number,
   logOffsets: StartupLogOffsets
 ): Promise<void> => {
@@ -201,7 +220,7 @@ const waitForRuntimeReady = async (
       return;
     }
 
-    const processes = await describeManagedServer();
+    const processes = await describeManagedServer(pm2);
     const process = processes[0];
     if (!process) {
       throw createStartupError(
@@ -226,22 +245,11 @@ const waitForRuntimeReady = async (
   throw createStartupError(`runtime readiness timed out after ${waitMs}ms`, logOffsets);
 };
 
-const waitForManagedServerExit = async (): Promise<void> => {
-  while (true) {
-    const processes = await describeManagedServer();
-    if (processes.length === 0) {
-      return;
-    }
-
-    await sleep(STARTUP_POLL_INTERVAL_MS);
-  }
-};
-
-const waitForManagedServerDeletion = async (waitMs: number): Promise<void> => {
+const waitForManagedServerDeletion = async (pm2: Pm2Module, waitMs: number): Promise<void> => {
   const deadline = Date.now() + waitMs;
 
   while (Date.now() <= deadline) {
-    const processes = await withPm2Connection(describeManagedServer);
+    const processes = await describeManagedServer(pm2);
     if (processes.length === 0) {
       return;
     }
@@ -257,6 +265,34 @@ const waitForManagedServerDeletion = async (waitMs: number): Promise<void> => {
   throw new Error(`Timed out waiting for the managed server to stop after ${waitMs}ms.`);
 };
 
+const deleteManagedServerInSession = async (
+  pm2: Pm2Module,
+  {
+    ignoreMissing = false,
+  }: {
+    ignoreMissing?: boolean;
+  } = {}
+): Promise<boolean> => {
+  const processes = await describeManagedServer(pm2);
+  if (processes.length === 0) {
+    return false;
+  }
+
+  try {
+    await removeManagedServer(pm2);
+  } catch (error) {
+    if (ignoreMissing && isMissingManagedServerError(error)) {
+      await waitForManagedServerDeletion(pm2, PM2_DELETE_WAIT_MS);
+      return false;
+    }
+
+    throw error;
+  }
+
+  await waitForManagedServerDeletion(pm2, PM2_DELETE_WAIT_MS);
+  return true;
+};
+
 const ensureLogDirectory = (): void => {
   mkdirSync(join(homedir(), ".coder-studio", "logs"), { recursive: true });
 };
@@ -307,56 +343,26 @@ export const deleteManagedServer = async ({
   ignoreMissing = false,
 }: {
   ignoreMissing?: boolean;
-} = {}): Promise<boolean> => {
-  const deleted = await withPm2Connection(async () => {
-    const processes = await describeManagedServer();
-    if (processes.length === 0) {
-      return false;
-    }
-
-    try {
-      await removeManagedServer();
-      return true;
-    } catch (error) {
-      if (ignoreMissing && isMissingManagedServerError(error)) {
-        return false;
-      }
-
-      throw error;
-    }
-  });
-
-  if (!deleted) {
-    return false;
-  }
-
-  await waitForManagedServerDeletion(PM2_DELETE_WAIT_MS);
-  return true;
-};
+} = {}): Promise<boolean> =>
+  withPm2Connection((pm2) => deleteManagedServerInSession(pm2, { ignoreMissing }));
 
 export const startManagedServer = async ({
   script,
   cwd,
   waitMs,
   args,
-}: StartManagedServerOptions): Promise<void> => {
-  // First try to delete any existing managed server
-  await deleteManagedServer({ ignoreMissing: true });
+}: StartManagedServerOptions): Promise<void> =>
+  withPm2Connection(async (pm2) => {
+    await deleteManagedServerInSession(pm2, { ignoreMissing: true });
 
-  // Wait for the old process to actually exit
-  await withPm2Connection(waitForManagedServerExit);
-
-  // Clear stale runtime config
-  if (readRuntimeConfig()) {
-    deleteRuntimeConfig();
-  }
-
-  ensureLogDirectory();
-  const { outFile, errFile } = getLogPaths();
-  const pm2 = await loadPm2();
+    if (readRuntimeConfig()) {
+      deleteRuntimeConfig();
+    }
 
-  await withPm2Connection(async () => {
+    ensureLogDirectory();
+    const { outFile, errFile } = getLogPaths();
     const logOffsets = captureStartupLogOffsets();
+
     await new Promise<void>((resolve, reject) => {
       pm2.start(
         {
@@ -386,13 +392,12 @@ export const startManagedServer = async ({
       );
     });
 
-    await waitForRuntimeReady(waitMs, logOffsets);
+    await waitForRuntimeReady(pm2, waitMs, logOffsets);
   });
-};
 
 export const getManagedServerStatus = async (): Promise<ManagedServerStatus> =>
-  withPm2Connection(async () => {
-    const processes = await describeManagedServer();
+  withPm2Connection(async (pm2) => {
+    const processes = await describeManagedServer(pm2);
     const process = processes[0];
 
     if (!process) {