Prevent stale heartbeat start locks

server/src/__tests__/heartbeat-start-lock.test.ts

@@ -0,0 +1,30 @@
+import { randomUUID } from "node:crypto";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { withAgentStartLockForTest } from "../services/heartbeat.ts";
+
+describe("heartbeat agent start lock", () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("does not let a stale start lock freeze later queued-run starts", async () => {
+    vi.useFakeTimers();
+
+    const agentId = randomUUID();
+    const firstStart = vi.fn(() => new Promise<void>(() => undefined));
+    const secondStart = vi.fn(async () => "started");
+
+    void withAgentStartLockForTest(agentId, firstStart);
+    await Promise.resolve();
+    expect(firstStart).toHaveBeenCalledTimes(1);
+
+    const secondStartResult = withAgentStartLockForTest(agentId, secondStart);
+    await Promise.resolve();
+    expect(secondStart).not.toHaveBeenCalled();
+
+    await vi.advanceTimersByTimeAsync(30_000);
+
+    await expect(secondStartResult).resolves.toBe("started");
+    expect(secondStart).toHaveBeenCalledTimes(1);
+  });
+});

server/src/services/heartbeat.ts

@@ -137,6 +137,7 @@ const MAX_RUN_EVENT_PAYLOAD_OBJECT_KEYS = 100;
 const MAX_RUN_EVENT_PAYLOAD_DEPTH = 6;
 const HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT = AGENT_DEFAULT_MAX_CONCURRENT_RUNS;
 const HEARTBEAT_MAX_CONCURRENT_RUNS_MAX = 10;
+const AGENT_START_LOCK_STALE_MS = 30_000;
 const LIVENESS_BOOKKEEPING_ACTIVITY_ACTIONS = [
   "environment.lease_acquired",
   "environment.lease_released",
@@ -146,7 +147,7 @@ const WAKE_COMMENT_IDS_KEY = "wakeCommentIds";
 const PAPERCLIP_WAKE_PAYLOAD_KEY = "paperclipWake";
 const PAPERCLIP_HARNESS_CHECKOUT_KEY = "paperclipHarnessCheckedOut";
 const DETACHED_PROCESS_ERROR_CODE = "process_detached";
-const startLocksByAgent = new Map<string, Promise<void>>();
+const startLocksByAgent = new Map<string, { promise: Promise<void>; startedAtMs: number }>();
 const REPO_ONLY_CWD_SENTINEL = "/__paperclip_repo_only__";
 const MANAGED_WORKSPACE_GIT_CLONE_TIMEOUT_MS = 10 * 60 * 1000;
 const MAX_INLINE_WAKE_COMMENTS = 8;
@@ -785,18 +786,49 @@ function normalizeMaxConcurrentRuns(value: unknown) {
   return Math.max(HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT, Math.min(HEARTBEAT_MAX_CONCURRENT_RUNS_MAX, parsed));
 }
 
+async function waitForAgentStartLock(agentId: string, lock: { promise: Promise<void>; startedAtMs: number }) {
+  const elapsedMs = Date.now() - lock.startedAtMs;
+  const remainingMs = AGENT_START_LOCK_STALE_MS - elapsedMs;
+  if (remainingMs <= 0) {
+    logger.warn({ agentId, staleMs: elapsedMs }, "agent start lock stale; continuing queued-run start");
+    return;
+  }
+
+  let timedOut = false;
+  let timeout: ReturnType<typeof setTimeout> | null = null;
+  await Promise.race([
+    lock.promise,
+    new Promise<void>((resolve) => {
+      timeout = setTimeout(() => {
+        timedOut = true;
+        resolve();
+      }, remainingMs);
+    }),
+  ]);
+  if (timeout) clearTimeout(timeout);
+
+  if (timedOut) {
+    logger.warn({ agentId, staleMs: AGENT_START_LOCK_STALE_MS }, "agent start lock timed out; continuing queued-run start");
+  }
+}
+
+export async function withAgentStartLockForTest<T>(agentId: string, fn: () => Promise<T>) {
+  return withAgentStartLock(agentId, fn);
+}
+
 async function withAgentStartLock<T>(agentId: string, fn: () => Promise<T>) {
-  const previous = startLocksByAgent.get(agentId) ?? Promise.resolve();
-  const run = previous.then(fn);
+  const previous = startLocksByAgent.get(agentId);
+  const waitForPrevious = previous ? waitForAgentStartLock(agentId, previous) : Promise.resolve();
+  const run = waitForPrevious.then(fn);
   const marker = run.then(
     () => undefined,
     () => undefined,
   );
-  startLocksByAgent.set(agentId, marker);
+  startLocksByAgent.set(agentId, { promise: marker, startedAtMs: Date.now() });
   try {
     return await run;
   } finally {
-    if (startLocksByAgent.get(agentId) === marker) {
+    if (startLocksByAgent.get(agentId)?.promise === marker) {
       startLocksByAgent.delete(agentId);
     }
   }