diff --git a/server/src/__tests__/heartbeat-queued-run-claim-isolation.test.ts b/server/src/__tests__/heartbeat-queued-run-claim-isolation.test.ts new file mode 100644 index 00000000000..82edfab7f69 --- /dev/null +++ b/server/src/__tests__/heartbeat-queued-run-claim-isolation.test.ts @@ -0,0 +1,245 @@ +import { randomUUID } from "node:crypto"; +import { sql } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; +import { + agents, + agentWakeupRequests, + companies, + createDb, + heartbeatRuns, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { heartbeatService } from "../services/heartbeat.ts"; +import { runningProcesses } from "../adapters/index.ts"; + +const mockAdapterExecute = vi.hoisted(() => + vi.fn(async () => ({ + exitCode: 0, + signal: null, + timedOut: false, + errorMessage: null, + summary: "Queued-run claim isolation test run.", + provider: "test", + model: "test-model", + })), +); + +vi.mock("../adapters/index.ts", async () => { + const actual = await vi.importActual("../adapters/index.ts"); + return { + ...actual, + getServerAdapter: vi.fn(() => ({ + supportsLocalAgentJwt: false, + execute: mockAdapterExecute, + })), + }; +}); + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +if (!embeddedPostgresSupport.supported) { + console.warn( + `Skipping embedded Postgres queued-run claim isolation tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, + ); +} + +describeEmbeddedPostgres("heartbeat queued-run claim isolation", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + let heartbeat!: ReturnType; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("heartbeat-queued-run-claim-isolation-"); + db = createDb(tempDb.connectionString); + heartbeat = heartbeatService(db); + }, 20_000); + + afterEach(async () => { + await heartbeat.drainActiveRunExecutions(); + mockAdapterExecute.mockClear(); + runningProcesses.clear(); + // Executed runs write to many company-scoped tables; clear them all. + await db.execute(sql`truncate table ${companies} cascade`); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function insertAgent() { + const companyId = randomUUID(); + const agentId = randomUUID(); + + await db.insert(companies).values({ + id: companyId, + name: "Claim Isolation Co", + status: "active", + issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + + await db.insert(agents).values({ + id: agentId, + companyId, + name: "Claim Isolation Agent", + role: "engineer", + status: "idle", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: { + heartbeat: { + enabled: true, + intervalSec: 60, + wakeOnDemand: true, + maxConcurrentRuns: 1, + }, + }, + permissions: {}, + }); + + return { companyId, agentId }; + } + + async function insertQueuedRun( + companyId: string, + agentId: string, + wake: Partial, + createdAt: Date, + ) { + const wakeupRequestId = randomUUID(); + const runId = randomUUID(); + + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "on_demand", + triggerDetail: "manual", + status: "queued", + runId, + ...wake, + }); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId, + invocationSource: "on_demand", + triggerDetail: "manual", + status: "queued", + wakeupRequestId, + createdAt, + }); + + return { runId, wakeupRequestId }; + } + + // A queued-comment interrupt wake whose receipt cannot be verified. The run + // identity check rejects it with a 403 every time the run is claimed. + function insertUnverifiableInterruptRun(companyId: string, agentId: string, createdAt = new Date(Date.now() - 60_000)) { + return insertQueuedRun(companyId, agentId, { + reason: "issue_commented", + idempotencyKey: `queued-comment-interrupt:${randomUUID()}`, + requestedByActorType: "user", + requestedByActorId: "board-user", + }, createdAt); + } + + // A system wake with no issue in a company without members. Its responsible + // user cannot be resolved yet (422), but a user added later would fix that. + function insertUnresolvableOwnerRun(companyId: string, agentId: string, createdAt = new Date(Date.now() - 60_000)) { + return insertQueuedRun(companyId, agentId, { + source: "automation", + triggerDetail: "system", + requestedByActorType: "system", + requestedByActorId: "heartbeat_test", + }, createdAt); + } + + // A manual wake from a user. It resolves its identity from the receipt, so it + // can always be claimed. + function insertClaimableRun(companyId: string, agentId: string, createdAt = new Date()) { + return insertQueuedRun(companyId, agentId, { + requestedByActorType: "user", + requestedByActorId: "board-user", + payload: { manualUserWake: true }, + }, createdAt); + } + + async function runStatus(runId: string) { + return db + .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode, error: heartbeatRuns.error }) + .from(heartbeatRuns) + .where(sql`${heartbeatRuns.id} = ${runId}`) + .then((rows) => rows[0] ?? null); + } + + it("cancels a queued run whose claim is rejected instead of failing recovery", async () => { + const { companyId, agentId } = await insertAgent(); + const { runId, wakeupRequestId } = await insertUnverifiableInterruptRun(companyId, agentId); + + await expect(heartbeat.resumeQueuedRuns()).resolves.toBeUndefined(); + + expect(await runStatus(runId)).toMatchObject({ + status: "cancelled", + errorCode: "queued_run_claim_rejected", + error: "Cancelled because the queued run cannot be claimed: Queued-message interrupt authority is unavailable", + }); + const wakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(sql`${agentWakeupRequests.id} = ${wakeupRequestId}`) + .then((rows) => rows[0] ?? null); + expect(wakeup).toMatchObject({ status: "cancelled" }); + + // Recovery runs again on the next cycle and on every restart. The run must + // stay settled instead of failing the claim loop again. + await expect(heartbeat.resumeQueuedRuns()).resolves.toBeUndefined(); + expect(await runStatus(runId)).toMatchObject({ status: "cancelled" }); + }); + + it("claims the agent's next queued run after a rejected one", async () => { + const { companyId, agentId } = await insertAgent(); + const { runId: rejectedRunId } = await insertUnverifiableInterruptRun(companyId, agentId); + const { runId: healthyRunId } = await insertClaimableRun(companyId, agentId); + + await expect(heartbeat.resumeQueuedRuns()).resolves.toBeUndefined(); + await heartbeat.drainActiveRunExecutions(); + + expect(await runStatus(rejectedRunId)).toMatchObject({ + status: "cancelled", + errorCode: "queued_run_claim_rejected", + }); + expect(mockAdapterExecute).toHaveBeenCalledOnce(); + expect(await runStatus(healthyRunId)).toMatchObject({ status: "succeeded" }); + }); + + it("keeps a recoverable rejection queued without blocking the runs behind it", async () => { + const { companyId, agentId } = await insertAgent(); + const { runId: deferredRunId } = await insertUnresolvableOwnerRun(companyId, agentId); + const { runId: healthyRunId } = await insertClaimableRun(companyId, agentId); + + await expect(heartbeat.resumeQueuedRuns()).resolves.toBeUndefined(); + await heartbeat.drainActiveRunExecutions(); + + expect(await runStatus(deferredRunId)).toMatchObject({ status: "queued", errorCode: null }); + expect(mockAdapterExecute).toHaveBeenCalledOnce(); + expect(await runStatus(healthyRunId)).toMatchObject({ status: "succeeded" }); + }); + + it("keeps processing other agents' queued runs after one agent's run is rejected", async () => { + const first = await insertAgent(); + const second = await insertAgent(); + const { runId: rejectedRunId } = await insertUnverifiableInterruptRun(first.companyId, first.agentId); + const { runId: healthyRunId } = await insertClaimableRun(second.companyId, second.agentId); + + await expect(heartbeat.resumeQueuedRuns()).resolves.toBeUndefined(); + await heartbeat.drainActiveRunExecutions(); + + expect(await runStatus(rejectedRunId)).toMatchObject({ status: "cancelled" }); + expect(await runStatus(healthyRunId)).toMatchObject({ status: "succeeded" }); + }); +}); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 3d0eb828219..5d8a03ebd33 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -19878,9 +19878,51 @@ export function heartbeatService( } } + // A 403 from claimQueuedRun comes from the run's own persisted identity + // (an unverifiable interrupt receipt, a manual wake with no user). Those rows + // do not change, so the claim fails the same way on every pass and restart. + function isPermanentClaimRejection(err: unknown): err is HttpError { + return err instanceof HttpError && err.status === 403; + } + + // Other 4xx rejections can clear later (a responsible user gets assigned, a + // conflicting claim finishes). Keep the run queued, but do not let it stop + // the rest of the queue or startup recovery. + function isDeferrableClaimRejection(err: unknown): err is HttpError { + return err instanceof HttpError && err.status >= 400 && err.status < 500; + } + + // Settle runs that can never be claimed. Letting the error escape stalls the + // agent's queue and, during startup recovery, stops the server from booting. + async function cancelRejectedQueuedRuns( + rejected: Array<{ run: typeof heartbeatRuns.$inferSelect; err: HttpError }>, + ) { + for (const { run, err } of rejected) { + logger.warn( + { err, runId: run.id, agentId: run.agentId, companyId: run.companyId }, + "cancelling queued heartbeat run whose claim was rejected", + ); + try { + await cancelRunInternal( + run.id, + `Cancelled because the queued run cannot be claimed: ${err.message}`, + { errorCode: "queued_run_claim_rejected" }, + ); + } catch (cancelErr) { + logger.error( + { err: cancelErr, runId: run.id }, + "failed to cancel queued heartbeat run whose claim was rejected; it stays queued for the next recovery pass", + ); + } + } + } + async function startNextQueuedRunForAgent(agentId: string) { if ((await getSchedulingSuppression()).suppressed) return []; const cutoff = await getWorktreeExecutionCutoff(); + // Cancelled after the start lock is released: cancelRunInternal promotes the + // agent's next queued run, which takes this same lock. + const rejectedClaims: Array<{ run: typeof heartbeatRuns.$inferSelect; err: HttpError }> = []; return withAgentStartLock(agentId, async () => { const agent = await getAgent(agentId); @@ -19992,7 +20034,21 @@ export function heartbeatService( const claimedRuns: Array = []; for (const queuedRun of prioritizedRuns) { if (claimedRuns.length >= availableSlots) break; - const claimed = await claimQueuedRun(queuedRun, companyAgents); + let claimed: typeof heartbeatRuns.$inferSelect | null; + try { + claimed = await claimQueuedRun(queuedRun, companyAgents); + } catch (err) { + if (isPermanentClaimRejection(err)) { + rejectedClaims.push({ run: queuedRun, err }); + continue; + } + if (!isDeferrableClaimRejection(err)) throw err; + logger.warn( + { err, runId: queuedRun.id, agentId: queuedRun.agentId, companyId: queuedRun.companyId }, + "queued heartbeat run claim was rejected; leaving it queued for the next recovery pass", + ); + continue; + } if (claimed) claimedRuns.push(claimed); } if (claimedRuns.length === 0) return []; @@ -20016,7 +20072,7 @@ export function heartbeatService( }); } return claimedRuns; - }); + }).finally(() => cancelRejectedQueuedRuns(rejectedClaims)); } // Await every background heartbeat execution that is currently in flight. A