diff --git a/src/main/service-graph/create-lifecycle.ts b/src/main/service-graph/create-lifecycle.ts index 6fac0576..703d9b43 100644 --- a/src/main/service-graph/create-lifecycle.ts +++ b/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 { - // 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((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, ]); diff --git a/src/main/worker-service.ts b/src/main/worker-service.ts index 1a6a2d4f..700d355a 100644 --- a/src/main/worker-service.ts +++ b/src/main/worker-service.ts @@ -32,6 +32,7 @@ export function createWorkerService({ let worker: ReturnType | null = null; let destroyed = false; + let shutdownPromise: Promise | null = null; const pendingRetryTimers = new Set>(); /** 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 { + 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((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"); +} diff --git a/src/main/worker/entry/index.ts b/src/main/worker/entry/index.ts index 0cfe3301..c453486e 100644 --- a/src/main/worker/entry/index.ts +++ b/src/main/worker/entry/index.ts @@ -25,6 +25,8 @@ interface ActiveRun { const active = new Map(); const queue: RunJob[] = []; const cancelledRuns = new Set(); +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 { + 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((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); } }); diff --git a/src/main/worker/entry/message-handler.ts b/src/main/worker/entry/message-handler.ts index 48ad061c..258ecf89 100644 --- a/src/main/worker/entry/message-handler.ts +++ b/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; } } diff --git a/src/main/worker/entry/types.ts b/src/main/worker/entry/types.ts index 97d054fd..b086f845 100644 --- a/src/main/worker/entry/types.ts +++ b/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 { diff --git a/src/main/worker/runtime/helpers.ts b/src/main/worker/runtime/helpers.ts index cce4ca0b..f07fb97d 100644 --- a/src/main/worker/runtime/helpers.ts +++ b/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): (text: string) => string { const sensitivePatterns: string[] = []; diff --git a/tests/unit/worker-completion-boundary.test.ts b/tests/unit/worker-completion-boundary.test.ts index 9b57ae9b..c68da7ee 100644 --- a/tests/unit/worker-completion-boundary.test.ts +++ b/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; + }; + 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 { + this.deps.post(runFinishedMessage(job.runId)); + } + + cancel(): void {} + + steer(): boolean { + return false; + } + + async cancelAll(): Promise {} + }, + })); + + 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); + }); }); diff --git a/tests/unit/worker-service-shutdown.test.ts b/tests/unit/worker-service-shutdown.test.ts new file mode 100644 index 00000000..72698174 --- /dev/null +++ b/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; + }; + 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", + })); + }); +});