fix(server): bound the shutdown steps before the run drain (#19165)

Bound initial shutdown waits and finish workspace repairs before teardown.

Share a 20-second budget between waits before the heartbeat run drain.
Retain the original workspace recovery stop promise and join it after the
run drain and finalizers, before application and database teardown. Keep
that join for hot restart when the run drain is skipped. Log step and
drain durations and bound telemetry flush separately after the drain.

Verification: all 154 tests across the three relevant shutdown/startup
suites and every server source/test typecheck layer passed. Both actual
server lifecycle variants reproduce premature teardown without the fix.
All applicable CI passed and exact-head Greptile Apex scored 5/5.

Co-Authored-By: Devin Foley <devinfoley@users.noreply.github.com>
Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/server-startup-feedback-export.test.ts

@@ -12,6 +12,8 @@ import { getStartupRecoveryState } from "../startup-recovery-state.js";
 import { idleWorkSnapshot } from "../services/task-admission.js";
 import { readIdleSleepSafety } from "../services/idle-sleep-safety.js";
 import * as idleLocalWork from "../services/idle-local-work.js";
+import * as workspaceRecovery from "../services/workspace-recovery.js";
+import * as workspaceRecoveryWork from "../services/workspace-recovery-work.js";
 import { idleBackupWakeMarker, markIdleBackupWakeRequired } from "../services/idle-database-backup.js";
 import { DELIVERY_QUEUES } from "../services/delivery-work-notifications.js";
 import { decisionService, decisionRetentionService, toolAccessService, githubConnectionEventService, reconcileBuiltInAgentsOnStartup, reconcileCodexLocalManagedHomesOnStartup, reconcilePersistedRuntimeServicesOnStartup } from "../services/index.js";
@@ -1705,6 +1707,71 @@ describe("startServer feedback export wiring", () => {
     } finally { release(); await starting; timer.mockRestore(); }
   });
 
+  it.each([false, true])("joins timed-out workspace repairs before teardown with hot-restart skipDrain=%s", async skipDrain => {
+    loadConfigMock.mockReturnValue(buildTestConfig({ heartbeatSchedulerEnabled: true }));
+    let releaseFirst!: () => void;
+    let releaseSecond!: () => void;
+    const first = new Promise<void>(resolve => { releaseFirst = resolve; });
+    const second = new Promise<void>(resolve => { releaseSecond = resolve; });
+    const scan = vi.spyOn(workspaceRecovery, "reconcileRetainedWorkspaces").mockResolvedValue([first, second]);
+    const register = vi.spyOn(workspaceRecoveryWork, "registerWorkspaceRecoveryWork");
+    const drain = vi.fn(async () => undefined);
+    const heartbeat = { ...heartbeatServiceMock,
+      drainRunningRunsForShutdown: drain,
+      prepareHotRestartShutdown: vi.fn(async () => ({ skipDrain })),
+    };
+    heartbeatServiceFactoryMock.mockReturnValueOnce(heartbeat);
+    const interval = vi.spyOn(globalThis, "setInterval").mockImplementation(() => 1 as unknown as ReturnType<typeof setInterval>);
+    let stopping: Promise<void> | undefined;
+    let stop: ReturnType<typeof vi.spyOn> | undefined;
+    try {
+      const server = await startServer();
+      await settleStartupHousekeeping();
+      expect(scan).toHaveBeenCalledOnce();
+      stop = vi.spyOn(register.mock.results[0]!.value, "stop");
+      const app = await createAppMock.mock.results.at(-1)!.value as {
+        locals: { paperclipShutdown?: () => Promise<void>; paperclipDatabaseClosing?: boolean };
+      };
+      const appShutdown = vi.fn(async () => undefined);
+      app.locals.paperclipShutdown = appShutdown;
+      vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] });
+      let stopped = false;
+      stopping = server.shutdown().then(() => { stopped = true; });
+      await settleStartupHousekeeping();
+      expect(stop).toHaveBeenCalledOnce();
+      await vi.advanceTimersByTimeAsync(20_000);
+      await settleStartupHousekeeping();
+      expect(logger.info).toHaveBeenCalledWith(
+        expect.objectContaining({ step: "workspace_recovery_stop" }),
+        "shutdown step timed out; continuing shutdown without it",
+      );
+      expect(drain).toHaveBeenCalledTimes(skipDrain ? 0 : 1);
+      expect(heartbeatServiceMock.drainActiveRunExecutions).toHaveBeenCalledTimes(skipDrain ? 0 : 1);
+      expect(appShutdown).not.toHaveBeenCalled();
+      expect(app.locals.paperclipDatabaseClosing).toBeUndefined();
+      expect(stopped).toBe(false);
+      releaseFirst();
+      await settleStartupHousekeeping();
+      expect(appShutdown).not.toHaveBeenCalled();
+      expect(stopped).toBe(false);
+      releaseSecond();
+      await stopping;
+      expect(appShutdown).toHaveBeenCalledOnce();
+      expect(app.locals.paperclipDatabaseClosing).toBe(true);
+      expect(stopped).toBe(true);
+      expect(stop).toHaveBeenCalledOnce();
+    } finally {
+      releaseFirst();
+      releaseSecond();
+      await stopping;
+      vi.useRealTimers();
+      stop?.mockRestore();
+      interval.mockRestore();
+      register.mockRestore();
+      scan.mockRestore();
+    }
+  });
+
   it.each(["retained cleanup", "workspace finalization", "environment settlement", "native run projection", "native result commit", "native recovery notices", "external conversation settlement", "native response presentation"] as const)("becomes ready while %s is held but joins it at shutdown", async owner => {
     loadConfigMock.mockReturnValue(buildTestConfig({ heartbeatSchedulerEnabled: true }));
     let releaseCleanup!: () => void;

server/src/index.ts

@@ -200,9 +200,11 @@ import { createDecisionRetentionNotifyOriginAgent, createDecisionWakeOriginAgent
 import {
   closeHttpListenerForShutdown,
   coordinateHeartbeatSchedulerShutdown,
+  createPreDrainShutdownBudget,
   drainRunExecutionFinalizersForShutdown,
   finalizeServerShutdown,
   loadWithoutCoordinatedShutdownSignalHooks,
+  runBoundedShutdownStep,
 } from "./shutdown.js";
 import { initializeCloudRuntimeIdentity } from "./services/cloud-runtime-identity.js";
 import { initializeAdapterSettings } from "./services/adapter-settings.js";
@@ -2019,6 +2021,8 @@ async function startServerWithDatabaseTeardown(
     signal: "SIGINT" | "SIGTERM",
     exitProcess: boolean,
   ) => {
+    const shutdownStartedAt = Date.now();
+    logger.info({ signal }, "shutdown signal received; starting ordered shutdown");
     setStartupRecoveryPhase("draining");
     for (const client of liveEventsWebSocketServer.clients) client.close(1012, "Server restarting");
     liveEventsWebSocketServer.close();
@@ -2026,33 +2030,51 @@ async function startServerWithDatabaseTeardown(
     // are still available. Railway must switch routing before signaling A.
     await closeHttpListenerForShutdown({ server, signal, timeoutMs: 30_000, log: logger,
       drainRequests: abort => drainHttpRequestWork(app, abort) });
-    await systemdNotify(["--stopping", `--status=Stopping after ${signal}`]);
+    const preDrain = createPreDrainShutdownBudget({ signal, log: logger });
+    await preDrain.run("systemd_notify", () =>
+      systemdNotify(["--stopping", `--status=Stopping after ${signal}`]));
     heartbeatSchedulerStopped = true;
     skillVersionRetentionWorker?.unsubscribe();
-    await taskDrainExpiryWork.stop();
+    await preDrain.run("task_drain_expiry_stop", () => taskDrainExpiryWork.stop());
     nativeRunnerEvidenceRetentionWorker?.unsubscribe();
-    await skillSourceSync.stop();
+    await preDrain.run("skill_source_sync_stop", () => skillSourceSync.stop());
     unsubscribeChatCompletions();
     unsubscribeNativeResponsePresentation?.();
     unsubscribeNativeWorkspaceFinalization?.();
     unsubscribeRetainedNativeSessionCleanup?.();
     // B can accept Stop/answers while A still owns an execution. Preserve the
     // owner command workers and their transport until the handoff below ends.
-    await deliveryWork.stop(deploymentHints ? new Set([
+    await preDrain.run("delivery_work_stop", () => deliveryWork.stop(deploymentHints ? new Set([
       DELIVERY_QUEUES.nativeCancellation, DELIVERY_QUEUES.nativeRuntimeResponses, DELIVERY_QUEUES.question,
       DELIVERY_QUEUES.deploymentNotifications, DELIVERY_QUEUES.liveInvalidations,
-    ]) : undefined);
-    await runEvidenceMigrationWork?.stop();
-    await providerTraceRetentionWork.stop();
-    await skillVersionRetentionWorker?.dispose();
-    await workspaceRecoveryWork.stop();
+    ]) : undefined));
+    await preDrain.run("run_evidence_migration_stop", async () => {
+      await runEvidenceMigrationWork?.stop();
+    });
+    await preDrain.run("provider_trace_retention_stop", () => providerTraceRetentionWork.stop());
+    await preDrain.run("skill_version_retention_dispose", async () => {
+      await skillVersionRetentionWorker?.dispose();
+    });
+    const workspaceRecoveryStop = workspaceRecoveryWork.stop();
+    await preDrain.run("workspace_recovery_stop", () => workspaceRecoveryStop);
     unregisterTerminalWorkspaceProbe?.();
     unsubscribeNativeWorkspaceSyncRetention?.();
 
+    const prepareHotRestart = prepareHotRestartShutdown;
     const heartbeatShutdown = await coordinateHeartbeatSchedulerShutdown({
       signal,
-      prepareHotRestartShutdown,
-      waitForHeartbeatSchedulerIdle,
+      prepareHotRestartShutdown: prepareHotRestart
+        ? async (preparationSignal) => {
+          const preparation = await preDrain.run("hot_restart_prepare", () =>
+            prepareHotRestart(preparationSignal));
+          if (preparation.outcome === "completed") return preparation.value;
+          if (preparation.outcome === "failed") throw preparation.error;
+          throw new Error("hot-restart shutdown preparation timed out");
+        }
+        : null,
+      waitForHeartbeatSchedulerIdle: async () => {
+        await preDrain.run("scheduler_idle", waitForHeartbeatSchedulerIdle);
+      },
     });
     const skipHeartbeatDrain = heartbeatShutdown.hotRestart?.skipDrain === true;
     const selectiveDrainRunIds = heartbeatShutdown.hotRestart?.drainRunIds ?? null;
@@ -2068,16 +2090,18 @@ async function startServerWithDatabaseTeardown(
       );
     }
 
-    const telemetryClient = getTelemetryClient();
-    if (telemetryClient) {
-      telemetryClient.stop();
-      await telemetryClient.flush();
-    }
-
     if (!skipHeartbeatDrain && drainHeartbeatRunsForShutdown) {
+      const drainStartedAt = Date.now();
+      logger.info(
+        { signal, sinceSignalMs: drainStartedAt - shutdownStartedAt },
+        "starting graceful heartbeat run drain",
+      );
       try {
         const drain = await drainHeartbeatRunsForShutdown(signal, selectiveDrainRunIds);
-        logger.info({ signal, drain }, "graceful heartbeat run drain complete");
+        logger.info(
+          { signal, drain, durationMs: Date.now() - drainStartedAt },
+          "graceful heartbeat run drain complete",
+        );
       } catch (err) {
         logger.error({ err, signal }, "graceful heartbeat run drain failed");
       }
@@ -2092,6 +2116,20 @@ async function startServerWithDatabaseTeardown(
       });
     }
 
+    await workspaceRecoveryStop;
+
+    const telemetryClient = getTelemetryClient();
+    if (telemetryClient) {
+      telemetryClient.stop();
+      await runBoundedShutdownStep({
+        step: "telemetry_flush",
+        signal,
+        timeoutMs: 5_000,
+        run: () => telemetryClient.flush(),
+        log: logger,
+      });
+    }
+
     await cancellationDelivery?.stop();
     await runtimeResponseDelivery?.stop();
     await deploymentHints?.stop();

server/src/shutdown.test.ts

@@ -3,10 +3,12 @@ import { describe, expect, it, vi } from "vitest";
 import {
   closeHttpListenerForShutdown,
   coordinateHeartbeatSchedulerShutdown,
+  createPreDrainShutdownBudget,
   drainRunExecutionFinalizersForShutdown,
   finalizeServerShutdown,
   finishServiceShutdown,
   loadWithoutCoordinatedShutdownSignalHooks,
+  runBoundedShutdownStep,
 } from "./shutdown.js";
 
 function deferred<T = void>() {
@@ -357,6 +359,131 @@ describe("drainRunExecutionFinalizersForShutdown", () => {
   });
 });
 
+describe("runBoundedShutdownStep", () => {
+  it("returns the step value and logs its duration", async () => {
+    const log = stubLogger();
+    await expect(runBoundedShutdownStep({
+      step: "delivery_work_stop",
+      signal: "SIGTERM",
+      timeoutMs: 1_000,
+      run: async () => "stopped",
+      log,
+    })).resolves.toMatchObject({ outcome: "completed", value: "stopped" });
+    expect(log.info).toHaveBeenCalledWith(
+      expect.objectContaining({ step: "delivery_work_stop", signal: "SIGTERM", durationMs: expect.any(Number) }),
+      "shutdown step completed",
+    );
+  });
+
+  it("stops waiting for a step that outlives its bound", async () => {
+    vi.useFakeTimers();
+    try {
+      const log = stubLogger();
+      const pending = runBoundedShutdownStep({
+        step: "scheduler_idle",
+        signal: "SIGTERM",
+        timeoutMs: 250,
+        run: () => new Promise<void>(() => undefined),
+        log,
+      });
+      await vi.advanceTimersByTimeAsync(250);
+      await expect(pending).resolves.toMatchObject({ outcome: "timed_out" });
+      expect(log.info).toHaveBeenCalledWith(
+        expect.objectContaining({ step: "scheduler_idle", timeoutMs: 250 }),
+        expect.stringContaining("timed out"),
+      );
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("logs a step that rejects and does not reject the shutdown", async () => {
+    const log = stubLogger();
+    const error = new Error("skill sync stop failed");
+    await expect(runBoundedShutdownStep({
+      step: "skill_source_sync_stop",
+      signal: "SIGINT",
+      timeoutMs: 1_000,
+      run: async () => { throw error; },
+      log,
+    })).resolves.toMatchObject({ outcome: "failed", error });
+    expect(log.error).toHaveBeenCalledWith(
+      expect.objectContaining({ step: "skill_source_sync_stop", err: error }),
+      expect.stringContaining("failed"),
+    );
+  });
+
+  it("contains a step that throws synchronously", async () => {
+    const result = await runBoundedShutdownStep({
+      step: "systemd_notify",
+      signal: "SIGTERM",
+      timeoutMs: 1_000,
+      run: () => { throw new Error("spawn failed"); },
+      log: stubLogger(),
+    });
+    expect(result.outcome).toBe("failed");
+  });
+});
+
+describe("createPreDrainShutdownBudget", () => {
+  it("gives each step only the time that is left in the shared budget", async () => {
+    vi.useFakeTimers();
+    try {
+      const log = stubLogger();
+      const budget = createPreDrainShutdownBudget({ signal: "SIGTERM", budgetMs: 1_000, log });
+      const first = budget.run("skill_source_sync_stop", () =>
+        new Promise<void>((resolve) => setTimeout(resolve, 700)));
+      await vi.advanceTimersByTimeAsync(700);
+      await expect(first).resolves.toMatchObject({ outcome: "completed" });
+
+      const second = budget.run("scheduler_idle", () => new Promise<void>(() => undefined));
+      await vi.advanceTimersByTimeAsync(300);
+      await expect(second).resolves.toMatchObject({ outcome: "timed_out" });
+      expect(log.info).toHaveBeenCalledWith(
+        expect.objectContaining({ step: "scheduler_idle", timeoutMs: 300 }),
+        expect.stringContaining("timed out"),
+      );
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("reaches the run drain within the budget when every pre-drain step hangs", async () => {
+    vi.useFakeTimers();
+    try {
+      const budget = createPreDrainShutdownBudget({ signal: "SIGTERM", budgetMs: 1_000, log: stubLogger() });
+      const hang = () => new Promise<void>(() => undefined);
+      const drain = vi.fn(async () => ({ interrupted: 1 }));
+      const shutdown = (async () => {
+        await budget.run("skill_source_sync_stop", hang);
+        await budget.run("delivery_work_stop", hang);
+        await coordinateHeartbeatSchedulerShutdown({
+          signal: "SIGTERM",
+          prepareHotRestartShutdown: async () => {
+            const preparation = await budget.run("hot_restart_prepare", () =>
+              new Promise<{ skipDrain: boolean }>(() => undefined));
+            if (preparation.outcome === "completed") return preparation.value;
+            throw new Error("hot-restart shutdown preparation timed out");
+          },
+          waitForHeartbeatSchedulerIdle: async () => {
+            await budget.run("scheduler_idle", hang);
+          },
+        });
+        return drain();
+      })();
+
+      const startedAt = Date.now();
+      await vi.advanceTimersByTimeAsync(1_000);
+      await vi.runAllTimersAsync();
+      await expect(shutdown).resolves.toEqual({ interrupted: 1 });
+      expect(drain).toHaveBeenCalledOnce();
+      expect(Date.now() - startedAt).toBeLessThan(1_010);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});
+
 describe("loadWithoutCoordinatedShutdownSignalHooks", () => {
   it("removes the eager signal handlers from the real embedded-postgres import", async () => {
     const before = {

server/src/shutdown.ts

@@ -7,6 +7,77 @@ type ShutdownLogger = {
   error(obj: object, msg: string): void;
 };
 
+export const PRE_DRAIN_SHUTDOWN_BUDGET_MS = 20_000;
+
+export type ShutdownStepResult<T> =
+  | { outcome: "completed"; value: T; durationMs: number }
+  | { outcome: "timed_out"; durationMs: number }
+  | { outcome: "failed"; error: unknown; durationMs: number };
+
+export async function runBoundedShutdownStep<T>(input: {
+  step: string;
+  signal: "SIGINT" | "SIGTERM";
+  timeoutMs: number;
+  run: () => Promise<T>;
+  log: ShutdownLogger;
+  now?: () => number;
+}): Promise<ShutdownStepResult<T>> {
+  const now = input.now ?? Date.now;
+  const startedAt = now();
+  const timeoutMs = Math.max(0, input.timeoutMs);
+  let timer: NodeJS.Timeout | null = null;
+  try {
+    const result = await Promise.race([
+      Promise.resolve()
+        .then(input.run)
+        .then(
+          (value) => ({ outcome: "completed" as const, value }),
+          (error: unknown) => ({ outcome: "failed" as const, error }),
+        ),
+      new Promise<{ outcome: "timed_out" }>((resolve) => {
+        timer = setTimeout(() => resolve({ outcome: "timed_out" }), timeoutMs);
+        timer.unref?.();
+      }),
+    ]);
+    const durationMs = now() - startedAt;
+    const fields = { signal: input.signal, step: input.step, durationMs, timeoutMs };
+    if (result.outcome === "completed") {
+      input.log.info(fields, "shutdown step completed");
+      return { outcome: "completed", value: result.value, durationMs };
+    }
+    if (result.outcome === "failed") {
+      input.log.error({ ...fields, err: result.error }, "shutdown step failed; continuing shutdown");
+      return { outcome: "failed", error: result.error, durationMs };
+    }
+    input.log.info(fields, "shutdown step timed out; continuing shutdown without it");
+    return { outcome: "timed_out", durationMs };
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
+}
+
+export function createPreDrainShutdownBudget(input: {
+  signal: "SIGINT" | "SIGTERM";
+  budgetMs?: number;
+  log: ShutdownLogger;
+  now?: () => number;
+}) {
+  const now = input.now ?? Date.now;
+  const deadline = now() + (input.budgetMs ?? PRE_DRAIN_SHUTDOWN_BUDGET_MS);
+  return {
+    run<T>(step: string, run: () => Promise<T>) {
+      return runBoundedShutdownStep({
+        step,
+        signal: input.signal,
+        timeoutMs: deadline - now(),
+        run,
+        log: input.log,
+        now,
+      });
+    },
+  };
+}
+
 export async function drainRunExecutionFinalizersForShutdown(input: {
   signal: "SIGINT" | "SIGTERM";
   drain: (() => Promise<void>) | null;