Preserve native follow-up deadlines across drain and resume

Co-Authored-By: Paperclip <noreply@paperclip.ing>

doc/execution-semantics.md

@@ -2696,6 +2696,9 @@ checks and current admission authority still govern the next turn.
 An active remote owner causes startup recovery to retain the obligation at the
 actual lease deadline. The existing recovery dispatch timer is rearmed on drain
 resume; successful reconciliation leaves no recurring recovery work. Upgrade
+Follow-up waits retain their original dispatch mode and owner across a drain,
+including when another heartbeat service facade resumes the work. A failed
+dispatch retains a bounded retry obligation rather than losing its deadline.
 diagnostics retain their existing `ready` meaning and add bounded retained-pointer
 header evidence. That inventory is preliminary, reports partial coverage, and
 does not certify workspace durability or volume removal.

server/src/services/heartbeat.ts

@@ -607,6 +607,9 @@ const nativeSessionResumeDispatchTimers = new Map<
   string,
   ReturnType<typeof setTimeout>
 >();
+// Routes and the scheduler can drain/resume through different service facades.
+// Rearm through the original owner so its database and dispatch mode survive.
+const nativeRecoveryDeadlines = new Map<string, { rearm: () => void }>();
 // Shared with HTTP admission so an idle hold fences work before inspection.
 export { computeTaskDrain, applyTaskDrain, startTaskDrain, stopTaskDrain } from "./task-admission.js";
 
@@ -1335,8 +1338,8 @@ export function heartbeatService(
   }
 
   function resumeQueuedRunsAfterDrain() {
-    for (const [runId, deadline] of nativeRecoveryDeadlines) {
-      scheduleNativeSessionResumeDispatch(runId, deadline, "restart");
+    for (const pending of nativeRecoveryDeadlines.values()) {
+      pending.rearm();
     }
     const promise = heartbeatQueue.resumeQueuedRunsAfterDrain();
     activeWakeupPromises.add(promise);
@@ -1715,13 +1718,13 @@ export function heartbeatService(
 
   // Native ownership waits survive a temporary drain without adding a polling
   // loop. A process restart rediscovers the same deadlines from coordinator rows.
-  const nativeRecoveryDeadlines = new Map<string, Date>();
   function scheduleNativeSessionResumeDispatch(
     runId: string,
     nextAttemptAt: Date,
     mode?: "restart",
   ) {
-    if (mode === "restart") nativeRecoveryDeadlines.set(runId, nextAttemptAt);
+    const pending = { rearm: () => scheduleNativeSessionResumeDispatch(runId, nextAttemptAt, mode) };
+    nativeRecoveryDeadlines.set(runId, pending);
     const prior = nativeSessionResumeDispatchTimers.get(runId);
     if (prior) clearTimeout(prior);
     const delayMs = Math.max(0, nextAttemptAt.getTime() - Date.now());
@@ -1730,8 +1733,9 @@ export function heartbeatService(
       nativeSessionResumeDispatchTimers.delete(runId);
       void (async () => {
         if ((await getSchedulingSuppression()).suppressed) return;
+        if (nativeRecoveryDeadlines.get(runId) !== pending) return;
+        nativeRecoveryDeadlines.delete(runId);
         if (mode === "restart") {
-          nativeRecoveryDeadlines.delete(runId);
           await recoverNativeRunsAfterRestart(new Date(), [runId]);
           return;
         }
@@ -1756,7 +1760,9 @@ export function heartbeatService(
           },
         });
       })().catch((error) => {
-        if (mode === "restart") scheduleNativeSessionResumeDispatch(runId, new Date(Date.now() + 30_000), "restart");
+        if (!nativeRecoveryDeadlines.has(runId)) {
+          scheduleNativeSessionResumeDispatch(runId, new Date(Date.now() + 30_000), mode);
+        }
         logger.error(
           { err: error, runId },
           "failed to dispatch scheduled native session resume",

server/src/services/heartbeat/native-recovery-deadlines.test.ts

@@ -0,0 +1,102 @@
+import type { Db } from "@paperclipai/db";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { heartbeatService, startTaskDrain, stopTaskDrain } from "../heartbeat.js";
+import type { HeartbeatRecoveryDependencies } from "./recovery.js";
+
+const hooks = vi.hoisted(() => ({
+  schedule: undefined as HeartbeatRecoveryDependencies["scheduleNativeSessionResumeDispatch"] | undefined,
+  restart: vi.fn(async () => []),
+  resume: vi.fn(async () => []),
+  queue: vi.fn(async () => undefined),
+}));
+
+vi.mock("./recovery.js", async original => {
+  const module = await original<typeof import("./recovery.js")>();
+  return { ...module, createHeartbeatRecovery: (db: Db, dependencies: HeartbeatRecoveryDependencies) => {
+    hooks.schedule = dependencies.scheduleNativeSessionResumeDispatch;
+    return { ...module.createHeartbeatRecovery(db, dependencies), recoverNativeRunsAfterRestart: hooks.restart };
+  } };
+});
+vi.mock("./queue.js", async original => {
+  const module = await original<typeof import("./queue.js")>();
+  return { ...module, createHeartbeatQueue: (...args: Parameters<typeof module.createHeartbeatQueue>) => ({
+    ...module.createHeartbeatQueue(...args), resumeQueuedRunsAfterDrain: hooks.queue,
+  }) };
+});
+vi.mock("../native-runtime/index.js", async original => ({
+  ...await original<typeof import("../native-runtime/index.js")>(),
+  dispatchNativeSessionResumptions: hooks.resume,
+}));
+
+describe("native deadline drain and resume", () => {
+  let heartbeat: ReturnType<typeof heartbeatService>;
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.useFakeTimers();
+    stopTaskDrain();
+    const db = new Proxy({}, { get: () => { throw new Error("Unexpected deadline database access"); } }) as Db;
+    heartbeat = heartbeatService(db, { runtimeEnv: {} });
+  });
+  afterEach(async () => {
+    stopTaskDrain();
+    await heartbeat?.resumeQueuedRunsAfterDrain();
+    await vi.runOnlyPendingTimersAsync();
+    await heartbeat?.drainActiveRunExecutions();
+    vi.useRealTimers();
+  });
+
+  it.each([undefined, "restart"] as const)("retains an overdue %s wait through drain with its original dispatch mode", async mode => {
+    const deadline = new Date(Date.now() + 100);
+    startTaskDrain({});
+    hooks.schedule!("waiting-run", deadline, mode);
+    await vi.advanceTimersByTimeAsync(100);
+    expect(hooks.restart).not.toHaveBeenCalled();
+    expect(hooks.resume).not.toHaveBeenCalled();
+    stopTaskDrain();
+    await heartbeat.resumeQueuedRunsAfterDrain();
+    await vi.advanceTimersByTimeAsync(0);
+    const called = mode === "restart" ? hooks.restart : hooks.resume;
+    const other = mode === "restart" ? hooks.resume : hooks.restart;
+    expect(called).toHaveBeenCalledOnce();
+    expect(other).not.toHaveBeenCalled();
+    await heartbeat.resumeQueuedRunsAfterDrain();
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(called).toHaveBeenCalledOnce();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("preserves a future follow-up deadline when draining cancels its timer", async () => {
+    const deadline = new Date(Date.now() + 100);
+    hooks.schedule!("future-run", deadline);
+    await heartbeat.drainActiveRunExecutions();
+    await heartbeat.resumeQueuedRunsAfterDrain();
+    await vi.advanceTimersByTimeAsync(99);
+    expect(hooks.resume).not.toHaveBeenCalled();
+    await vi.advanceTimersByTimeAsync(1);
+    expect(hooks.resume).toHaveBeenCalledOnce();
+    expect(hooks.restart).not.toHaveBeenCalled();
+  });
+
+  it("rearms the original owner when another service facade ends the drain", async () => {
+    hooks.schedule!("other-facade-run", new Date(Date.now() + 100));
+    const db = new Proxy({}, { get: () => { throw new Error("Unexpected replacement facade database access"); } }) as Db;
+    const other = heartbeatService(db, { runtimeEnv: {} });
+    await other.drainActiveRunExecutions();
+    await vi.advanceTimersByTimeAsync(100);
+    await other.resumeQueuedRunsAfterDrain();
+    await vi.advanceTimersByTimeAsync(0);
+    expect(hooks.resume).toHaveBeenCalledOnce();
+    expect(hooks.restart).not.toHaveBeenCalled();
+  });
+
+  it("retries a failed follow-up obligation without changing its dispatch mode", async () => {
+    hooks.resume.mockRejectedValueOnce(new Error("temporary database failure"));
+    hooks.schedule!("failed-follow-up", new Date());
+    await vi.advanceTimersByTimeAsync(0);
+    expect(hooks.resume).toHaveBeenCalledOnce();
+    await vi.advanceTimersByTimeAsync(30_000);
+    expect(hooks.resume).toHaveBeenCalledTimes(2);
+    expect(hooks.restart).not.toHaveBeenCalled();
+    expect(vi.getTimerCount()).toBe(0);
+  });
+});