diff --git a/server/src/services/legacy-controller-lease.test.ts b/server/src/services/legacy-controller-lease.test.ts index e9224cc71b..0eb33ccec9 100644 --- a/server/src/services/legacy-controller-lease.test.ts +++ b/server/src/services/legacy-controller-lease.test.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; import { eq, sql } from "drizzle-orm"; -import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { agents, companies, createDb, heartbeatRuns } from "@paperclipai/db"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; import { heartbeatService } from "./heartbeat.js"; @@ -91,6 +91,20 @@ const support = await getEmbeddedPostgresTestSupport(); expect(await renewLegacyControllerLease(db, run)).toBe(false); expect(await revokeExpiredLegacyController(db, run)).toBe(false); }); + it("rejects dispatch at the lease deadline even if the database query never settles", async () => { + const run = await seed(); + const hungDb = { update: () => ({ set: () => ({ where: () => ({ returning: () => new Promise(() => {}) }) }) }) } as unknown as typeof db; + vi.useFakeTimers(); + const controller = new AbortController(); + const watch = watchLegacyControllerLease(hungDb, { ...run, controllerLeaseExpiresAt: new Date(Date.now() + 100) }, controller); + try { + const checked = expect(watch.assertOwned("dispatching")).rejects.toThrow("lease lost"); + await vi.advanceTimersByTimeAsync(101); + await checked; + expect(controller.signal.aborted).toBe(true); + } finally { watch.stop(); vi.useRealTimers(); } + }); + it("leaves native controller ownership to the native coordinator", () => { expect(legacyControllerClaim("native")).toEqual({}); }); diff --git a/server/src/services/legacy-controller-lease.ts b/server/src/services/legacy-controller-lease.ts index fc5a3dc71c..4fc746fe48 100644 --- a/server/src/services/legacy-controller-lease.ts +++ b/server/src/services/legacy-controller-lease.ts @@ -70,7 +70,7 @@ export function watchLegacyControllerLease(db: Db, run: Run, controller: AbortCo } let stopped = false; let pending = false; - const lost = () => controller.abort(new Error("Legacy controller lease lost")); + const lost = () => { if (!stopped) controller.abort(new Error("Legacy controller lease lost")); }; let deadline = setTimeout(lost, Math.max(0, (run.controllerLeaseExpiresAt?.getTime() ?? 0) - Date.now())); deadline.unref(); @@ -78,7 +78,17 @@ export function watchLegacyControllerLease(db: Db, run: Run, controller: AbortCo if (stopped) return; controller.signal.throwIfAborted(); const startedAt = Date.now(); - const renewed = await renewLegacyControllerLease(db, run, stage); + let onAbort!: () => void; + const aborted = new Promise((_, reject) => { + onAbort = () => reject(controller.signal.reason); + controller.signal.addEventListener("abort", onAbort, { once: true }); + }); + let renewed: boolean; + try { + renewed = await Promise.race([renewLegacyControllerLease(db, run, stage), aborted]); + } finally { + controller.signal.removeEventListener("abort", onAbort); + } if (stopped) return; if (!renewed) { lost();