diff --git a/server/src/__tests__/heartbeat-start-lock.test.ts b/server/src/__tests__/heartbeat-start-lock.test.ts new file mode 100644 index 00000000000..dc52673f2d6 --- /dev/null +++ b/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(() => 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); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 8dbc3b74b82..ffb05b1d618 100644 --- a/server/src/services/heartbeat.ts +++ b/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>(); +const startLocksByAgent = new Map; 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; 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 | null = null; + await Promise.race([ + lock.promise, + new Promise((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(agentId: string, fn: () => Promise) { + return withAgentStartLock(agentId, fn); +} + async function withAgentStartLock(agentId: string, fn: () => Promise) { - 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); } }