diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index e10d52e8a7..65a0e3ae9c 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1640,6 +1640,12 @@ export { isClosedIsolatedExecutionWorkspace, } from "./execution-workspace-guards.js"; +export { + MAX_TASK_DRAIN_TTL_MS, + startTaskDrainRequestSchema, + type StartTaskDrainRequest, +} from "./validators/instance.js"; + export { instanceSettingsSchema, instanceGeneralSettingsSchema, diff --git a/packages/shared/src/validators/instance.ts b/packages/shared/src/validators/instance.ts index 0869c7de6c..1dcf70d467 100644 --- a/packages/shared/src/validators/instance.ts +++ b/packages/shared/src/validators/instance.ts @@ -125,6 +125,15 @@ export const issueGraphLivenessAutoRecoveryRequestSchema = z.object({ .optional(), }).strict(); +// The longest time a task drain can run before it expires on its own. A +// caller can send a shorter `ttlMs`, but not a longer one — the request must +// fail instead of the server silently clamping the value. +export const MAX_TASK_DRAIN_TTL_MS = 24 * 60 * 60 * 1000; + +export const startTaskDrainRequestSchema = z.object({ + ttlMs: z.number().int().positive().max(MAX_TASK_DRAIN_TTL_MS).nullable().optional(), +}).strict(); + export type InstanceGeneralSettings = z.infer; // The patch schema removes each default so an absent key stays absent. Declare // the type from the full settings type, so every field keeps its precise type. @@ -140,6 +149,7 @@ export type PatchInstanceSettings = z.infer; export type IssueGraphLivenessAutoRecoveryRequest = z.infer< typeof issueGraphLivenessAutoRecoveryRequestSchema >; +export type StartTaskDrainRequest = z.infer; export const instanceSettingsSchema = z.object({ id: z.string().guid(), diff --git a/server/src/__tests__/heartbeat-scheduling-suppression.test.ts b/server/src/__tests__/heartbeat-scheduling-suppression.test.ts index f583e60fb8..f0259f086b 100644 --- a/server/src/__tests__/heartbeat-scheduling-suppression.test.ts +++ b/server/src/__tests__/heartbeat-scheduling-suppression.test.ts @@ -1,10 +1,16 @@ -import { describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it } from "vitest"; import { resolveHeartbeatSchedulingSuppression, resolveSkillTestRunCompletionForHeartbeatOutcome, + startTaskDrain, + stopTaskDrain, } from "../services/heartbeat.ts"; describe("heartbeat scheduling suppression", () => { + afterEach(() => { + stopTaskDrain(); + }); + it("suppresses heartbeat scheduling for worktree runtimes", () => { expect(resolveHeartbeatSchedulingSuppression({ PAPERCLIP_IN_WORKTREE: "true", @@ -57,6 +63,26 @@ describe("heartbeat scheduling suppression", () => { }); }); + it("suppresses heartbeat scheduling while a task drain is active", () => { + startTaskDrain({}); + expect(resolveHeartbeatSchedulingSuppression({})).toEqual({ + suppressed: true, + reason: "task_drain", + }); + }); + + it("still suppresses database restore even when a task drain is active", () => { + startTaskDrain({}); + expect( + resolveHeartbeatSchedulingSuppression({ + PAPERCLIP_DATABASE_RESTORE_IN_PROGRESS: "1", + }), + ).toEqual({ + suppressed: true, + reason: "database_restore_in_progress", + }); + }); + it("maps unsuccessful heartbeat outcomes to terminal skill test run outcomes", () => { expect(resolveSkillTestRunCompletionForHeartbeatOutcome("succeeded", null)).toBeNull(); expect(resolveSkillTestRunCompletionForHeartbeatOutcome("cancelled", null)).toEqual({ diff --git a/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts b/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts new file mode 100644 index 0000000000..ddbb642f94 --- /dev/null +++ b/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts @@ -0,0 +1,722 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + activityLog, + agents, + agentRuntimeState, + agentWakeupRequests, + companies, + companySkills, + createDb, + heartbeatRunEvents, + heartbeatRuns, + issues, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { heartbeatService, getTaskDrainStatus, startTaskDrain, stopTaskDrain } from "../services/heartbeat.ts"; +import { subscribeCompanyLiveEvents } from "../services/live-events.ts"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +if (!embeddedPostgresSupport.supported) { + console.warn( + `Skipping embedded Postgres task-drain admission release tests on this host: ${embeddedPostgresSupport.reason ?? "unsupported environment"}`, + ); +} + +describeEmbeddedPostgres("heartbeat task-drain admission release", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("heartbeat-task-drain-admission-release-"); + db = createDb(tempDb.connectionString); + }, 60_000); + + function isHeartbeatRunDependentFkError(error: unknown) { + const message = error instanceof Error ? `${error.message} ${String(error.cause ?? "")}` : String(error); + return ( + message.includes("heartbeat_run_events_run_id_heartbeat_runs_id_fk") || + message.includes("activity_log_run_id_heartbeat_runs_id_fk") + ); + } + + async function deleteHeartbeatRunsWithDependents() { + for (let attempt = 0; attempt < 5; attempt += 1) { + await db.delete(heartbeatRunEvents); + await db.delete(activityLog); + try { + await db.delete(heartbeatRuns); + return; + } catch (error) { + if (!isHeartbeatRunDependentFkError(error) || attempt === 4) throw error; + await new Promise((resolve) => setTimeout(resolve, 25)); + } + } + } + + afterEach(async () => { + stopTaskDrain(); + await deleteHeartbeatRunsWithDependents(); + await db.delete(agentWakeupRequests); + await db.delete(issues); + await db.delete(agentRuntimeState); + await db.delete(agents); + await db.delete(companySkills); + await db.delete(companies); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seedQueuedRun() { + const companyId = randomUUID(); + const agentId = randomUUID(); + const issueId = randomUUID(); + const runId = randomUUID(); + const wakeupRequestId = randomUUID(); + + await db.insert(companies).values({ + id: companyId, + name: "Paperclip", + issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + + await db.insert(agents).values({ + id: agentId, + companyId, + name: "Drain Race Agent", + role: "engineer", + status: "idle", + adapterType: "process", + adapterConfig: { + command: process.execPath, + args: ["-e", "process.exit(0)"], + }, + runtimeConfig: { + heartbeat: { + enabled: true, + intervalSec: 60, + wakeOnDemand: true, + }, + }, + permissions: {}, + }); + + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Work claimed just before a drain trips", + status: "todo", + priority: "high", + assigneeAgentId: agentId, + responsibleUserId: "responsible-user", + }); + + await db.insert(agentWakeupRequests).values({ + id: wakeupRequestId, + companyId, + agentId, + source: "assignment", + status: "queued", + }); + + await db.insert(heartbeatRuns).values({ + id: runId, + companyId, + agentId, + invocationSource: "assignment", + triggerDetail: "system", + status: "queued", + wakeupRequestId, + contextSnapshot: { issueId, wakeReason: "issue_assigned" }, + }); + + return { companyId, agentId, issueId, runId, wakeupRequestId }; + } + + it("releases the run, wakeup, and issue lock when a task drain trips right after the run is claimed", async () => { + const { companyId, issueId, runId, wakeupRequestId } = await seedQueuedRun(); + const heartbeat = heartbeatService(db); + + // The claim path publishes a "heartbeat.run.status" live event with + // status "running" the moment it flips the run row, before the run is + // dispatched to executeRun's second suppression check. Starting the + // drain from that same event reproduces the gap the fix closes: the + // drain trips after the first admission check passed but before the + // second one runs. + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + const run = await db + .select({ + status: heartbeatRuns.status, + startedAt: heartbeatRuns.startedAt, + responsibleUserId: heartbeatRuns.responsibleUserId, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(run).toMatchObject({ status: "queued", startedAt: null, responsibleUserId: null }); + + const wakeup = await db + .select({ status: agentWakeupRequests.status, claimedAt: agentWakeupRequests.claimedAt }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(wakeup).toMatchObject({ status: "queued", claimedAt: null }); + + const issue = await db + .select({ + executionRunId: issues.executionRunId, + executionAgentNameKey: issues.executionAgentNameKey, + executionLockedAt: issues.executionLockedAt, + }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(issue).toMatchObject({ + executionRunId: null, + executionAgentNameKey: null, + executionLockedAt: null, + }); + + const status = getTaskDrainStatus(); + expect(status.draining).toBe(true); + expect(status.activeRuns).toBe(0); + expect(status.pendingWakes).toBe(0); + expect(status.quiescent).toBe(true); + + // The released run is not orphaned: once the drain lifts, the normal + // admission path picks it back up and it runs to completion. + stopTaskDrain(); + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + + const finished = await db + .select({ status: heartbeatRuns.status }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(finished?.status).toBe("succeeded"); + }, 20_000); + + // Wraps db.transaction so the callback's tx object throws the moment code + // calls tx.update(table) for a table named in tablesByCall — this makes a + // real Postgres transaction roll back exactly like a genuine write failure + // partway through, without touching any other table's update path. + // tablesByCall maps a 0-based db.transaction() call index (in call order) + // to the table that call should fail on; a call index with no entry runs + // every update for real. For example { 0: issues, 1: agentWakeupRequests } + // fails only the issue-lock write in the first transaction + // (releaseRunClaimedJustBeforeSuppression) and only the wakeup write in + // the second (failRunClaimedJustBeforeSuppression's own transaction). + function withFailingTransactionalUpdate(realDb: typeof db, tablesByCall: Record) { + let callIndex = 0; + return new Proxy(realDb, { + get(target, prop, receiver) { + if (prop !== "transaction") return Reflect.get(target, prop, receiver); + return (fn: (tx: unknown) => Promise) => { + const failingTable = tablesByCall[callIndex]; + callIndex += 1; + return target.transaction((tx) => { + const txProxy = new Proxy(tx as object, { + get(txTarget, txProp, txReceiver) { + if (txProp === "update") { + return (table: unknown) => { + if (failingTable !== undefined && table === failingTable) { + throw new Error("simulated transactional write failure"); + } + return (txTarget as any).update(table); + }; + } + return Reflect.get(txTarget, txProp, txReceiver); + }, + }); + return fn(txProxy); + }); + }; + }, + }) as typeof db; + } + + for (const [label, failingTable] of [ + ["the wakeup-request update", agentWakeupRequests], + ["the issue-lock update", issues], + ] as const) { + it(`fails the run instead of leaving it claimed when ${label} fails`, async () => { + const { companyId, issueId, runId, wakeupRequestId } = await seedQueuedRun(); + // Fault only the first (release) transaction, so the fallback's own + // transaction runs for real and this test proves it can still reach + // "failed" on its own — atomicity of the fallback itself is covered + // separately below. + const failingDb = withFailingTransactionalUpdate(db, { 0: failingTable }); + const heartbeat = heartbeatService(failingDb); + + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + // The atomic release transaction rolled back (a non-atomic release + // would show a partial mix of "queued" and "claimed" instead), so + // executeRun's fallback takes over and fails the run outright. A + // stuck "running" run here would keep the wakeup claimed and the + // issue locked forever while active tracking already reports zero + // active runs — the false-quiescence bug this test guards against. + const run = await db + .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(run?.status).toBe("failed"); + expect(run?.errorCode).toBe("claim_release_failed"); + + const wakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(wakeup?.status).toBe("failed"); + + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(issue?.executionRunId).toBeNull(); + + // The database converged to the same "not active" conclusion active + // tracking already reached, so quiescence now reads true because it + // is genuinely true, not because the database was never checked. + const status = getTaskDrainStatus(); + expect(status.activeRuns).toBe(0); + expect(status.quiescent).toBe(true); + }, 20_000); + } + + // Wraps db.transaction so the release transaction (call 0) fails on + // releaseFailingTable exactly like withFailingTransactionalUpdate above — + // this forces the fallback to run. The fallback's own transaction (call 1) + // first writes a terminal outcome to the run, wakeup, and issue-lock rows + // before it runs its real update. This stands in for a concurrent path (a + // cancellation, the orphan reaper) that reaches a terminal status — and + // finishes releasing the same three rows this fallback also guards — while + // the fallback was still waiting to run its own update. + function withRunTerminalizedBeforeFallbackUpdate( + realDb: typeof db, + releaseFailingTable: unknown, + ids: { runId: string; wakeupRequestId: string; issueId: string }, + ) { + let callIndex = 0; + return new Proxy(realDb, { + get(target, prop, receiver) { + if (prop !== "transaction") return Reflect.get(target, prop, receiver); + return (fn: (tx: unknown) => Promise) => { + const isReleaseCall = callIndex === 0; + const isFallbackCall = callIndex === 1; + callIndex += 1; + return target.transaction(async (tx) => { + if (isReleaseCall) { + const txProxy = new Proxy(tx as object, { + get(txTarget, txProp, txReceiver) { + if (txProp === "update") { + return (table: unknown) => { + if (table === releaseFailingTable) { + throw new Error("simulated transactional write failure"); + } + return (txTarget as any).update(table); + }; + } + return Reflect.get(txTarget, txProp, txReceiver); + }, + }); + return fn(txProxy); + } + if (isFallbackCall) { + const now = new Date(); + const txDb = tx as typeof db; + await txDb + .update(heartbeatRuns) + .set({ + status: "cancelled", + finishedAt: now, + error: "Cancelled while a task drain was pending", + errorCode: "cancelled", + updatedAt: now, + }) + .where(eq(heartbeatRuns.id, ids.runId)); + await txDb + .update(agentWakeupRequests) + .set({ status: "cancelled", finishedAt: now, updatedAt: now }) + .where(eq(agentWakeupRequests.id, ids.wakeupRequestId)); + await txDb + .update(issues) + .set({ executionRunId: null, executionAgentNameKey: null, executionLockedAt: null, updatedAt: now }) + .where(eq(issues.id, ids.issueId)); + } + return fn(tx); + }); + }; + }, + }) as typeof db; + } + + it("leaves a run's outcome untouched when another path already terminalized it before the fallback runs", async () => { + const { companyId, issueId, runId, wakeupRequestId } = await seedQueuedRun(); + + const failingDb = withRunTerminalizedBeforeFallbackUpdate(db, issues, { runId, wakeupRequestId, issueId }); + const heartbeat = heartbeatService(failingDb); + + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + // The other path's outcome survives untouched. Before the fix, the + // fallback's unconditional update matched this already-terminal row and + // overwrote it with "failed" / "claim_release_failed", losing the real + // cause. + const run = await db + .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(run?.status).toBe("cancelled"); + expect(run?.errorCode).toBe("cancelled"); + + const wakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(wakeup?.status).toBe("cancelled"); + + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(issue?.executionRunId).toBeNull(); + + // The run is not active by any measure: not in the live execution-promise + // tracking (it already settled) and not in the stuck claim-release + // marker (the fallback found no row to update, so it never throws). + // Quiescence must be able to reach true. + const status = getTaskDrainStatus(); + expect(status.activeRuns).toBe(0); + expect(status.quiescent).toBe(true); + }, 20_000); + + for (const [label, failingTable] of [ + ["the wakeup-request update", agentWakeupRequests], + ["the issue-lock update", issues], + ] as const) { + it(`leaves the run claimed instead of a partial write when the fallback's own ${label} fails`, async () => { + const { companyId, issueId, runId, wakeupRequestId } = await seedQueuedRun(); + // Fault the release transaction (call 0) on the issue lock so the + // fallback engages, then fault the fallback's own transaction + // (call 1) on a different table. Before the fix, the fallback wrote + // the run row with a plain, unconditional update before it ever + // touched the wakeup or issue rows — that write would have committed + // here regardless of what came after it. With the fallback's writes + // in one transaction, a failure anywhere inside it must roll back + // everything, including the run-status write that ran first. + const failingDb = withFailingTransactionalUpdate(db, { 0: issues, 1: failingTable }); + const heartbeat = heartbeatService(failingDb); + + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + // Both transactions rolled back, so the database still shows the run + // exactly as the admission claim left it — claimed, not a mix of + // "failed" run row with a still-claimed wakeup or a still-locked issue. + const run = await db + .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(run?.status).toBe("running"); + expect(run?.errorCode).toBeNull(); + + const wakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(wakeup?.status).toBe("claimed"); + + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(issue?.executionRunId).toBe(runId); + + // The database still holds the claim, so task-drain must not report + // quiescent for it. Before the fix, executeRun's rejection here was + // caught by the dispatch site's generic handler, which removed this + // run's execution promise from active tracking regardless — reporting + // quiescent while the run, wakeup, and issue lock were all still + // durably claimed. + const status = getTaskDrainStatus(); + expect(status.activeRuns).toBeGreaterThanOrEqual(1); + expect(status.quiescent).toBe(false); + + // The run's row is still "running", so the orphan reaper (which the + // failing-transaction proxy no longer intercepts past call index 1) + // finds it, finalizes the run, wakeup, and issue lock for real, and + // this fix drops the in-memory marker along with them. Before the + // fix, this marker survived the reap and quiescent stayed false + // until the process restarted. + const reapResult = await heartbeat.reapOrphanedRuns(); + expect(reapResult.runIds).toContain(runId); + + const reapedRun = await db + .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(reapedRun?.status).toBe("failed"); + expect(reapedRun?.errorCode).toBe("process_lost"); + + const reapedWakeup = await db + .select({ status: agentWakeupRequests.status }) + .from(agentWakeupRequests) + .where(eq(agentWakeupRequests.id, wakeupRequestId)) + .then((rows) => rows[0] ?? null); + expect(reapedWakeup?.status).toBe("failed"); + + // The issue is still "todo" and assigned to the same agent, so the + // reaper's normal self-heal path queues a fresh recovery run for it + // instead of leaving the lock empty — that recovery is unrelated to + // this fix and stays queued (not running) because the drain is still + // active, so it does not itself count toward activeRuns below. + const reapedIssue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(reapedIssue?.executionRunId).not.toBe(runId); + + const statusAfterReap = getTaskDrainStatus(); + expect(statusAfterReap.activeRuns).toBe(0); + expect(statusAfterReap.quiescent).toBe(true); + }, 20_000); + } + + // Wraps a db handle (which may already be wrapped by + // withFailingTransactionalUpdate) so a call to db.insert(table) throws + // once armed.value is true. Lets a test fail one specific later cleanup + // step without touching any insert that happens earlier. + function withFailingInsertWhenArmed(realDb: typeof db, table: unknown, armed: { value: boolean }) { + return new Proxy(realDb, { + get(target, prop, receiver) { + if (prop !== "insert") return Reflect.get(target, prop, receiver); + return (insertTable: unknown) => { + if (armed.value && insertTable === table) { + throw new Error("simulated insert failure"); + } + return (target as any).insert(insertTable); + }; + }, + }) as typeof db; + } + + // Wraps a db handle (which may already be wrapped by + // withFailingTransactionalUpdate) so a db.transaction() callback's own + // tx.update(table) call throws once armed.value is true. This mirrors + // withFailingInsertWhenArmed above, but for a table a step updates inside + // its own transaction (releaseIssueExecutionAndPromote updates the issues + // table this way) instead of a plain top-level insert. + function withFailingTransactionalUpdateWhenArmed(realDb: typeof db, table: unknown, armed: { value: boolean }) { + return new Proxy(realDb, { + get(target, prop, receiver) { + if (prop !== "transaction") return Reflect.get(target, prop, receiver); + return (fn: (tx: unknown) => Promise) => + target.transaction((tx) => { + const txProxy = new Proxy(tx as object, { + get(txTarget, txProp, txReceiver) { + if (txProp === "update") { + return (updateTable: unknown) => { + if (armed.value && updateTable === table) { + throw new Error("simulated issue-lock release failure"); + } + return (txTarget as any).update(updateTable); + }; + } + return Reflect.get(txTarget, txProp, txReceiver); + }, + }); + return fn(txProxy); + }); + }, + }) as typeof db; + } + + it("clears the stuck claim-release marker even when later reap cleanup rejects", async () => { + const { companyId, runId } = await seedQueuedRun(); + // Reuse the same setup as "leaves the run claimed instead of a partial + // write" above: both the release transaction and the fallback's own + // transaction fail, so the run stays "running" and its claim-release + // marker keeps task drain non-quiescent until the orphan reaper picks + // the run up. + const transactionFailingDb = withFailingTransactionalUpdate(db, { 0: issues, 1: agentWakeupRequests }); + const armedRunEventInsertFailure = { value: false }; + const failingDb = withFailingInsertWhenArmed(transactionFailingDb, heartbeatRunEvents, armedRunEventInsertFailure); + const heartbeat = heartbeatService(failingDb); + + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + const claimedStatus = getTaskDrainStatus(); + expect(claimedStatus.activeRuns).toBeGreaterThanOrEqual(1); + expect(claimedStatus.quiescent).toBe(false); + + // Arm the failure only now, so it hits the reap loop's own run-event + // insert — a cleanup step that runs after the run's row already reaches + // a terminal status — instead of any insert during the claim race above. + armedRunEventInsertFailure.value = true; + await expect(heartbeat.reapOrphanedRuns()).rejects.toThrow("simulated insert failure"); + + // The run's row reached "failed" before the injected failure, and the + // marker must have cleared right after that point, not only after every + // later cleanup step succeeds — the bug this test guards against. + const reapedRun = await db + .select({ status: heartbeatRuns.status }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(reapedRun?.status).toBe("failed"); + + const statusAfterFailedReap = getTaskDrainStatus(); + expect(statusAfterFailedReap.activeRuns).toBe(0); + expect(statusAfterFailedReap.quiescent).toBe(true); + }, 20_000); + + // This test's failure (the issue-lock release itself rejects) never + // resolves the run's marker within this run of the process — see + // stuckClaimReleaseRunIds's own comment in heartbeat.ts: that is the + // documented, accepted outcome when the lock release itself keeps + // failing, not a bug. Because the marker is process-memory state with no + // per-test reset, this test runs last in the file so its permanently + // stuck marker cannot affect another test's activeRuns count. + it("keeps the stuck claim-release marker active when the reap loop's own issue-lock release rejects", async () => { + const { companyId, issueId, runId } = await seedQueuedRun(); + // Same admission-race setup as "leaves the run claimed instead of a + // partial write" above: both the release transaction and the fallback's + // own transaction fail, so the run stays "running" and its + // claim-release marker keeps task drain non-quiescent until the orphan + // reaper picks the run up. + const transactionFailingDb = withFailingTransactionalUpdate(db, { 0: issues, 1: agentWakeupRequests }); + const armedIssueReleaseFailure = { value: false }; + const failingDb = withFailingTransactionalUpdateWhenArmed(transactionFailingDb, issues, armedIssueReleaseFailure); + const heartbeat = heartbeatService(failingDb); + + const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { + const payload = event.payload as { runId?: string; status?: string }; + if (event.type === "heartbeat.run.status" && payload.runId === runId && payload.status === "running") { + startTaskDrain({}); + } + }); + + try { + await heartbeat.resumeQueuedRuns(); + await heartbeat.drainActiveRunExecutions(); + } finally { + unsubscribe(); + } + + const claimedStatus = getTaskDrainStatus(); + expect(claimedStatus.activeRuns).toBeGreaterThanOrEqual(1); + expect(claimedStatus.quiescent).toBe(false); + + // Arm the failure only now, so it hits releaseIssueExecutionAndPromote's + // own issue-lock update inside the reap loop, after the run's row + // already reached "failed" — not the admission-time issue-lock write + // exercised above. + armedIssueReleaseFailure.value = true; + await expect(heartbeat.reapOrphanedRuns()).rejects.toThrow("simulated issue-lock release failure"); + + const reapedRun = await db + .select({ status: heartbeatRuns.status }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId)) + .then((rows) => rows[0] ?? null); + expect(reapedRun?.status).toBe("failed"); + + // The run's row reached a terminal status, but its issue-lock release + // itself failed, so the issue is still locked to this run. The marker + // must stay active and keep reporting this instance non-quiescent — the + // failure this test guards against clears it as soon as the row reaches + // a terminal status, before the lock is actually released. + const issue = await db + .select({ executionRunId: issues.executionRunId }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0] ?? null); + expect(issue?.executionRunId).toBe(runId); + + const statusAfterFailedRelease = getTaskDrainStatus(); + expect(statusAfterFailedRelease.activeRuns).toBeGreaterThanOrEqual(1); + expect(statusAfterFailedRelease.quiescent).toBe(false); + }, 20_000); +}); diff --git a/server/src/__tests__/heartbeat-task-drain.test.ts b/server/src/__tests__/heartbeat-task-drain.test.ts new file mode 100644 index 0000000000..b7040d80a5 --- /dev/null +++ b/server/src/__tests__/heartbeat-task-drain.test.ts @@ -0,0 +1,113 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { MAX_TASK_DRAIN_TTL_MS } from "@paperclipai/shared"; +import { + getTaskDrainGeneration, + getTaskDrainStatus, + resolveHeartbeatSchedulingSuppression, + restoreTaskDrainIfCurrent, + startTaskDrain, + stopTaskDrain, +} from "../services/heartbeat.ts"; + +describe("heartbeat task drain", () => { + afterEach(() => { + stopTaskDrain(); + vi.useRealTimers(); + }); + + it("start_task_drain_suppresses_admission", () => { + startTaskDrain({}); + expect(resolveHeartbeatSchedulingSuppression({})).toEqual({ + suppressed: true, + reason: "task_drain", + }); + }); + + it("stop_task_drain_restores_admission", () => { + startTaskDrain({}); + expect(stopTaskDrain()).toEqual({ wasActive: true }); + expect(resolveHeartbeatSchedulingSuppression({})).toEqual({ + suppressed: false, + reason: null, + }); + expect(stopTaskDrain()).toEqual({ wasActive: false }); + }); + + it("null_ttl_produces_no_expiry", () => { + const { expiresAt } = startTaskDrain({ ttlMs: null }); + expect(expiresAt).toBeNull(); + expect(getTaskDrainStatus().expiresAt).toBeNull(); + }); + + it("ttl_above_the_maximum_clamps_to_24_hours", () => { + const { startedAt, expiresAt } = startTaskDrain({ ttlMs: MAX_TASK_DRAIN_TTL_MS * 10 }); + expect(expiresAt).not.toBeNull(); + expect((expiresAt as Date).getTime() - startedAt.getTime()).toBe(MAX_TASK_DRAIN_TTL_MS); + }); + + it("an_expired_ttl_ends_the_drain_and_restores_admission", () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z")); + startTaskDrain({ ttlMs: 1000 }); + expect(resolveHeartbeatSchedulingSuppression({})).toEqual({ + suppressed: true, + reason: "task_drain", + }); + + vi.setSystemTime(new Date("2026-01-01T00:00:01.001Z")); + expect(resolveHeartbeatSchedulingSuppression({})).toEqual({ + suppressed: false, + reason: null, + }); + expect(getTaskDrainStatus().draining).toBe(false); + }); + + it("status_reports_quiescent_when_both_promise_sets_are_empty", () => { + startTaskDrain({}); + const status = getTaskDrainStatus(); + expect(status.draining).toBe(true); + expect(status.activeRuns).toBe(0); + expect(status.pendingWakes).toBe(0); + expect(status.quiescent).toBe(true); + }); + + it("a_stale_restore_does_not_clobber_a_newer_concurrent_mutation", () => { + // Simulate a route's own mutation, whose caller wants to roll it back + // on a failed audit write. + startTaskDrain({ ttlMs: null }); + const staleGeneration = getTaskDrainGeneration(); + + // A concurrent request supersedes that mutation with its own drain + // before the first caller's rollback runs. + const { startedAt: newerStartedAt } = startTaskDrain({ ttlMs: 5_000 }); + + const restored = restoreTaskDrainIfCurrent(staleGeneration, { draining: false, ttlMs: null }); + + expect(restored).toBe(false); + const status = getTaskDrainStatus(); + expect(status.draining).toBe(true); + expect(status.startedAt).toEqual(newerStartedAt); + }); + + it("restore_applies_when_no_newer_mutation_happened", () => { + startTaskDrain({ ttlMs: 10_000 }); + const generation = getTaskDrainGeneration(); + + const restored = restoreTaskDrainIfCurrent(generation, { draining: false, ttlMs: null }); + + expect(restored).toBe(true); + expect(getTaskDrainStatus().draining).toBe(false); + }); + + it("restore_reinstates_a_prior_active_drain", () => { + startTaskDrain({ ttlMs: null }); + const generation = getTaskDrainGeneration(); + + const restored = restoreTaskDrainIfCurrent(generation, { draining: true, ttlMs: 30_000 }); + + expect(restored).toBe(true); + const status = getTaskDrainStatus(); + expect(status.draining).toBe(true); + expect(status.expiresAt).not.toBeNull(); + }); +}); diff --git a/server/src/__tests__/instance-settings-routes.test.ts b/server/src/__tests__/instance-settings-routes.test.ts index a201627c23..5818d7536d 100644 --- a/server/src/__tests__/instance-settings-routes.test.ts +++ b/server/src/__tests__/instance-settings-routes.test.ts @@ -14,6 +14,11 @@ const mockInstanceSettingsService = vi.hoisted(() => ({ const mockHeartbeatService = vi.hoisted(() => ({ buildIssueGraphLivenessAutoRecoveryPreview: vi.fn(), reconcileIssueGraphLiveness: vi.fn(), + startTaskDrain: vi.fn(), + stopTaskDrain: vi.fn(), + getTaskDrainStatus: vi.fn(), + getTaskDrainGeneration: vi.fn(), + restoreTaskDrainIfCurrent: vi.fn(), })); const mockEnvironmentService = vi.hoisted(() => ({ getById: vi.fn(), @@ -21,12 +26,14 @@ const mockEnvironmentService = vi.hoisted(() => ({ update: vi.fn(), })); const mockLogActivity = vi.hoisted(() => vi.fn()); +const mockPublishActivity = vi.hoisted(() => vi.fn()); function registerModuleMocks() { vi.doMock("../services/index.js", () => ({ heartbeatService: () => mockHeartbeatService, instanceSettingsService: () => mockInstanceSettingsService, logActivity: mockLogActivity, + publishActivity: mockPublishActivity, })); vi.doMock("../services/environments.js", () => ({ environmentService: () => mockEnvironmentService, @@ -36,6 +43,15 @@ function registerModuleMocks() { // Identity object the mocked db.transaction hands to writers; tests assert // both the marker clear and the settings update receive THIS same tx. const TX_SENTINEL = { __tx: true }; +// Module-scoped (not rebuilt per createApp call) so a test can assert how +// many times a request opened a transaction — the task-drain audit writes +// for every company must share ONE transaction, not one each. +const mockDb = { + // Runs the callback with a sentinel tx and propagates throws, so a + // failing write inside rejects the whole request exactly like a real + // transaction rollback. + transaction: vi.fn(async (fn: (tx: unknown) => Promise) => fn(TX_SENTINEL)), +}; async function createApp(actor: any) { const [{ errorHandler }, { instanceSettingsRoutes }] = await Promise.all([ @@ -48,12 +64,6 @@ async function createApp(actor: any) { req.actor = actor; next(); }); - const mockDb = { - // Runs the callback with a sentinel tx and propagates throws, so a - // failing write inside rejects the whole request exactly like a real - // transaction rollback. - transaction: vi.fn(async (fn: (tx: unknown) => Promise) => fn(TX_SENTINEL)), - }; app.use("/api", instanceSettingsRoutes(mockDb as any)); app.use(errorHandler); return app; @@ -77,11 +87,24 @@ describe("instance settings routes", () => { mockInstanceSettingsService.listCompanyIds.mockReset(); mockHeartbeatService.buildIssueGraphLivenessAutoRecoveryPreview.mockReset(); mockHeartbeatService.reconcileIssueGraphLiveness.mockReset(); + mockHeartbeatService.startTaskDrain.mockReset(); + mockHeartbeatService.stopTaskDrain.mockReset(); + mockHeartbeatService.getTaskDrainStatus.mockReset(); + mockHeartbeatService.getTaskDrainGeneration.mockReset(); + mockHeartbeatService.restoreTaskDrainIfCurrent.mockReset(); mockEnvironmentService.getById.mockReset(); mockEnvironmentService.findManagedSandboxEnvironment.mockReset(); mockEnvironmentService.findManagedSandboxEnvironment.mockResolvedValue(null); mockEnvironmentService.update.mockReset(); + mockPublishActivity.mockReset(); mockLogActivity.mockReset(); + // Mirrors the real logActivity: push a publication for the transaction + // to publish once it commits, so route-level tests can prove publish + // happens only after every company's write in the same request lands. + mockLogActivity.mockImplementation((_db: unknown, input: { companyId: string }, postCommitPublications?: unknown[]) => { + postCommitPublications?.push({ companyId: input.companyId, payload: input, pluginEvent: null }); + return Promise.resolve({ id: `activity-${input.companyId}` }); + }); mockInstanceSettingsService.get.mockResolvedValue({ id: "instance-settings-1", defaultEnvironmentId: null, @@ -924,4 +947,336 @@ describe("instance settings routes", () => { expect(experimental.status).toBe(200); }); }); + + describe("task drain", () => { + const adminActor = { + type: "board", + userId: "admin-1", + source: "session", + isInstanceAdmin: true, + companyIds: ["company-1"], + }; + const nonAdminActor = { + type: "board", + userId: "user-1", + source: "session", + isInstanceAdmin: false, + companyIds: ["company-1"], + }; + const idleStatus = { + draining: false, + startedAt: null, + expiresAt: null, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + }; + + afterEach(() => { + // A drain the mock left active must not carry over into an unrelated + // test, so every test starts from the idle status again. + mockHeartbeatService.getTaskDrainStatus.mockReset(); + mockHeartbeatService.startTaskDrain.mockReset(); + mockHeartbeatService.stopTaskDrain.mockReset(); + mockHeartbeatService.getTaskDrainGeneration.mockReset(); + mockHeartbeatService.restoreTaskDrainIfCurrent.mockReset(); + }); + + it("returns the idle status", async () => { + mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); + const app = await createApp(nonAdminActor); + + const res = await request(app).get("/api/instance/task-drain"); + + expect(res.status).toBe(200); + expect(res.body).toEqual(idleStatus); + }); + + it("starts a drain and writes an activity record for every company in one transaction", async () => { + const startedAt = "2026-08-29T00:00:00.000Z"; + const expiresAt = "2026-08-29T06:00:00.000Z"; + mockHeartbeatService.startTaskDrain.mockReturnValue({ startedAt, expiresAt }); + const app = await createApp(adminActor); + + const res = await request(app) + .post("/api/instance/task-drain") + .send({ ttlMs: 21_600_000 }); + + expect(res.status).toBe(200); + expect(mockHeartbeatService.startTaskDrain).toHaveBeenCalledWith({ ttlMs: 21_600_000 }); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + expect(mockLogActivity).toHaveBeenCalledTimes(2); + for (const call of mockLogActivity.mock.calls) { + expect(call[0]).toBe(TX_SENTINEL); + expect(call[1]).toMatchObject({ action: "instance.task_drain.started" }); + } + // Publish only runs after the shared transaction commits. + expect(mockPublishActivity).toHaveBeenCalledTimes(2); + }); + + it("starts an indefinite drain when the caller sends no ttlMs", async () => { + mockHeartbeatService.startTaskDrain.mockReturnValue({ startedAt: "2026-08-29T00:00:00.000Z", expiresAt: null }); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBe(200); + expect(mockHeartbeatService.startTaskDrain).toHaveBeenCalledWith({ ttlMs: null }); + }); + + it("ends the drain and writes an activity record for every company in one transaction", async () => { + mockHeartbeatService.stopTaskDrain.mockReturnValue({ wasActive: true }); + const app = await createApp(adminActor); + + const res = await request(app).delete("/api/instance/task-drain"); + + expect(res.status).toBe(200); + expect(mockHeartbeatService.stopTaskDrain).toHaveBeenCalledWith(); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + expect(mockLogActivity).toHaveBeenCalledTimes(2); + for (const call of mockLogActivity.mock.calls) { + expect(call[0]).toBe(TX_SENTINEL); + expect(call[1]).toMatchObject({ action: "instance.task_drain.stopped" }); + } + expect(mockPublishActivity).toHaveBeenCalledTimes(2); + }); + + it("does not start a drain when the company list read fails", async () => { + mockInstanceSettingsService.listCompanyIds.mockRejectedValue(new Error("db unavailable")); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + }); + + it("commits no activity record for any company when one company's audit write fails", async () => { + // company-1 succeeds, company-2 fails. A real transaction rolls both + // back together; here we prove the route puts both writes in the + // SAME transaction (rather than firing one independent write per + // company) and never publishes a record for the company that did + // succeed before the shared transaction rejected. + mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); + mockHeartbeatService.startTaskDrain.mockReturnValue({ + startedAt: "2026-08-29T00:00:00.000Z", + expiresAt: null, + }); + mockLogActivity.mockImplementation((_db: unknown, input: { companyId: string }) => ( + input.companyId === "company-2" + ? Promise.reject(new Error("activity insert failed")) + : Promise.resolve({ id: `activity-${input.companyId}` }) + )); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + for (const call of mockLogActivity.mock.calls) { + expect(call[0]).toBe(TX_SENTINEL); + } + expect(mockPublishActivity).not.toHaveBeenCalled(); + }); + + it("reverts the drain when the activity log write fails", async () => { + mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); + mockHeartbeatService.startTaskDrain.mockReturnValue({ + startedAt: "2026-08-29T00:00:00.000Z", + expiresAt: null, + }); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(3); + mockLogActivity.mockRejectedValue(new Error("activity insert failed")); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.startTaskDrain).toHaveBeenCalledTimes(1); + // The rollback restores through the generation-guarded primitive + // (not a direct start/stop call), so a concurrent mutation that + // superseded this one after the generation was stamped is never + // clobbered by this restore. + expect(mockHeartbeatService.restoreTaskDrainIfCurrent).toHaveBeenCalledWith(3, { + draining: false, + ttlMs: null, + }); + }); + + it("restores the prior drain when a POST over an active drain fails to write its audit record", async () => { + const priorExpiresAt = new Date(Date.now() + 60_000); + mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ + draining: true, + startedAt: new Date(Date.now() - 60_000), + expiresAt: priorExpiresAt, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + }); + mockHeartbeatService.startTaskDrain.mockReturnValue({ + startedAt: "2026-08-29T00:00:00.000Z", + expiresAt: "2026-08-29T06:00:00.000Z", + }); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(9); + mockLogActivity.mockRejectedValue(new Error("activity insert failed")); + const app = await createApp(adminActor); + + const res = await request(app) + .post("/api/instance/task-drain") + .send({ ttlMs: 21_600_000 }); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.stopTaskDrain).not.toHaveBeenCalled(); + expect(mockHeartbeatService.startTaskDrain).toHaveBeenCalledTimes(1); + expect(mockHeartbeatService.restoreTaskDrainIfCurrent).toHaveBeenCalledTimes(1); + const [generationArg, restoreArg] = mockHeartbeatService.restoreTaskDrainIfCurrent.mock.calls[0]; + expect(generationArg).toBe(9); + expect(restoreArg.draining).toBe(true); + expect(restoreArg.ttlMs).toBeGreaterThan(0); + expect(restoreArg.ttlMs).toBeLessThanOrEqual(60_000); + }); + + it("still reports the started drain when publishing its committed audit record fails", async () => { + // The audit row already committed by the time publish runs, so a + // publish failure must not roll the in-memory drain back — that + // would desync it from the row a client can already read. It also + // must not turn the response into a false failure: the caller asked + // to start a drain, and the drain did start. + mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); + const drain = { startedAt: "2026-08-29T00:00:00.000Z", expiresAt: null }; + mockHeartbeatService.startTaskDrain.mockReturnValue(drain); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(5); + mockPublishActivity.mockImplementation(() => { + throw new Error("live event bus unavailable"); + }); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBe(200); + expect(res.body).toEqual(drain); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + expect(mockLogActivity).toHaveBeenCalledTimes(2); + expect(mockHeartbeatService.restoreTaskDrainIfCurrent).not.toHaveBeenCalled(); + }); + + it("does not stop the drain when the company list read fails", async () => { + mockInstanceSettingsService.listCompanyIds.mockRejectedValue(new Error("db unavailable")); + const app = await createApp(adminActor); + + const res = await request(app).delete("/api/instance/task-drain"); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.stopTaskDrain).not.toHaveBeenCalled(); + }); + + it("restores an active drain when the activity log write fails", async () => { + const expiresAt = new Date(Date.now() + 60_000); + mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ + draining: true, + startedAt: new Date(), + expiresAt, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + }); + mockHeartbeatService.stopTaskDrain.mockReturnValue({ wasActive: true }); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(11); + mockLogActivity.mockRejectedValue(new Error("activity insert failed")); + const app = await createApp(adminActor); + + const res = await request(app).delete("/api/instance/task-drain"); + + expect(res.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.stopTaskDrain).toHaveBeenCalledTimes(1); + expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + expect(mockHeartbeatService.restoreTaskDrainIfCurrent).toHaveBeenCalledTimes(1); + const [generationArg, restoreArg] = mockHeartbeatService.restoreTaskDrainIfCurrent.mock.calls[0]; + expect(generationArg).toBe(11); + expect(restoreArg.draining).toBe(true); + expect(restoreArg.ttlMs).toBeGreaterThan(0); + expect(restoreArg.ttlMs).toBeLessThanOrEqual(60_000); + }); + + it("still reports the stopped drain when publishing its committed audit record fails", async () => { + const expiresAt = new Date(Date.now() + 60_000); + mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ + draining: true, + startedAt: new Date(), + expiresAt, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + }); + mockHeartbeatService.stopTaskDrain.mockReturnValue({ wasActive: true }); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(7); + mockPublishActivity.mockImplementation(() => { + throw new Error("live event bus unavailable"); + }); + const app = await createApp(adminActor); + + const res = await request(app).delete("/api/instance/task-drain"); + + expect(res.status).toBe(200); + expect(res.body).toEqual({ wasActive: true }); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + expect(mockLogActivity).toHaveBeenCalledTimes(2); + expect(mockHeartbeatService.restoreTaskDrainIfCurrent).not.toHaveBeenCalled(); + }); + + it("still publishes the second company's record when the first company's publish fails", async () => { + // Each committed audit record publishes independently, so one + // company's publish failure must not stop the rest from publishing. + const drain = { startedAt: "2026-08-29T00:00:00.000Z", expiresAt: null }; + mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); + mockHeartbeatService.startTaskDrain.mockReturnValue(drain); + mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(5); + mockPublishActivity.mockImplementation((publication: { companyId: string }) => { + if (publication.companyId === "company-1") throw new Error("live event bus unavailable"); + }); + const app = await createApp(adminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBe(200); + expect(res.body).toEqual(drain); + expect(mockPublishActivity).toHaveBeenCalledTimes(2); + expect(mockPublishActivity.mock.calls.map(([publication]: [{ companyId: string }]) => publication.companyId)).toEqual([ + "company-1", + "company-2", + ]); + }); + + it("rejects a board actor without instance admin rights", async () => { + const app = await createApp(nonAdminActor); + + const res = await request(app).post("/api/instance/task-drain").send({}); + + expect(res.status).toBe(403); + expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + }); + + it("rejects a ttl above the maximum", async () => { + const app = await createApp(adminActor); + + const res = await request(app) + .post("/api/instance/task-drain") + .send({ ttlMs: 24 * 60 * 60 * 1000 + 1 }); + + expect(res.status).toBe(400); + expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + }); + + it("rejects a zero or negative ttl", async () => { + const app = await createApp(adminActor); + + const zeroRes = await request(app).post("/api/instance/task-drain").send({ ttlMs: 0 }); + expect(zeroRes.status).toBe(400); + + const negativeRes = await request(app).post("/api/instance/task-drain").send({ ttlMs: -1 }); + expect(negativeRes.status).toBe(400); + + expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + }); + }); }); diff --git a/server/src/routes/instance-settings.ts b/server/src/routes/instance-settings.ts index 181ced8bf8..a362911982 100644 --- a/server/src/routes/instance-settings.ts +++ b/server/src/routes/instance-settings.ts @@ -5,12 +5,20 @@ import { patchInstanceSettingsSchema, patchInstanceExperimentalSettingsSchema, patchInstanceGeneralSettingsSchema, + startTaskDrainRequestSchema, } from "@paperclipai/shared"; import { forbidden } from "../errors.js"; import { isCloudManagedInstance } from "../services/cloud-instance.js"; import { getHiddenSettings } from "../services/settings-visibility.js"; import { validate } from "../middleware/validate.js"; -import { heartbeatService, instanceSettingsService, logActivity } from "../services/index.js"; +import { logger } from "../middleware/logger.js"; +import { + heartbeatService, + instanceSettingsService, + logActivity, + publishActivity, + type ActivityPublication, +} from "../services/index.js"; import { environmentService } from "../services/environments.js"; import { assertEnvironmentSelectionForCompany } from "./environment-selection.js"; import { assertBoardOrgAccess, getActorInfo } from "./authz.js"; @@ -54,6 +62,22 @@ async function assertNoHiddenSettingChanges( } } +/** + * Publish activity events for an already-committed mutation. The audit row + * exists in the database no matter what happens here, so a publish failure + * must not turn into a route error: that would report the mutation as + * failed to the caller when it in fact succeeded. Log and swallow instead. + */ +function publishActivitiesBestEffort(publications: ActivityPublication[], action: string) { + for (const publication of publications) { + try { + publishActivity(publication); + } catch (err) { + logger.error({ err, action, companyId: publication.companyId }, "failed to publish activity event"); + } + } +} + function assertCanManageInstanceSettings(req: Request) { if (req.actor.type !== "board") { throw forbidden("Board access required"); @@ -291,5 +315,140 @@ export function instanceSettingsRoutes(db: Db) { }, ); + router.get("/instance/task-drain", async (req, res) => { + assertBoardOrgAccess(req); + res.json(heartbeat.getTaskDrainStatus()); + }); + + router.post( + "/instance/task-drain", + validate(startTaskDrainRequestSchema), + async (req, res) => { + assertCanManageInstanceSettings(req); + const actor = getActorInfo(req); + // Read the company list, an operation that can fail, before the + // process-local drain mutation below, so a failed read never leaves + // that mutation in place with no audit record of it. + const companyIds = await svc.listCompanyIds(); + // A POST over an already-active drain replaces it. Capture that prior + // state before the mutation, so a failed audit write below can restore + // it instead of clearing task-drain state the operator still relies on. + const priorStatus = heartbeat.getTaskDrainStatus(); + const drain = heartbeat.startTaskDrain({ ttlMs: req.body.ttlMs ?? null }); + // Stamp the generation right after this call's own mutation (no + // await runs between the two, so nothing else can mutate the drain + // in between), so a later restore can tell whether a concurrent + // request has already superseded it. + const generation = heartbeat.getTaskDrainGeneration(); + // One transaction for every company's audit row, so a write that + // succeeds for one company and fails for another never leaves a + // partial activity history behind — either every company gets the + // record, or none does. + const postCommitActivityPublications: ActivityPublication[] = []; + try { + await db.transaction((tx) => + Promise.all( + companyIds.map((companyId) => + logActivity(tx as unknown as Db, { + companyId, + actorType: actor.actorType, + actorId: actor.actorId, + agentId: actor.agentId, + runId: actor.runId, + agentApiKeyId: actor.agentApiKeyId, + action: "instance.task_drain.started", + entityType: "instance_settings", + entityId: "default", + details: { + startedAt: drain.startedAt, + expiresAt: drain.expiresAt, + }, + }, postCommitActivityPublications), + ), + ), + ); + } catch (err) { + // The audit record did not commit, so undo the in-memory drain this + // call started. If a drain was already active, this call replaced + // it — restore that prior drain (best-effort: the remaining TTL + // carries over, but the original start time does not) instead of + // clearing task-drain state the operator still relies on. Guard the + // restore with the generation stamped above: if a concurrent + // request has already mutated the drain again, this restore must + // not overwrite that newer state with the state captured here. + const remainingTtlMs = priorStatus.expiresAt + ? Math.max(0, priorStatus.expiresAt.getTime() - Date.now()) + : null; + heartbeat.restoreTaskDrainIfCurrent(generation, { + draining: priorStatus.draining, + ttlMs: remainingTtlMs, + }); + throw err; + } + // The audit record already committed, so a failure to publish it here + // is not a reason to undo the drain: reverting the in-memory state at + // this point would desync it from the committed row. Publish outside + // the try above so this failure cannot reach the restore path, and + // swallow a publish failure so it cannot turn a committed mutation + // into a false 500 either. + publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.started"); + res.json(drain); + }, + ); + + router.delete("/instance/task-drain", async (req, res) => { + assertCanManageInstanceSettings(req); + const actor = getActorInfo(req); + const companyIds = await svc.listCompanyIds(); + const priorStatus = heartbeat.getTaskDrainStatus(); + const result = heartbeat.stopTaskDrain(); + // See the POST handler above for why the generation is stamped here. + const generation = heartbeat.getTaskDrainGeneration(); + // See the POST handler above for why this is one transaction. + const postCommitActivityPublications: ActivityPublication[] = []; + try { + await db.transaction((tx) => + Promise.all( + companyIds.map((companyId) => + logActivity(tx as unknown as Db, { + companyId, + actorType: actor.actorType, + actorId: actor.actorId, + agentId: actor.agentId, + runId: actor.runId, + agentApiKeyId: actor.agentApiKeyId, + action: "instance.task_drain.stopped", + entityType: "instance_settings", + entityId: "default", + details: { + wasActive: result.wasActive, + }, + }, postCommitActivityPublications), + ), + ), + ); + } catch (err) { + // Restore the drain this call ended (best-effort: the remaining TTL + // carries over, but the original start time does not) so a failed + // audit write does not silently end a drain the operator still relies + // on to hold new run admission. See the POST handler above for why + // the restore is guarded by the generation stamped above. + if (priorStatus.draining) { + const remainingTtlMs = priorStatus.expiresAt + ? Math.max(0, priorStatus.expiresAt.getTime() - Date.now()) + : null; + heartbeat.restoreTaskDrainIfCurrent(generation, { draining: true, ttlMs: remainingTtlMs }); + } + throw err; + } + // See the POST handler above for why publish runs outside the try, and + // why a publish failure here is swallowed instead of failing the route: + // the audit record already committed, so a publish failure here must + // not undo a drain-stop that is already correct in the database, and + // must not report the stop as failed when it succeeded. + publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.stopped"); + res.json(result); + }); + return router; } diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index ba442d6a24..2407224284 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -172,6 +172,7 @@ import { patchInstanceExperimentalSettingsSchema, patchInstanceSettingsSchema, issueGraphLivenessAutoRecoveryRequestSchema, + startTaskDrainRequestSchema, // Resource memberships updateDocumentResourceMembershipSchema, updateResourceMembershipSchema, @@ -4160,6 +4161,31 @@ registry.registerPath({ responses: { 200: r.ok(), 400: r.badRequest, 401: r.unauthorized }, }); +registry.registerPath({ + method: "get", + path: "/api/instance/task-drain", + tags: ["instance"], + summary: "Get the task-drain status", + responses: { 200: r.ok(), 401: r.unauthorized }, +}); + +registry.registerPath({ + method: "post", + path: "/api/instance/task-drain", + tags: ["instance"], + summary: "Start a task drain, so new run admission holds until active runs finish", + request: { body: jsonBody(startTaskDrainRequestSchema) }, + responses: { 200: r.ok(), 400: r.badRequest, 401: r.unauthorized, 403: r.forbidden }, +}); + +registry.registerPath({ + method: "delete", + path: "/api/instance/task-drain", + tags: ["instance"], + summary: "End a task drain and restore run admission", + responses: { 200: r.ok(), 401: r.unauthorized, 403: r.forbidden }, +}); + // ─── Board chat (Conference Room Chat, experimental) ────────────────────────── registry.registerPath({ diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index ca32a904cf..fca8c3bb9b 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -9,6 +9,7 @@ import { AGENT_DEFAULT_MAX_CONCURRENT_RUNS, ISSUE_CONTINUATION_SUMMARY_DOCUMENT_KEY, ISSUE_DISPOSITION_REPAIR_RETRY_REASON, + MAX_TASK_DRAIN_TTL_MS, MODEL_PROFILE_KEYS, PROVIDER_QUOTA_MONITOR_SERVICE_NAME, envBindingSchema, @@ -858,6 +859,123 @@ const activeRunExecutionPromises = new Set>(); // can await a wake that is still before run registration. A caller that tears // down a shared database (a test afterEach) then cannot race a late wake. const activeWakeupPromises = new Set>(); +// executeRun's task-drain suppression branch can fail to release a run's +// claim two ways in a row: the atomic release transaction fails, then its +// fallback transaction (which marks the run "failed" instead) also fails. +// When that happens the run, wakeup, and issue lock are left in an unknown +// durable state. The dispatch site still removes this run's execution +// promise from activeRunExecutionPromises once executeRun settles — +// drainActiveRunExecutions loops on that set's size, so an entry that never +// clears would hang it forever — so a promise-based set cannot carry this +// signal. Track the runId here instead. getTaskDrainStatus() folds this set +// into its active-run count, so it keeps reporting a non-quiescent instance +// for this run. There is no in-process retry for a fallback that already +// failed once, so the run's row stays "running" until reapOrphanedRuns picks +// it up as an orphan and finalizes it — that is also where this marker gets +// removed, right after the reaper finishes the durable issue-lock cleanup +// for the run (releaseIssueExecutionAndPromote, or handing the run's pending +// work to a retry), and before any later, lock-unrelated cleanup runs. An +// entry here only outlives that reap, and so only clears on a process +// restart, if the reaper itself never runs again for this run, or if +// classification, retry scheduling, or the lock release itself keeps +// failing for it. +const stuckClaimReleaseRunIds = new Set(); +// Thrown by executeRun's task-drain suppression branch when both the atomic +// claim release and its fallback fail a durable write for the same run. The +// dispatch site catches this to add the run to stuckClaimReleaseRunIds +// instead of treating it as an ordinary execution failure. +class RunClaimReleaseUnresolvedFailure extends Error { + constructor(runId: string, cause: unknown) { + const causeMessage = cause instanceof Error ? cause.message : String(cause); + super(`Run claim release failed durably for run ${runId}: ${causeMessage}`); + this.name = "RunClaimReleaseUnresolvedFailure"; + } +} +// Task drain: an operator-controlled hold on new run admission, so a caller +// can wait for active work to finish before it stops the process. The state +// lives in process memory only — a process restart clears it — and it sits at +// module scope like activeRunExecutions above, so both the pure +// resolveHeartbeatSchedulingSuppression() check and every heartbeatService() +// instance see the same drain. +let taskDrainState: { startedAt: Date; expiresAt: Date | null } | null = null; +// Bumped on every explicit task-drain mutation (a start or a stop), so a +// caller can tell whether a later mutation has already superseded its own. +// A route rollback reads this right after its own mutation, then passes it +// to restoreTaskDrainIfCurrent() before it restores prior state on a failed +// audit write — so the rollback never overwrites a newer concurrent +// mutation with stale state. +let taskDrainGeneration = 0; + +function readTaskDrain(now: Date): { startedAt: Date; expiresAt: Date | null } | null { + if (taskDrainState && taskDrainState.expiresAt !== null && taskDrainState.expiresAt.getTime() <= now.getTime()) { + taskDrainState = null; + } + return taskDrainState; +} + +export function startTaskDrain(opts: { ttlMs?: number | null } = {}): { startedAt: Date; expiresAt: Date | null } { + const startedAt = new Date(); + const ttlMs = opts.ttlMs ?? null; + const expiresAt = ttlMs === null ? null : new Date(startedAt.getTime() + Math.min(ttlMs, MAX_TASK_DRAIN_TTL_MS)); + taskDrainState = { startedAt, expiresAt }; + taskDrainGeneration += 1; + return taskDrainState; +} + +export function stopTaskDrain(): { wasActive: boolean } { + const wasActive = readTaskDrain(new Date()) !== null; + taskDrainState = null; + taskDrainGeneration += 1; + return { wasActive }; +} + +/** The current task-drain mutation count. See taskDrainGeneration above. */ +export function getTaskDrainGeneration(): number { + return taskDrainGeneration; +} + +/** + * Restore a captured task-drain state, but only if no other mutation has + * happened since expectedGeneration was read. Returns false, and leaves the + * current state untouched, when a newer mutation has already superseded it. + */ +export function restoreTaskDrainIfCurrent( + expectedGeneration: number, + restore: { draining: boolean; ttlMs: number | null }, +): boolean { + if (taskDrainGeneration !== expectedGeneration) return false; + if (restore.draining) { + startTaskDrain({ ttlMs: restore.ttlMs }); + } else { + stopTaskDrain(); + } + return true; +} + +export function getTaskDrainStatus(): { + draining: boolean; + startedAt: Date | null; + expiresAt: Date | null; + activeRuns: number; + pendingWakes: number; + quiescent: boolean; +} { + const state = readTaskDrain(new Date()); + // Fold in stuckClaimReleaseRunIds so a run whose claim release failed + // durably (see the set's own comment) keeps this read non-quiescent, even + // though its execution promise already left activeRunExecutionPromises. + const activeRuns = activeRunExecutionPromises.size + stuckClaimReleaseRunIds.size; + const pendingWakes = activeWakeupPromises.size; + return { + draining: state !== null, + startedAt: state?.startedAt ?? null, + expiresAt: state?.expiresAt ?? null, + activeRuns, + pendingWakes, + quiescent: activeRuns === 0 && pendingWakes === 0, + }; +} + const INLINE_BASE64_IMAGE_DATA_RE = /("type":"image","source":\{"type":"base64","data":")([A-Za-z0-9+/=]{1024,})(")/g; type RuntimeConfigSecretResolver = Pick< ReturnType, @@ -6750,7 +6868,7 @@ function isTruthyRuntimeEnvValue(value: string | undefined) { export function resolveHeartbeatSchedulingSuppression( env: Record = process.env, overrides: { allowWorktreeRunExecution?: boolean } = {}, -): { suppressed: boolean; reason: "worktree_instance" | "database_restore_in_progress" | null } { +): { suppressed: boolean; reason: "worktree_instance" | "database_restore_in_progress" | "task_drain" | null } { if (isTruthyRuntimeEnvValue(env.PAPERCLIP_IN_WORKTREE) && !overrides.allowWorktreeRunExecution) { return { suppressed: true, reason: "worktree_instance" }; } @@ -6760,6 +6878,9 @@ export function resolveHeartbeatSchedulingSuppression( ) { return { suppressed: true, reason: "database_restore_in_progress" }; } + if (readTaskDrain(new Date()) !== null) { + return { suppressed: true, reason: "task_drain" }; + } return { suppressed: false, reason: null }; } @@ -12792,6 +12913,142 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) return claimed; } + // startNextQueuedRunForAgent checks admission suppression once, then claims + // runs (sets status "running"), then dispatches each to executeRun, which + // checks suppression again before it does any work. Suppression (task + // drain, worktree mode, a database restore) can start in the gap between + // those two checks. When executeRun's check catches that, the run is + // already claimed — release it back to "queued" so it does not keep a + // running execution lock that nothing will ever process. This runs inside + // the same promise startNextQueuedRunForAgent already tracks in + // activeRunExecutionPromises, so getTaskDrainStatus() keeps reporting the + // run as active until the release finishes. + // + // The run row, the wakeup request, and the issue execution lock all guard + // the same claim, so one transaction commits all three writes together. A + // partial write (for example the run flips to "queued" but the wakeup or + // issue update then fails) would let the execution promise clear from + // activeRunExecutionPromises while the wakeup stayed "claimed" or the + // issue stayed locked to a queued run — task-drain status would then read + // quiescent while the database still held part of the old claim. + async function releaseRunClaimedJustBeforeSuppression(runId: string) { + const now = new Date(); + await db.transaction(async (tx) => { + const released = await tx + .update(heartbeatRuns) + .set({ status: "queued", startedAt: null, responsibleUserId: null, updatedAt: now }) + .where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.status, "running"))) + .returning() + .then((rows) => rows[0] ?? null); + if (!released) return; + + if (released.wakeupRequestId) { + await tx + .update(agentWakeupRequests) + .set({ status: "queued", claimedAt: null, updatedAt: now }) + .where(eq(agentWakeupRequests.id, released.wakeupRequestId)); + } + + const context = parseObject(released.contextSnapshot); + const issueId = readNonEmptyString(context.issueId); + if (issueId) { + await tx + .update(issues) + .set({ executionRunId: null, executionAgentNameKey: null, executionLockedAt: null, updatedAt: now }) + .where(and( + eq(issues.id, issueId), + eq(issues.companyId, released.companyId), + eq(issues.executionRunId, released.id), + )); + } + }); + } + + // Fallback for when the atomic release above itself fails (a genuine write + // error, not a normal no-op). executeRun's caller removes this run's + // promise from activeRunExecutionPromises as soon as executeRun settles, + // whether it resolves or rejects — so task-drain quiescence is about to + // read "no active runs" regardless of what happens here. If the run, + // wakeup, and issue lock stayed at "running"/"claimed"/locked, that read + // would be false: the database would still hold a claim nothing is + // tracking anymore. Fail the run outright instead, so the database + // reaches the same "not active" conclusion active tracking already + // reached. This does not retry the "queued" release: a run that could not + // even release cleanly is treated as failed, not requeued. + // + // The run row, the wakeup request, and the issue execution lock all guard + // the same claim, so — same as the atomic release above — one transaction + // commits all three writes together. A partial write here would leave the + // same false-quiescence gap this fallback exists to close. Live-event and + // plugin-event publishing run after the transaction commits, so a publish + // failure cannot roll back the durable claim writes. + // + // The update below carries the same status: "running" condition the atomic + // release above uses. While this fallback waits, a concurrent path (a + // cancellation, the orphan reaper) can move the run to a terminal status. + // Without the condition this update would match that row and overwrite its + // real outcome. With it, the update matches no row, so the function returns + // early below and leaves the terminal run, its wakeup request, and its + // issue lock untouched. + async function failRunClaimedJustBeforeSuppression(runId: string, cause: unknown) { + const now = new Date(); + const causeMessage = cause instanceof Error ? cause.message : String(cause); + + const failed = await db.transaction(async (tx) => { + const updated = await tx + .update(heartbeatRuns) + .set({ + status: "failed", + finishedAt: now, + error: `Failed to release the run claim before task-drain suppression: ${causeMessage}`, + errorCode: "claim_release_failed", + updatedAt: now, + }) + .where(and(eq(heartbeatRuns.id, runId), eq(heartbeatRuns.status, "running"))) + .returning() + .then((rows) => rows[0] ?? null); + if (!updated) return null; + + if (updated.wakeupRequestId) { + await tx + .update(agentWakeupRequests) + .set({ + status: "failed", + finishedAt: now, + error: "Run claim release failed before task-drain suppression", + updatedAt: now, + }) + .where(eq(agentWakeupRequests.id, updated.wakeupRequestId)); + } + + const context = parseObject(updated.contextSnapshot); + const issueId = readNonEmptyString(context.issueId); + if (issueId) { + await tx + .update(issues) + .set({ executionRunId: null, executionAgentNameKey: null, executionLockedAt: null, updatedAt: now }) + .where(and( + eq(issues.id, issueId), + eq(issues.companyId, updated.companyId), + eq(issues.executionRunId, updated.id), + )); + } + + return updated; + }); + if (!failed) return; + + if (isHeartbeatRunTerminalStatus(failed.status)) { + clearHeartbeatRunRuntimeStatus(failed.id); + } + publishLiveEvent({ + companyId: failed.companyId, + type: "heartbeat.run.status", + payload: buildHeartbeatRunStatusLiveEventPayload(failed), + }); + publishRunLifecyclePluginEvent(failed); + } + async function cancelQueuedRunForBlockedDependencies( run: typeof heartbeatRuns.$inferSelect, issueId: string, @@ -13806,6 +14063,19 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) if (!retriedRun) { await releaseIssueExecutionAndPromote(finalizedRun); } + // The run's row reached a terminal status above, and the block above + // just finished the durable issue-lock cleanup for it — either + // releaseIssueExecutionAndPromote cleared executionRunId/checkoutRunId, + // or the retry took over the run's pending work. Only now is it safe to + // drop this run's stuck claim-release marker (see stuckClaimReleaseRunIds + // above): a failure in classification, retry scheduling, or the release + // call above throws before this point, so the marker stays active and + // task-drain keeps reporting this instance non-quiescent while the + // issue may still be locked. Clearing it here, rather than after the + // unrelated cleanup below (event logging, agent-status finalization, + // queue promotion), keeps it from getting stuck on a failure in one of + // those instead. + stuckClaimReleaseRunIds.delete(run.id); await appendRunEvent(finalizedRun, await nextRunEventSeq(finalizedRun.id), { eventType: "lifecycle", @@ -14076,6 +14346,14 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) for (const claimedRun of claimedRuns) { const execution = executeRun(claimedRun.id).catch((err) => { + if (err instanceof RunClaimReleaseUnresolvedFailure) { + // The run's claim release failed durably (both the atomic release + // and its fallback failed a write), already logged inside + // executeRun. Track the runId so getTaskDrainStatus() keeps + // reporting this run as active — see stuckClaimReleaseRunIds. + stuckClaimReleaseRunIds.add(claimedRun.id); + return; + } logger.error({ err, runId: claimedRun.id }, "queued heartbeat execution failed"); }); // Register the in-flight execution so drainActiveRunExecutions() can await @@ -14084,6 +14362,11 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) // have landed before a caller (e.g. a test's afterEach) mutates the DB. activeRunExecutionPromises.add(execution); void execution.finally(() => { + // Always remove the settled execution promise itself — drainActiveRunExecutions + // loops on activeRunExecutionPromises.size, so an entry that never + // clears here would hang it forever. A run added to + // stuckClaimReleaseRunIds above stays reported as active through + // that separate set instead. activeRunExecutionPromises.delete(execution); }); } @@ -14129,7 +14412,26 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) } async function executeRun(runId: string) { - if ((await getSchedulingSuppression()).suppressed) return; + if ((await getSchedulingSuppression()).suppressed) { + try { + await releaseRunClaimedJustBeforeSuppression(runId); + } catch (err) { + logger.error( + { err, runId }, + "failed to release run claimed just before task-drain suppression; failing the run instead", + ); + try { + await failRunClaimedJustBeforeSuppression(runId, err); + } catch (fallbackErr) { + logger.error( + { err: fallbackErr, runId }, + "failed to fail the run after its claim release also failed; the run, wakeup, and issue lock are in an unknown state, so task-drain will keep reporting this run as active until the process restarts", + ); + throw new RunClaimReleaseUnresolvedFailure(runId, fallbackErr); + } + } + return; + } let run = await getRun(runId); if (!run) return; @@ -19695,6 +19997,11 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) resolveSchedulingSuppression: getSchedulingSuppression, drainRunningRunsForShutdown, drainActiveRunExecutions, + startTaskDrain, + stopTaskDrain, + getTaskDrainStatus, + getTaskDrainGeneration, + restoreTaskDrainIfCurrent, promoteDueScheduledRetries, retryScheduledRetryNow,