refactor(shutdown): coordinate worker drain handshake

src/main/service-graph/create-lifecycle.ts

@@ -9,7 +9,7 @@ import type { createOrchestrationService } from "../orchestration-service";
 import type { createWorkerService } from "../worker-service";
 import type { PiRuntimeServices } from "../service-graph-pi-runtime";
 
-/** Overall budget for awaited shutdown drain (worker exit + API close). */
+/** Overall budget for awaited shutdown drain. */
 const SHUTDOWN_BUDGET_MS = 5_000;
 
 export interface LifecycleDeps {
@@ -38,10 +38,9 @@ export interface LifecycleDeps {
 
 /**
  * Build the app lifecycle service and wire a deterministic shutdown drain
- * that awaits the utility-process worker's exit (and API server close)
- * before calling app.exit(). On 5s timeout we log the stuck run ids and
- * exit(1) so operators see an unclean shutdown rather than a silent
- * force-kill.
+ * that awaits the utility-process worker's ready-to-exit handshake
+ * before calling app.exit(). On timeout we log the stuck run ids and
+ * exit(1) so operators see an unclean shutdown.
  */
 export function createLifecycle(deps: LifecycleDeps) {
   const {
@@ -72,8 +71,6 @@ export function createLifecycle(deps: LifecycleDeps) {
   }
 
   async function drain(): Promise<number> {
-    // Stop accepting new scheduled work first; cancel in-flight runs in the
-    // worker via the utility-process shutdown postMessage inside destroy().
     try { orchestrationService.stopAllSchedulers(); } catch (err) { logger.warn("[lifecycle] stopAllSchedulers failed:", err); }
     try { automationService.destroy(); } catch (err) { logger.warn("[lifecycle] automationService.destroy failed:", err); }
     try { apiServerRef.current?.stop(); } catch (err) { logger.warn("[lifecycle] apiServer.stop failed:", err); }
@@ -81,22 +78,12 @@ export function createLifecycle(deps: LifecycleDeps) {
     try { browserManagerRef.current?.shutdown(); } catch (err) { logger.warn("[lifecycle] browserManager.shutdown failed:", err); }
     browserManagerRef.current = null;
 
-    // workerService.destroy() sends "shutdown" to the utility process and has
-    // its own 2s SIGKILL safety net. Awaiting the worker's exit event closes
-    // the remaining event-loop holder (the utilityProcess handle).
-    const workerExit = new Promise<void>((resolve) => {
-      const worker = workerService.getWorker();
-      if (!worker) { resolve(); return; }
-      worker.once("exit", () => resolve());
-    });
-    workerService.destroy();
-
     const timeout = new Promise<"timeout">((resolve) => {
       const handle = setTimeout(() => resolve("timeout"), SHUTDOWN_BUDGET_MS);
       if (typeof handle.unref === "function") handle.unref();
     });
     const outcome = await Promise.race([
-      workerExit.then(() => "drained" as const),
+      workerService.shutdown().then(() => "drained" as const),
       timeout,
     ]);
 

src/main/worker-service.ts

@@ -32,6 +32,7 @@ export function createWorkerService({
 
   let worker: ReturnType<typeof utilityProcess.fork> | null = null;
   let destroyed = false;
+  let shutdownPromise: Promise<void> | null = null;
   const pendingRetryTimers = new Set<ReturnType<typeof setTimeout>>();
 
   /** Crash-loop backoff state. Resets to 0 whenever a worker lives long enough
@@ -110,6 +111,16 @@ export function createWorkerService({
     });
 
     worker.on("message", async (message) => {
+      if (isShutdownReadyMessage(message)) {
+        try {
+          worker?.postMessage({ type: "shutdown-exit" });
+        } catch (error) {
+          logger.debug(
+            `[worker-service] shutdown-exit postMessage failed: ${error instanceof Error ? error.message : String(error)}`,
+          );
+        }
+        return;
+      }
       await handleWorkerMessage(
         message,
         worker,
@@ -140,66 +151,58 @@ export function createWorkerService({
     });
   }
 
-  function destroy() {
-    destroyed = true;
-
-    // Clear all pending retry timers so they never fire after quit
-    for (const timer of pendingRetryTimers) {
-      clearTimeout(timer);
-    }
-    pendingRetryTimers.clear();
-
-    if (worker) {
-      // Tell the worker to kill all its child CLI processes and exit.
-      // Previous approach sent individual cancel-run messages then immediately
-      // killed the worker — a race condition that left CLI processes as orphans
-      // because the worker died before it could forward SIGTERM to its children.
-      try {
-        worker.postMessage({ type: "shutdown" });
-      } catch (error) {
-        logger.debug(
-          `[worker-service] shutdown postMessage failed: ${error instanceof Error ? error.message : String(error)}`,
-        );
-      }
-
-      // Safety net: force-kill worker after 2s if shutdown didn't complete.
-      // Track the timer so we can cancel it the moment the worker cleanly
-      // exits — otherwise the pending timer keeps the event loop alive.
-      const ref = worker;
-      const killTimer = setTimeout(() => {
-        try {
-          ref.kill();
-        } catch (error) {
-          logger.debug(
-            `[worker-service] worker kill() failed (likely already dead): ${error instanceof Error ? error.message : String(error)}`,
-          );
-        }
-      }, 2000);
-      ref.once("exit", () => clearTimeout(killTimer));
-      worker = null;
-    }
-
-    // Mark all running/queued runs as interrupted so DB state is clean on restart
+  function collectShutdownRuns(): Array<{
+    id: string;
+    taskId: string | null;
+    workspaceId: string | null;
+  }> {
     try {
       const snapshot = db.projection.snapshot.listSnapshot();
-      for (const run of snapshot.runs) {
-        if (run.status === "running" || run.status === "queued") {
-          db.core.runs.finishRun({
-            runId: run.id,
-            status: "interrupted",
-            summary: "Application shut down.",
-            errorMessage: null,
-            exitCode: null,
-            signal: null,
-            model: null,
-            unattributedCost: true,
-          });
-          if (run.workspaceId) {
-            db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id);
-          }
-          if (run.taskId) {
-            db.core.tasks.releaseTaskLock(run.taskId, run.id);
-          }
+      return snapshot.runs
+        .filter((run) => run.status === "running" || run.status === "queued")
+        .map((run) => ({ id: run.id, taskId: run.taskId, workspaceId: run.workspaceId }));
+    } catch (error) {
+      logger.warn(
+        `[worker-service] shutdown run collection failed: ${error instanceof Error ? error.message : String(error)}`,
+      );
+      return [];
+    }
+  }
+
+  function interruptShutdownRuns(runs: Array<{
+    id: string;
+    taskId: string | null;
+    workspaceId: string | null;
+  }>): void {
+    try {
+      for (const run of runs) {
+        let status: string;
+        try {
+          status = db.core.runs.getRun(run.id).status;
+        } catch (error) {
+          logger.warn(
+            `[worker-service] shutdown cleanup skipped missing run ${run.id}: ${error instanceof Error ? error.message : String(error)}`,
+          );
+          continue;
+        }
+        if (status !== "running" && status !== "queued") {
+          continue;
+        }
+        db.core.runs.finishRun({
+          runId: run.id,
+          status: "interrupted",
+          summary: "Application shut down.",
+          errorMessage: null,
+          exitCode: null,
+          signal: null,
+          model: null,
+          unattributedCost: true,
+        });
+        if (run.workspaceId) {
+          db.core.workspaces.releaseWorkspaceLock(run.workspaceId, run.id);
+        }
+        if (run.taskId) {
+          db.core.tasks.releaseTaskLock(run.taskId, run.id);
         }
       }
     } catch (error) {
@@ -209,9 +212,52 @@ export function createWorkerService({
     }
   }
 
+  function shutdown(): Promise<void> {
+    if (shutdownPromise) return shutdownPromise;
+    destroyed = true;
+    const shutdownRuns = collectShutdownRuns();
+
+    for (const timer of pendingRetryTimers) {
+      clearTimeout(timer);
+    }
+    pendingRetryTimers.clear();
+
+    if (!worker) {
+      interruptShutdownRuns(shutdownRuns);
+      shutdownPromise = Promise.resolve();
+      return shutdownPromise;
+    }
+
+    const ref = worker;
+    shutdownPromise = new Promise<void>((resolve) => {
+      ref.once("exit", () => {
+        if (worker === ref) worker = null;
+        interruptShutdownRuns(shutdownRuns);
+        resolve();
+      });
+      try {
+        ref.postMessage({ type: "shutdown" });
+      } catch (error) {
+        logger.debug(
+          `[worker-service] shutdown postMessage failed: ${error instanceof Error ? error.message : String(error)}`,
+        );
+      }
+    });
+    return shutdownPromise;
+  }
+
+  function destroy() {
+    void shutdown();
+  }
+
   return {
     createWorker,
     getWorker,
+    shutdown,
     destroy,
   };
 }
+
+function isShutdownReadyMessage(message: unknown): message is { type: "shutdown-ready" } {
+  return Boolean(message && typeof message === "object" && (message as { type?: unknown }).type === "shutdown-ready");
+}

src/main/worker/entry/index.ts

@@ -25,6 +25,8 @@ interface ActiveRun {
 const active = new Map<string, ActiveRun>();
 const queue: RunJob[] = [];
 const cancelledRuns = new Set<string>();
+let shuttingDown = false;
+let drainAckResolver: (() => void) | null = null;
 
 function post(message: ChildMessage) {
   parentPort.postMessage(message);
@@ -57,6 +59,7 @@ function isWorkspaceBusy(lockKey: string): boolean {
 }
 
 function enqueue(job: RunJob) {
+  if (shuttingDown) return;
   queue.push(job);
   post({
     type: "run-status",
@@ -73,6 +76,7 @@ function enqueue(job: RunJob) {
  * `startRun`. No async window between claim and dispatch.
  */
 function tickQueue() {
+  if (shuttingDown) return;
   let i = 0;
   while (i < queue.length) {
     const job = queue[i];
@@ -96,6 +100,17 @@ async function startRun(job: RunJob) {
 function acknowledgeRunCompletion(runId: string) {
   const released = active.delete(runId);
   if (released) tickQueue();
+  if (shuttingDown && active.size === 0 && drainAckResolver) {
+    drainAckResolver();
+    drainAckResolver = null;
+  }
+}
+
+function waitForCompletionAcks(): Promise<void> {
+  if (active.size === 0) return Promise.resolve();
+  return new Promise((resolve) => {
+    drainAckResolver = resolve;
+  });
 }
 
 function cancelRun(runId: string) {
@@ -151,17 +166,18 @@ parentPort.on("message", (event) => {
     return;
   }
   if (message.type === "shutdown") {
+    if (shuttingDown) return;
+    shuttingDown = true;
     void (async () => {
-      // Graceful drain: abort every active run and await its promise
-      // settlement so that final `run-finished` messages and
-      // `persistMessage` RPCs reach the parent before we exit. Bounded by
-      // a 5s budget so a stuck run does not wedge the shutdown handshake
-      // (the parent has its own kill-timer as a safety net).
-      const drainBudgetMs = 5_000;
-      const timeout = new Promise<void>((resolve) => setTimeout(resolve, drainBudgetMs));
-      await Promise.race([runner.cancelAll(), timeout]);
+      queue.splice(0);
+      await runner.cancelAll();
+      await waitForCompletionAcks();
       rpc.cancelAll("Worker shutting down.");
-      process.exit(0);
+      post({ type: "shutdown-ready" });
     })();
+    return;
+  }
+  if (message.type === "shutdown-exit") {
+    process.exit(0);
   }
 });

src/main/worker/entry/message-handler.ts

@@ -118,5 +118,7 @@ export async function handleWorkerMessage(
       deps.getRuntimeMessageInjector?.()?.handleSteerAck(event.correlationId, event.accepted);
       return;
     }
+    case "shutdown-ready":
+      return;
   }
 }

src/main/worker/entry/types.ts

@@ -93,6 +93,9 @@ export const workerMessageSchema = z.union([
     correlationId: z.string(),
     accepted: z.boolean(),
   }),
+  z.object({
+    type: z.literal("shutdown-ready"),
+  }),
 ]);
 
 export interface WorkerServiceDependencies {

src/main/worker/runtime/helpers.ts

@@ -49,6 +49,7 @@ export type ParentMessage =
   | { type: "cancel-run"; runId: string }
   | { type: "run-completion-ack"; runId: string }
   | { type: "shutdown" }
+  | { type: "shutdown-exit" }
   | { type: "steer-run"; runId: string; correlationId: string; message: AgentMessage }
   | { type: "rpc-response"; rpcId: string; response: { ok: true; result: unknown } | { ok: false; error: string } };
 
@@ -89,7 +90,8 @@ export type ChildMessage =
     }
   | { type: "rpc-request"; rpcId: string; runId: string; method: string; payload: unknown }
   | { type: "run-agent-state"; runId: string; state: RunAgentLifecycleState }
-  | { type: "steer-ack"; runId: string; correlationId: string; accepted: boolean };
+  | { type: "steer-ack"; runId: string; correlationId: string; accepted: boolean }
+  | { type: "shutdown-ready" };
 
 export function buildRedactor(env: Record<string, string>): (text: string) => string {
   const sensitivePatterns: string[] = [];

tests/unit/worker-completion-boundary.test.ts

@@ -285,4 +285,62 @@ describe("worker completion boundary", () => {
       expect.objectContaining({ type: "run-finished", runId: NEXT_RUN_ID }),
     );
   });
+
+  it("reports shutdown-ready only after parent acknowledges run completion", async () => {
+    const parentPort = new EventEmitter() as EventEmitter & {
+      postMessage: ReturnType<typeof vi.fn>;
+    };
+    parentPort.postMessage = vi.fn();
+    Object.defineProperty(process, "parentPort", {
+      configurable: true,
+      value: parentPort,
+    });
+    const exitSpy = vi.spyOn(process, "exit").mockImplementation((() => undefined) as never);
+
+    vi.doMock("@main/worker/runtime", () => ({
+      WorkerRunner: class TestWorkerRunner {
+        constructor(
+          private readonly deps: {
+            post: (message: unknown) => void;
+          },
+        ) {}
+
+        async run(job: RunJob): Promise<void> {
+          this.deps.post(runFinishedMessage(job.runId));
+        }
+
+        cancel(): void {}
+
+        steer(): boolean {
+          return false;
+        }
+
+        async cancelAll(): Promise<void> {}
+      },
+    }));
+
+    await import("@main/worker/entry/index");
+
+    parentPort.emit("message", { data: { type: "enqueue-run", payload: runJob(RUN_ID) } });
+    await Promise.resolve();
+    await Promise.resolve();
+
+    parentPort.emit("message", { data: { type: "shutdown" } });
+    await Promise.resolve();
+    await Promise.resolve();
+
+    expect(parentPort.postMessage).toHaveBeenCalledWith(
+      expect.objectContaining({ type: "run-finished", runId: RUN_ID }),
+    );
+    expect(parentPort.postMessage).not.toHaveBeenCalledWith({ type: "shutdown-ready" });
+
+    parentPort.emit("message", { data: { type: "run-completion-ack", runId: RUN_ID } });
+    await Promise.resolve();
+    await Promise.resolve();
+
+    expect(parentPort.postMessage).toHaveBeenCalledWith({ type: "shutdown-ready" });
+
+    parentPort.emit("message", { data: { type: "shutdown-exit" } });
+    expect(exitSpy).toHaveBeenCalledWith(0);
+  });
 });

tests/unit/worker-service-shutdown.test.ts

@@ -0,0 +1,86 @@
+import { EventEmitter } from "node:events";
+import { describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  forkMock: vi.fn(),
+}));
+
+vi.mock("electron", () => ({
+  app: { isPackaged: false },
+  utilityProcess: { fork: mocks.forkMock },
+}));
+
+import { createWorkerService } from "../../src/main/worker-service";
+
+function createWorker() {
+  const worker = new EventEmitter() as EventEmitter & {
+    stdout: EventEmitter;
+    stderr: EventEmitter;
+    postMessage: ReturnType<typeof vi.fn>;
+  };
+  worker.stdout = new EventEmitter();
+  worker.stderr = new EventEmitter();
+  worker.postMessage = vi.fn();
+  return worker;
+}
+
+describe("worker service shutdown", () => {
+  it("keeps completed run outcomes when cleaning shutdown runs", async () => {
+    const worker = createWorker();
+    mocks.forkMock.mockReturnValueOnce(worker);
+    const finishRun = vi.fn();
+    const db = {
+      projection: {
+        snapshot: {
+          listSnapshot: vi.fn(() => ({
+            runs: [
+              { id: "run-succeeded", status: "running", taskId: null, workspaceId: null },
+              { id: "run-queued", status: "queued", taskId: null, workspaceId: null },
+            ],
+          })),
+        },
+      },
+      core: {
+        runs: {
+          getRun: vi.fn((id: string) => ({
+            id,
+            status: id === "run-succeeded" ? "succeeded" : "queued",
+          })),
+          finishRun,
+        },
+        workspaces: { releaseWorkspaceLock: vi.fn() },
+        tasks: { releaseTaskLock: vi.fn() },
+      },
+    };
+    const service = createWorkerService({
+      db,
+      logger: { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() },
+      workerPath: "/tmp/worker.js",
+      isQuitting: () => false,
+      publishDomainChanged: vi.fn(),
+      emitEvent: vi.fn(),
+      wakeAgentIfPossible: vi.fn(() => null),
+      dispatchRunToWorker: vi.fn(),
+      processRunCompletion: vi.fn(),
+      onRunStarted: vi.fn(),
+      onRunCompleted: vi.fn(),
+      notify: vi.fn(),
+      budgetService: undefined,
+      sessionService: undefined,
+      getRuntimeBridge: vi.fn(() => null),
+      getRuntimeMessageInjector: vi.fn(() => null),
+    } as never);
+
+    service.createWorker();
+    const shutdown = service.shutdown();
+    worker.emit("exit", 0);
+    await shutdown;
+
+    expect(worker.postMessage).toHaveBeenCalledWith({ type: "shutdown" });
+    expect(finishRun).toHaveBeenCalledTimes(1);
+    expect(finishRun).toHaveBeenCalledWith(expect.objectContaining({
+      runId: "run-queued",
+      status: "interrupted",
+    }));
+  });
+});