fix: bound controller lease checks by the abort deadline

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

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({});
   });

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<never>((_, 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();