From 6154e00f2644827981dd790494dbda130fffbbc8 Mon Sep 17 00:00:00 2001 From: Nicky Leach Date: Sat, 29 Aug 2026 08:31:11 -0700 Subject: [PATCH] feat(server): add a task-drain admission hold to the instance API (#12485) ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - The server admits agent work through heartbeat scheduling and execution paths > - Operators need to stop new work before maintenance or a graceful shutdown > - A process restart alone does not provide a reusable admission control primitive > - This pull request adds an instance API that holds new task admission and reports process quiescence > - The benefit is a small, auditable control that lets operators wait for active work without a restart ## Linked Issues or Issue Description **Problem or motivation** Operators cannot hold new task admission without restarting the Paperclip process. A restart can interrupt maintenance flows and does not provide a status signal for active work. **Proposed solution** Add `GET /instance/task-drain`, `POST /instance/task-drain`, and `DELETE /instance/task-drain`. The server keeps the drain state in process memory, applies it to every scheduling suppression path, supports an optional TTL up to 24 hours, and reports active wake and run counts. **Alternatives considered** A timer would clear the drain after its TTL, but it could keep the Node.js event loop open during shutdown. A database row would add storage and query work for process-local state. The implementation uses lazy expiry and process memory instead. **Roadmap alignment** The change supports the roadmap goal for enforced outcomes and safe recovery actions. It does not duplicate a listed roadmap item. **Additional context** This is a server and shared-package change. It adds no user interface and no database migration. ## What Changed - Add process-local task-drain state with lazy TTL expiry. - Add task-drain admission suppression to the shared heartbeat resolver. - Add instance routes to read, start, and stop a task drain. - Add validation for positive TTL values and the shared 24-hour maximum. - Add activity records for drain mutations and tests for status, access control, validation, and suppression. ## Verification - Run `pnpm exec vitest run --project @paperclipai/server server/src/__tests__/heartbeat-task-drain.test.ts server/src/__tests__/instance-settings-routes.test.ts server/src/__tests__/heartbeat-scheduling-suppression.test.ts`. - Run `pnpm --filter @paperclipai/shared exec tsc --noEmit`. - Run `pnpm --filter @paperclipai/server exec tsc --noEmit` and compare its known pre-existing errors with the base commit. - Confirm that pull request CI reaches a terminal green state. ## Risks The drain state exists only in process memory, so a restart clears it. This behavior matches the process-local design. A drain without a TTL remains active until an operator calls the delete route. The status route reads in-memory activity sets and does not query stale database rows. ## Model Used OpenAI Codex, GPT-5, extended reasoning with tool use and code execution. The exact runtime context window is not exposed by the execution environment. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge --------- Co-authored-by: Paperclip --- packages/shared/src/index.ts | 6 + packages/shared/src/validators/instance.ts | 10 + .../heartbeat-scheduling-suppression.test.ts | 28 +- ...tbeat-task-drain-admission-release.test.ts | 722 ++++++++++++++++++ .../__tests__/heartbeat-task-drain.test.ts | 113 +++ .../instance-settings-routes.test.ts | 367 ++++++++- server/src/routes/instance-settings.ts | 161 +++- server/src/routes/openapi.ts | 26 + server/src/services/heartbeat.ts | 311 +++++++- 9 files changed, 1734 insertions(+), 10 deletions(-) create mode 100644 server/src/__tests__/heartbeat-task-drain-admission-release.test.ts create mode 100644 server/src/__tests__/heartbeat-task-drain.test.ts 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,