diff --git a/doc/execution-semantics.md b/doc/execution-semantics.md index e44d4da8e00..3fc99801642 100644 --- a/doc/execution-semantics.md +++ b/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. diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 4f69300ef43..ffa4d03d046 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -607,6 +607,9 @@ const nativeSessionResumeDispatchTimers = new Map< string, ReturnType >(); +// 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 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(); 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", diff --git a/server/src/services/heartbeat/native-recovery-deadlines.test.ts b/server/src/services/heartbeat/native-recovery-deadlines.test.ts new file mode 100644 index 00000000000..e39daad1b52 --- /dev/null +++ b/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(); + 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(); + return { ...module, createHeartbeatQueue: (...args: Parameters) => ({ + ...module.createHeartbeatQueue(...args), resumeQueuedRunsAfterDrain: hooks.queue, + }) }; +}); +vi.mock("../native-runtime/index.js", async original => ({ + ...await original(), + dispatchNativeSessionResumptions: hooks.resume, +})); + +describe("native deadline drain and resume", () => { + let heartbeat: ReturnType; + 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); + }); +});