diff --git a/server/src/__tests__/heartbeat-company-wip-limit-race.test.ts b/server/src/__tests__/heartbeat-company-wip-limit-race.test.ts new file mode 100644 index 000000000000..90b2c33c7f74 --- /dev/null +++ b/server/src/__tests__/heartbeat-company-wip-limit-race.test.ts @@ -0,0 +1,466 @@ +import { randomUUID } from "node:crypto"; +import { eq, sql } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; +import { + activityLog, + agents, + agentRuntimeState, + agentWakeupRequests, + companySkills, + companies, + createDb, + documentRevisions, + documents, + environmentLeases, + environments, + executionWorkspaces, + heartbeatRunEvents, + heartbeatRuns, + issueComments, + issueDocuments, + issueTreeHolds, + issues, + workspaceOperations, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { heartbeatService } from "../services/heartbeat.ts"; +import { withCompanyWipSlot } from "../services/company-wip-limit.ts"; +import { runningProcesses } from "../adapters/index.ts"; + +/** + * The company WIP limit has to be a reservation, not a count. + * + * A test that calls the gate twice in sequence passes against the broken code and + * proves nothing: the bug only exists while two claims overlap. Both tests here + * force a genuine overlap — the first claim is pinned open on a deferred while the + * second one runs — which is the only shape that can tell a reservation apart from + * an observation. + */ + +const mockAdapterExecute = vi.hoisted(() => + vi.fn(async () => ({ + exitCode: 0, + signal: null, + timedOut: false, + errorMessage: null, + summary: "Company WIP limit race 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 company WIP limit race tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, + ); +} + +function deferred() { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +} + +/** Settles to "pending" unless `promise` wins within a few macrotask turns. */ +async function settledWithin(promise: Promise, turns = 10): Promise<"pending" | { value: T }> { + const pending = Symbol("pending"); + for (let turn = 0; turn < turns; turn += 1) { + const winner = await Promise.race([ + promise.then((value) => ({ value })), + new Promise((resolve) => setTimeout(() => resolve(pending), 10)), + ]); + if (winner !== pending) return winner as { value: T }; + } + return "pending"; +} + +async function waitForCondition(fn: () => Promise, timeoutMs = 5_000) { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await fn()) return true; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + return fn(); +} + +describeEmbeddedPostgres("company WIP limit is a reservation, not a count", () => { + let db!: ReturnType; + let heartbeat!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("paperclip-heartbeat-company-wip-race-"); + db = createDb(tempDb.connectionString); + heartbeat = heartbeatService(db, { companyMaxConcurrentRuns: 1 }); + }, 20_000); + + afterEach(async () => { + let idlePolls = 0; + for (let attempt = 0; attempt < 100; attempt += 1) { + const runs = await db.select({ status: heartbeatRuns.status }).from(heartbeatRuns); + const hasActiveRun = runs.some((run) => run.status === "queued" || run.status === "running"); + if (!hasActiveRun) { + idlePolls += 1; + if (idlePolls >= 3) break; + } else { + idlePolls = 0; + } + await new Promise((resolve) => setTimeout(resolve, 50)); + } + const runIds = await db + .select({ id: heartbeatRuns.id }) + .from(heartbeatRuns) + .then((rows) => rows.map((row) => row.id)); + await Promise.all(runIds.map((runId) => heartbeat.waitForRunExecutionDrain(runId))); + mockAdapterExecute.mockReset(); + mockAdapterExecute.mockImplementation(async () => ({ + exitCode: 0, + signal: null, + timedOut: false, + errorMessage: null, + summary: "Company WIP limit race test run.", + provider: "test", + model: "test-model", + })); + runningProcesses.clear(); + await db.delete(environmentLeases); + await db.delete(activityLog); + await db.delete(companySkills); + await db.delete(issueComments); + await db.delete(issueDocuments); + await db.delete(documentRevisions); + await db.delete(documents); + await db.delete(issueTreeHolds); + await db.delete(issues); + await db.delete(heartbeatRunEvents); + await db.delete(heartbeatRuns); + await db.delete(agentWakeupRequests); + await db.delete(agentRuntimeState); + await db.delete(agents); + await db.delete(environments); + await db.delete(workspaceOperations); + await db.delete(executionWorkspaces); + for (let attempt = 0; attempt < 5; attempt += 1) { + try { + await db.transaction(async (tx) => { + await tx.delete(companySkills); + await tx.delete(companies); + }); + break; + } catch (error) { + if (attempt === 4) throw error; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + } + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seedCompany(companyId: string, prefixSeed: string) { + await db.insert(companies).values({ + id: companyId, + name: "Paperclip", + issuePrefix: `${prefixSeed}${companyId.replace(/-/g, "").slice(0, 5).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + defaultResponsibleUserId: "responsible-user", + }); + } + + async function seedAgent(companyId: string, agentId: string, name: string, role: string) { + await db.insert(agents).values({ + id: agentId, + companyId, + name, + role, + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: { heartbeat: { wakeOnDemand: true, maxConcurrentRuns: 1 } }, + permissions: {}, + }); + } + + it("holds the second claimant outside the company slot until the first claim commits", async () => { + const companyId = randomUUID(); + const firstAgentId = randomUUID(); + const secondAgentId = randomUUID(); + const firstRunId = randomUUID(); + const secondRunId = randomUUID(); + + await seedCompany(companyId, "R"); + await seedAgent(companyId, firstAgentId, "Implementer", "engineer"); + await seedAgent(companyId, secondAgentId, "Reviewer", "qa"); + await db.insert(heartbeatRuns).values([ + { id: firstRunId, companyId, agentId: firstAgentId, status: "queued", invocationSource: "on_demand" }, + { id: secondRunId, companyId, agentId: secondAgentId, status: "queued", invocationSource: "on_demand" }, + ]); + + const holdFirstClaim = deferred(); + let secondClaimRan = false; + + // Claimant 1 parks *before* writing its running row. That is the shape of the + // real bug: in claimQueuedRun about eleven database round trips sit between the + // slot decision and the UPDATE, so a second claimant gets to look at the world + // while the first has decided but not yet written. A count-based gate lets the + // second one through here; a reservation must not. + const firstClaim = withCompanyWipSlot({ db, companyId, limit: 1 }, async (tx) => { + await holdFirstClaim.promise; + return tx + .update(heartbeatRuns) + .set({ status: "running", startedAt: new Date(), updatedAt: new Date() }) + .where(eq(heartbeatRuns.id, firstRunId)) + .returning() + .then((rows) => rows[0] ?? null); + }); + + // Wait until claimant 1 really holds the advisory lock, so the overlap is real + // and not a scheduling accident. + await waitForCondition(async () => { + const locks = await db.execute( + sql`select count(*)::int as count from pg_locks where locktype = 'advisory' and granted`, + ); + return Number((locks as unknown as Array<{ count: number }>)[0]?.count ?? 0) > 0; + }); + + const secondClaim = withCompanyWipSlot({ db, companyId, limit: 1 }, async (tx) => { + secondClaimRan = true; + return tx + .update(heartbeatRuns) + .set({ status: "running", startedAt: new Date(), updatedAt: new Date() }) + .where(eq(heartbeatRuns.id, secondRunId)) + .returning() + .then((rows) => rows[0] ?? null); + }); + + // The gate is a reservation: claimant 2 cannot even decide yet, because + // claimant 1 holds the lock and has not committed its running row. + expect(await settledWithin(secondClaim)).toBe("pending"); + expect(secondClaimRan).toBe(false); + + holdFirstClaim.resolve(); + await expect(firstClaim).resolves.toMatchObject({ id: firstRunId, status: "running" }); + + // Now it can decide, and the answer is "no slot" — without ever running its claim. + await expect(secondClaim).resolves.toBeNull(); + expect(secondClaimRan).toBe(false); + + const statuses = await db + .select({ id: heartbeatRuns.id, status: heartbeatRuns.status }) + .from(heartbeatRuns) + .then((rows) => Object.fromEntries(rows.map((row) => [row.id, row.status]))); + expect(statuses[firstRunId]).toBe("running"); + expect(statuses[secondRunId]).toBe("queued"); + }, 30_000); + + it("starts exactly one run when two agents of one company wake at the same instant", async () => { + const companyId = randomUUID(); + const implementerId = randomUUID(); + const reviewerId = randomUUID(); + const implementerIssueId = randomUUID(); + const reviewerIssueId = randomUUID(); + + await seedCompany(companyId, "E"); + await seedAgent(companyId, implementerId, "Implementer", "engineer"); + await seedAgent(companyId, reviewerId, "Reviewer", "qa"); + await db.insert(issues).values([ + { + id: implementerIssueId, + companyId, + title: "Implement the thing", + status: "todo", + priority: "high", + assigneeAgentId: implementerId, + responsibleUserId: "responsible-user", + }, + { + id: reviewerIssueId, + companyId, + title: "Review the thing", + status: "todo", + priority: "high", + assigneeAgentId: reviewerId, + responsibleUserId: "responsible-user", + }, + ]); + + // Whoever wins parks in the adapter, so the winner stays `running` for the + // whole assertion window and the company slot is genuinely occupied. + const releaseAdapter = deferred(); + mockAdapterExecute.mockImplementation(async () => { + await releaseAdapter.promise; + return { + exitCode: 0, + signal: null, + timedOut: false, + errorMessage: null, + summary: "Company WIP limit race test run.", + provider: "test", + model: "test-model", + }; + }); + + try { + const [implementerRun, reviewerRun] = await Promise.all([ + heartbeat.wakeup(implementerId, { + source: "assignment", + triggerDetail: "system", + reason: "issue_assigned", + payload: { issueId: implementerIssueId }, + contextSnapshot: { issueId: implementerIssueId, wakeReason: "issue_assigned" }, + }), + heartbeat.wakeup(reviewerId, { + source: "assignment", + triggerDetail: "system", + reason: "issue_assigned", + payload: { issueId: reviewerIssueId }, + contextSnapshot: { issueId: reviewerIssueId, wakeReason: "issue_assigned" }, + }), + ]); + expect(implementerRun).not.toBeNull(); + expect(reviewerRun).not.toBeNull(); + + const oneRunning = await waitForCondition(async () => { + const rows = await db.select({ status: heartbeatRuns.status }).from(heartbeatRuns); + return rows.filter((row) => row.status === "running").length === 1; + }); + expect(oneRunning).toBe(true); + + // Let the winner reach the adapter, then let any second claim that was going + // to happen, happen. + const winnerReachedAdapter = await waitForCondition( + async () => mockAdapterExecute.mock.calls.length >= 1, + 30_000, + ); + expect(winnerReachedAdapter).toBe(true); + await new Promise((resolve) => setTimeout(resolve, 500)); + + const rows = await db + .select({ id: heartbeatRuns.id, agentId: heartbeatRuns.agentId, status: heartbeatRuns.status }) + .from(heartbeatRuns); + expect(rows).toHaveLength(2); + const running = rows.filter((row) => row.status === "running"); + const queued = rows.filter((row) => row.status === "queued"); + expect(running).toHaveLength(1); + expect(queued).toHaveLength(1); + expect(running[0]!.agentId).not.toBe(queued[0]!.agentId); + expect(mockAdapterExecute).toHaveBeenCalledTimes(1); + } finally { + releaseAdapter.resolve(); + } + }, 40_000); + + it("re-dispatches the losing agent company-wide when the winning run finishes", async () => { + const companyId = randomUUID(); + const implementerId = randomUUID(); + const reviewerId = randomUUID(); + const implementerIssueId = randomUUID(); + const reviewerIssueId = randomUUID(); + + await seedCompany(companyId, "S"); + await seedAgent(companyId, implementerId, "Implementer", "engineer"); + await seedAgent(companyId, reviewerId, "Reviewer", "qa"); + await db.insert(issues).values([ + { + id: implementerIssueId, + companyId, + title: "Implement the thing", + status: "todo", + priority: "high", + assigneeAgentId: implementerId, + responsibleUserId: "responsible-user", + }, + { + id: reviewerIssueId, + companyId, + title: "Review the thing", + status: "todo", + priority: "high", + assigneeAgentId: reviewerId, + responsibleUserId: "responsible-user", + }, + ]); + + const releaseWinner = deferred(); + mockAdapterExecute.mockImplementationOnce(async () => { + await releaseWinner.promise; + return { + exitCode: 0, + signal: null, + timedOut: false, + errorMessage: null, + summary: "Winning run complete.", + provider: "test", + model: "test-model", + }; + }); + + try { + await Promise.all([ + heartbeat.wakeup(implementerId, { + source: "assignment", + triggerDetail: "system", + reason: "issue_assigned", + payload: { issueId: implementerIssueId }, + contextSnapshot: { issueId: implementerIssueId, wakeReason: "issue_assigned" }, + }), + heartbeat.wakeup(reviewerId, { + source: "assignment", + triggerDetail: "system", + reason: "issue_assigned", + payload: { issueId: reviewerIssueId }, + contextSnapshot: { issueId: reviewerIssueId, wakeReason: "issue_assigned" }, + }), + ]); + + await waitForCondition(async () => { + const rows = await db.select({ status: heartbeatRuns.status }).from(heartbeatRuns); + return rows.filter((row) => row.status === "running").length === 1; + }); + + const loserAgentId = await db + .select({ agentId: heartbeatRuns.agentId }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.status, "queued")) + .then((rows) => rows[0]?.agentId ?? null); + expect(loserAgentId).not.toBeNull(); + + releaseWinner.resolve(); + + // No timer is running in this test, so the only thing that can move the loser + // off `queued` is the company-wide sweep on run completion. + const loserDispatched = await waitForCondition(async () => { + const row = await db + .select({ status: heartbeatRuns.status }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.agentId, loserAgentId!)) + .then((rows) => rows[0] ?? null); + return row?.status !== "queued"; + }, 15_000); + expect(loserDispatched).toBe(true); + expect(mockAdapterExecute.mock.calls.length).toBeGreaterThanOrEqual(2); + } finally { + releaseWinner.resolve(); + } + }, 60_000); +}); diff --git a/server/src/services/company-wip-limit.ts b/server/src/services/company-wip-limit.ts index 7ca7803b5dd0..09d765d3ebbd 100644 --- a/server/src/services/company-wip-limit.ts +++ b/server/src/services/company-wip-limit.ts @@ -49,7 +49,9 @@ export async function countRunningRunsForCompany(db: Db, companyId: string): Pro } /** - * Single entry point used by the scheduler: may this company start another run? + * Fast path used by the scheduler before it does any work: is it worth looking + * for a queued run at all? Advisory only — it is not the gate. See + * `withCompanyWipSlot`. */ export async function companyHasFreeWipSlot(input: { db: Db; @@ -60,3 +62,53 @@ export async function companyHasFreeWipSlot(input: { const running = await countRunningRunsForCompany(input.db, input.companyId); return hasCompanySlot(input.limit, running); } + +/** Advisory-lock key. One lock per company, so companies never contend. */ +function companyWipLockKey(companyId: string): string { + return `paperclip:company-wip:${companyId}`; +} + +/** + * Reserves a company slot for the duration of `claim`. + * + * A bare count is not a gate. The scheduler counts running runs, then does about + * eleven more database round trips (invokability, budgets, daily caps, tree + * holds, dependency readiness, staleness) before the conditional UPDATE that + * flips a run to `running`. The only mutex in that window is `withAgentStartLock`, + * which is in-process and keyed by agent id — two agents of one company take two + * different locks and both sail through. Both read zero running runs, both claim, + * and the company spends two runs where the operator configured one. It fails on + * money, so the check has to be a reservation, not an observation. + * + * Pushing `(select count(*) …) < limit` into the UPDATE's WHERE does not fix it: + * under READ COMMITTED the two implicit transactions each take their snapshot + * before the other commits, so both subqueries still see zero. The serialization + * point has to be a lock. `pg_advisory_xact_lock` is the right one — it lives in + * the database, so it holds across server processes, and Postgres drops it when + * the transaction ends, including when a process dies, which a counter table + * would not do. + * + * The transaction body is deliberately only three statements: take the lock, + * re-count under it, run the caller's UPDATE. No budget checks, no adapter calls. + * A company-wide lock held across slow work would serialize the whole company and + * could deadlock against the longer-lived `paperclip:folders:*` locks. + * + * Returns null when the company has no free slot; the caller's run stays queued. + * When the limit is disabled this touches neither a transaction nor the database + * and hands the caller the plain connection, so upstream behaviour is unchanged. + */ +export async function withCompanyWipSlot( + input: { db: Db; companyId: string; limit: number }, + claim: (tx: Db) => Promise, +): Promise { + if (!isCompanyWipLimitEnabled(input.limit)) return claim(input.db); + return input.db.transaction(async (tx) => { + const txDb = tx as unknown as Db; + await txDb.execute( + sql`select pg_advisory_xact_lock(hashtextextended(${companyWipLockKey(input.companyId)}, 0))`, + ); + const running = await countRunningRunsForCompany(txDb, input.companyId); + if (!hasCompanySlot(input.limit, running)) return null; + return claim(txDb); + }); +} diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index 0d03b0349ddd..132d0867deec 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -70,7 +70,12 @@ import { import { conflict, HttpError, notFound } from "../errors.js"; import { getStartupTraceContext } from "../instrumentation.js"; import { logger } from "../middleware/logger.js"; -import { companyHasFreeWipSlot, normalizeCompanyWipLimit } from "./company-wip-limit.js"; +import { + companyHasFreeWipSlot, + isCompanyWipLimitEnabled, + normalizeCompanyWipLimit, + withCompanyWipSlot, +} from "./company-wip-limit.js"; import { createGitRemoteAuthProvider, describeGitAuthFailure, @@ -12483,17 +12488,26 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) issueContext: issueId ? await getIssueExecutionContext(run.companyId, issueId) : null, routineEnvContext: { routineId: null, env: null, responsibleUserId: null }, }); - const claimed = await db - .update(heartbeatRuns) - .set({ - status: "running", - responsibleUserId, - startedAt: run.startedAt ?? claimedAt, - updatedAt: claimedAt, - }) - .where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "queued"))) - .returning() - .then((rows) => rows[0] ?? null); + // The company-wide slot is reserved here, not at the earlier count: the check + // at startNextQueuedRunForAgent is a fast path, and ~11 round trips separate + // it from this UPDATE. withCompanyWipSlot re-counts under a per-company + // advisory lock and returns null when the company is full, leaving the run + // queued. Disabled limit => no transaction, plain `db`. See company-wip-limit.ts. + const claimed = await withCompanyWipSlot( + { db, companyId: run.companyId, limit: companyMaxConcurrentRuns }, + (tx) => + tx + .update(heartbeatRuns) + .set({ + status: "running", + responsibleUserId, + startedAt: run.startedAt ?? claimedAt, + updatedAt: claimedAt, + }) + .where(and(eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.status, "queued"))) + .returning() + .then((rows) => rows[0] ?? null), + ); if (!claimed) return null; publishLiveEvent({ @@ -13289,6 +13303,9 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run), }); await startNextQueuedRunForAgent(run.agentId); + // Reaping frees the company slot the dead run held, and the agent waiting + // on it is a different one. Same reason as the completion path. + await startNextQueuedRunsForCompanyPeers(run.companyId, run.agentId); runningProcesses.delete(run.id); reaped.push(run.id); } @@ -13538,6 +13555,36 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) }); } + // Company-wide follow-up dispatch, run when a run finishes. + // + // The company WIP slot is a reservation taken in claimQueuedRun, so the agent + // that loses the race keeps its run `queued`. Every completion path re-dispatches + // only the finishing agent (startNextQueuedRunForAgent), and the loser is by + // definition a *different* agent — without this sweep the freed slot would sit + // idle until the next resumeQueuedRuns timer tick, turning an over-spend bug into + // a latency bug. Sweeps the company's other agents that still have queued runs. + // + // No-op unless a company limit is configured, so upstream scheduling is untouched. + async function startNextQueuedRunsForCompanyPeers(companyId: string, finishedAgentId: string) { + if (!isCompanyWipLimitEnabled(companyMaxConcurrentRuns)) return; + if ((await getSchedulingSuppression()).suppressed) return; + const cutoff = await getWorktreeExecutionCutoff(); + + const queuedRuns = await db + .select({ agentId: heartbeatRuns.agentId }) + .from(heartbeatRuns) + .where(and( + eq(heartbeatRuns.companyId, companyId), + eq(heartbeatRuns.status, "queued"), + ne(heartbeatRuns.agentId, finishedAgentId), + cutoff ? gte(heartbeatRuns.createdAt, cutoff) : undefined, + )); + + for (const agentId of [...new Set(queuedRuns.map((row) => row.agentId))]) { + await startNextQueuedRunForAgent(agentId); + } + } + // Await every background heartbeat execution that is currently in flight. A // draining run can, in its finally block, promote and dispatch the next queued // run for the same agent — that follow-up execution is registered in the set @@ -16384,6 +16431,7 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) } activeRunExecutions.delete(run.id); await startNextQueuedRunForAgent(run.agentId); + await startNextQueuedRunsForCompanyPeers(run.companyId, run.agentId); } } @@ -18577,6 +18625,9 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) wasFirstHeartbeat: timerClaimWasFirstHeartbeat(run), }); await startNextQueuedRunForAgent(run.agentId); + // Cancelling frees the company slot too, and the agent that lost the race for + // it is a different one. Same reason as the completion path. + await startNextQueuedRunsForCompanyPeers(run.companyId, run.agentId); return cancelled; }