diff --git a/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts b/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts index ddbb642f94..cefe6cff38 100644 --- a/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts +++ b/server/src/__tests__/heartbeat-task-drain-admission-release.test.ts @@ -225,10 +225,9 @@ describeEmbeddedPostgres("heartbeat task-drain admission release", () => { // 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). + // every update for real. For example { 0: issues } fails only the + // issue-lock write inside releaseRunClaimedJustBeforeSuppression's + // transaction. function withFailingTransactionalUpdate(realDb: typeof db, tablesByCall: Record) { let callIndex = 0; return new Proxy(realDb, { @@ -258,141 +257,12 @@ describeEmbeddedPostgres("heartbeat task-drain admission release", () => { }) 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 () => { + it("leaves the run row running for the orphan reaper when the atomic release fails", async () => { const { companyId, issueId, runId, wakeupRequestId } = await seedQueuedRun(); - - const failingDb = withRunTerminalizedBeforeFallbackUpdate(db, issues, { runId, wakeupRequestId, issueId }); + // Fault the release transaction on the issue-lock write, so executeRun's + // suppression branch catches the failure, logs it, and returns instead + // of throwing. There is no in-process fallback or retry for this path. + const failingDb = withFailingTransactionalUpdate(db, { 0: issues }); const heartbeat = heartbeatService(failingDb); const unsubscribe = subscribeCompanyLiveEvents(companyId, (event) => { @@ -409,305 +279,22 @@ describeEmbeddedPostgres("heartbeat task-drain admission release", () => { 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. + // The release transaction rolled back, so the run, wakeup, and issue + // lock stay exactly as the admission claim left them. const run = await db - .select({ status: heartbeatRuns.status, errorCode: heartbeatRuns.errorCode }) + .select({ status: heartbeatRuns.status }) .from(heartbeatRuns) .where(eq(heartbeatRuns.id, runId)) .then((rows) => rows[0] ?? null); - expect(run?.status).toBe("cancelled"); - expect(run?.errorCode).toBe("cancelled"); + expect(run?.status).toBe("running"); 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"); + 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).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) @@ -715,8 +302,36 @@ describeEmbeddedPostgres("heartbeat task-drain admission release", () => { .then((rows) => rows[0] ?? null); expect(issue?.executionRunId).toBe(runId); - const statusAfterFailedRelease = getTaskDrainStatus(); - expect(statusAfterFailedRelease.activeRuns).toBeGreaterThanOrEqual(1); - expect(statusAfterFailedRelease.quiescent).toBe(false); + // executeRun did not throw, so the dispatch site removed this run's + // execution promise from active tracking like it does for any other + // completed run. Task-drain now reads quiescent even though the + // database still holds the run claimed: this reading counts in-process + // work only. + const status = getTaskDrainStatus(); + expect(status.activeRuns).toBe(0); + expect(status.quiescent).toBe(true); + + // The run's row is still "running", so the orphan reaper finds it, + // finalizes it, and releases the issue lock on its own cycle. + 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"); + + // 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 pointed at the failed run. + 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); }, 20_000); }); diff --git a/server/src/__tests__/heartbeat-task-drain.test.ts b/server/src/__tests__/heartbeat-task-drain.test.ts index b7040d80a5..522a032ab5 100644 --- a/server/src/__tests__/heartbeat-task-drain.test.ts +++ b/server/src/__tests__/heartbeat-task-drain.test.ts @@ -1,10 +1,7 @@ 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"; @@ -39,12 +36,6 @@ describe("heartbeat task drain", () => { 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")); @@ -71,43 +62,4 @@ describe("heartbeat task drain", () => { 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 5818d7536d..56e6908808 100644 --- a/server/src/__tests__/instance-settings-routes.test.ts +++ b/server/src/__tests__/instance-settings-routes.test.ts @@ -14,11 +14,10 @@ const mockInstanceSettingsService = vi.hoisted(() => ({ const mockHeartbeatService = vi.hoisted(() => ({ buildIssueGraphLivenessAutoRecoveryPreview: vi.fn(), reconcileIssueGraphLiveness: vi.fn(), - startTaskDrain: vi.fn(), + computeTaskDrain: vi.fn(), + applyTaskDrain: vi.fn(), stopTaskDrain: vi.fn(), getTaskDrainStatus: vi.fn(), - getTaskDrainGeneration: vi.fn(), - restoreTaskDrainIfCurrent: vi.fn(), })); const mockEnvironmentService = vi.hoisted(() => ({ getById: vi.fn(), @@ -87,11 +86,10 @@ describe("instance settings routes", () => { mockInstanceSettingsService.listCompanyIds.mockReset(); mockHeartbeatService.buildIssueGraphLivenessAutoRecoveryPreview.mockReset(); mockHeartbeatService.reconcileIssueGraphLiveness.mockReset(); - mockHeartbeatService.startTaskDrain.mockReset(); + mockHeartbeatService.computeTaskDrain.mockReset(); + mockHeartbeatService.applyTaskDrain.mockReset(); mockHeartbeatService.stopTaskDrain.mockReset(); mockHeartbeatService.getTaskDrainStatus.mockReset(); - mockHeartbeatService.getTaskDrainGeneration.mockReset(); - mockHeartbeatService.restoreTaskDrainIfCurrent.mockReset(); mockEnvironmentService.getById.mockReset(); mockEnvironmentService.findManagedSandboxEnvironment.mockReset(); mockEnvironmentService.findManagedSandboxEnvironment.mockResolvedValue(null); @@ -976,10 +974,9 @@ describe("instance settings routes", () => { // 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.computeTaskDrain.mockReset(); + mockHeartbeatService.applyTaskDrain.mockReset(); mockHeartbeatService.stopTaskDrain.mockReset(); - mockHeartbeatService.getTaskDrainGeneration.mockReset(); - mockHeartbeatService.restoreTaskDrainIfCurrent.mockReset(); }); it("returns the idle status", async () => { @@ -992,10 +989,9 @@ describe("instance settings routes", () => { 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 }); + it("writes an activity record for every company, then applies the same drain values, in one transaction", async () => { + const drain = { startedAt: "2026-08-29T00:00:00.000Z", expiresAt: "2026-08-29T06:00:00.000Z" }; + mockHeartbeatService.computeTaskDrain.mockReturnValue(drain); const app = await createApp(adminActor); const res = await request(app) @@ -1003,52 +999,71 @@ describe("instance settings routes", () => { .send({ ttlMs: 21_600_000 }); expect(res.status).toBe(200); - expect(mockHeartbeatService.startTaskDrain).toHaveBeenCalledWith({ ttlMs: 21_600_000 }); + expect(res.body).toEqual(drain); + expect(mockHeartbeatService.computeTaskDrain).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" }); + expect(call[1]).toMatchObject({ + action: "instance.task_drain.started", + details: { startedAt: drain.startedAt, expiresAt: drain.expiresAt }, + }); } - // Publish only runs after the shared transaction commits. + // The mutation applies the same values the audit rows already carry, + // and only after the shared transaction commits. + expect(mockHeartbeatService.applyTaskDrain).toHaveBeenCalledTimes(1); + expect(mockHeartbeatService.applyTaskDrain).toHaveBeenCalledWith(drain); 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 }); + mockHeartbeatService.computeTaskDrain.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 }); + expect(mockHeartbeatService.computeTaskDrain).toHaveBeenCalledWith({ ttlMs: null }); }); - it("ends the drain and writes an activity record for every company in one transaction", async () => { - mockHeartbeatService.stopTaskDrain.mockReturnValue({ wasActive: true }); + it("writes an activity record for every company, then stops the drain, using one wasActive value throughout", async () => { + mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ + draining: true, + startedAt: new Date(), + expiresAt: null, + activeRuns: 0, + pendingWakes: 0, + quiescent: 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(res.body).toEqual({ wasActive: true }); 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(call[1]).toMatchObject({ + action: "instance.task_drain.stopped", + details: { wasActive: true }, + }); } + // The mutation runs only after the shared transaction commits. + expect(mockHeartbeatService.stopTaskDrain).toHaveBeenCalledWith(); expect(mockPublishActivity).toHaveBeenCalledTimes(2); }); - it("does not start a drain when the company list read fails", async () => { + it("does not apply 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(); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); }); it("commits no activity record for any company when one company's audit write fails", async () => { @@ -1057,8 +1072,7 @@ describe("instance settings routes", () => { // 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({ + mockHeartbeatService.computeTaskDrain.mockReturnValue({ startedAt: "2026-08-29T00:00:00.000Z", expiresAt: null, }); @@ -1079,73 +1093,26 @@ describe("instance settings routes", () => { expect(mockPublishActivity).not.toHaveBeenCalled(); }); - it("reverts the drain when the activity log write fails", async () => { - mockHeartbeatService.getTaskDrainStatus.mockReturnValue(idleStatus); - mockHeartbeatService.startTaskDrain.mockReturnValue({ + it("does not apply the drain when the audit transaction rejects", async () => { + mockHeartbeatService.computeTaskDrain.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); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); }); 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); + // publish failure must not turn the response into a false failure: + // the caller asked to start a drain, and the drain did start. const drain = { startedAt: "2026-08-29T00:00:00.000Z", expiresAt: null }; - mockHeartbeatService.startTaskDrain.mockReturnValue(drain); - mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(5); + mockHeartbeatService.computeTaskDrain.mockReturnValue(drain); mockPublishActivity.mockImplementation(() => { throw new Error("live event bus unavailable"); }); @@ -1157,7 +1124,7 @@ describe("instance settings routes", () => { expect(res.body).toEqual(drain); expect(mockDb.transaction).toHaveBeenCalledTimes(1); expect(mockLogActivity).toHaveBeenCalledTimes(2); - expect(mockHeartbeatService.restoreTaskDrainIfCurrent).not.toHaveBeenCalled(); + expect(mockHeartbeatService.applyTaskDrain).toHaveBeenCalledWith(drain); }); it("does not stop the drain when the company list read fails", async () => { @@ -1170,46 +1137,33 @@ describe("instance settings routes", () => { 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); + it("does not stop the drain when the audit transaction rejects", async () => { mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ draining: true, startedAt: new Date(), - expiresAt, + expiresAt: null, 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); + expect(mockHeartbeatService.stopTaskDrain).not.toHaveBeenCalled(); }); 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, + expiresAt: null, activeRuns: 0, pendingWakes: 0, quiescent: true, }); - mockHeartbeatService.stopTaskDrain.mockReturnValue({ wasActive: true }); - mockHeartbeatService.getTaskDrainGeneration.mockReturnValue(7); mockPublishActivity.mockImplementation(() => { throw new Error("live event bus unavailable"); }); @@ -1221,16 +1175,14 @@ describe("instance settings routes", () => { expect(res.body).toEqual({ wasActive: true }); expect(mockDb.transaction).toHaveBeenCalledTimes(1); expect(mockLogActivity).toHaveBeenCalledTimes(2); - expect(mockHeartbeatService.restoreTaskDrainIfCurrent).not.toHaveBeenCalled(); + expect(mockHeartbeatService.stopTaskDrain).toHaveBeenCalledWith(); }); 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); + mockHeartbeatService.computeTaskDrain.mockReturnValue(drain); mockPublishActivity.mockImplementation((publication: { companyId: string }) => { if (publication.companyId === "company-1") throw new Error("live event bus unavailable"); }); @@ -1247,13 +1199,96 @@ describe("instance settings routes", () => { ]); }); + it("queues an overlapping DELETE behind a POST whose audit transaction is still pending, so the audit order and the live state always agree", async () => { + // Model the real drain state instead of a canned return value, so + // this test can prove the applied state, the audit rows, and the + // response all agree even when the earlier request's transaction + // resolves after the later request's would have. + const liveState = { draining: false }; + mockHeartbeatService.computeTaskDrain.mockReturnValue({ + startedAt: "2026-08-29T00:00:00.000Z", + expiresAt: null, + }); + mockHeartbeatService.applyTaskDrain.mockImplementation(() => { + liveState.draining = true; + }); + mockHeartbeatService.stopTaskDrain.mockImplementation(() => { + const wasActive = liveState.draining; + liveState.draining = false; + return { wasActive }; + }); + mockHeartbeatService.getTaskDrainStatus.mockImplementation(() => ({ + draining: liveState.draining, + startedAt: liveState.draining ? "2026-08-29T00:00:00.000Z" : null, + expiresAt: null, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + })); + + // The POST's transaction call blocks until the test releases it. The + // DELETE's transaction call, if it is ever reached, commits at once — + // so if the route let the two requests race, the DELETE would commit + // its audit row, and reach stopTaskDrain, first. + const transactionCalls: string[] = []; + let releasePostTransaction: (() => void) | undefined; + let sawFirstCall = false; + mockDb.transaction.mockImplementation((fn: (tx: unknown) => Promise) => { + if (!sawFirstCall) { + sawFirstCall = true; + transactionCalls.push("post-start"); + return new Promise((resolve) => { + releasePostTransaction = () => { + transactionCalls.push("post-commit"); + resolve(fn(TX_SENTINEL)); + }; + }); + } + transactionCalls.push("delete-start-and-commit"); + return fn(TX_SENTINEL); + }); + + const app = await createApp(adminActor); + + // supertest only sends the request once something calls .then() on + // it, so kick both off eagerly instead of waiting for the final + // Promise.all below to do it — otherwise neither request would even + // reach the (still-pending) POST transaction during the wait. + const postPromise = request(app).post("/api/instance/task-drain").send({}); + postPromise.then(() => {}, () => {}); + const deletePromise = request(app).delete("/api/instance/task-drain"); + deletePromise.then(() => {}, () => {}); + // Give both requests time to reach as far as they can go before the + // POST's transaction is released. + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(transactionCalls).toEqual(["post-start"]); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); + expect(mockHeartbeatService.stopTaskDrain).not.toHaveBeenCalled(); + + releasePostTransaction?.(); + const [postRes, deleteRes] = await Promise.all([postPromise, deletePromise]); + + expect(postRes.status).toBe(200); + expect(deleteRes.status).toBe(200); + // The DELETE's transaction only starts once the POST's has committed + // and the POST has already applied its drain — audit order matches + // application order. + expect(transactionCalls).toEqual(["post-start", "post-commit", "delete-start-and-commit"]); + // The DELETE observed the drain the POST applied, so its audit row + // and its response correctly report an active drain, and the live + // state ends idle — not stuck "draining" from a stale POST that + // applied after the DELETE that was meant to be the final word. + expect(deleteRes.body).toEqual({ wasActive: true }); + expect(liveState.draining).toBe(false); + }); + 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(); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); }); it("rejects a ttl above the maximum", async () => { @@ -1264,7 +1299,7 @@ describe("instance settings routes", () => { .send({ ttlMs: 24 * 60 * 60 * 1000 + 1 }); expect(res.status).toBe(400); - expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); }); it("rejects a zero or negative ttl", async () => { @@ -1276,7 +1311,7 @@ describe("instance settings routes", () => { const negativeRes = await request(app).post("/api/instance/task-drain").send({ ttlMs: -1 }); expect(negativeRes.status).toBe(400); - expect(mockHeartbeatService.startTaskDrain).not.toHaveBeenCalled(); + expect(mockHeartbeatService.applyTaskDrain).not.toHaveBeenCalled(); }); }); }); diff --git a/server/src/routes/instance-settings.ts b/server/src/routes/instance-settings.ts index a362911982..e7cb6d30e9 100644 --- a/server/src/routes/instance-settings.ts +++ b/server/src/routes/instance-settings.ts @@ -88,6 +88,31 @@ function assertCanManageInstanceSettings(req: Request) { throw forbidden("Instance admin access required"); } +// A task-drain start or stop reads the live drain state, writes an audit +// transaction, and only then mutates the process-local drain state. The +// audit write is an async gap: two overlapping requests can commit their +// transactions in one order but reach the in-memory mutation in the other +// order, so a stale transition would win, the audit log would not match the +// live state, and the response for the newer request would not match what +// actually ended up live. Run each request's whole read-audit-apply +// sequence through this queue so overlapping requests execute one at a +// time, in the order they enter it: audit order and apply order then always +// agree, and each response reports exactly the state its own request +// produced. +let taskDrainTransitionQueue: Promise = Promise.resolve(); + +function withTaskDrainTransition(run: () => Promise): Promise { + const turn = taskDrainTransitionQueue.then(run); + // Normalize to a settled void promise for the next caller in line, so a + // rejected transition (a failed audit write, for example) cannot wedge + // every later transition behind it. + taskDrainTransitionQueue = turn.then( + () => undefined, + () => undefined, + ); + return turn; +} + export function instanceSettingsRoutes(db: Db) { const router = Router(); const svc = instanceSettingsService(db); @@ -326,26 +351,23 @@ export function instanceSettingsRoutes(db: Db) { 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 { + const ttlMs = req.body.ttlMs ?? null; + // The whole read-audit-apply sequence runs as one queued transition + // (see withTaskDrainTransition above), so an overlapping start or + // stop cannot commit its audit row, or apply its live state, out of + // order against this one. computeTaskDrain runs inside the turn so + // startedAt reflects the moment this request actually took effect, + // not the moment it arrived and was queued behind another transition. + const drain = await withTaskDrainTransition(async () => { + const computed = heartbeat.computeTaskDrain({ ttlMs }); + // 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. The drain mutation below runs only after + // this transaction commits, so a failed write leaves the live + // drain untouched and there is no partial state to roll back. + const postCommitActivityPublications: ActivityPublication[] = []; await db.transaction((tx) => Promise.all( companyIds.map((companyId) => @@ -360,38 +382,22 @@ export function instanceSettingsRoutes(db: Db) { entityType: "instance_settings", entityId: "default", details: { - startedAt: drain.startedAt, - expiresAt: drain.expiresAt, + startedAt: computed.startedAt, + expiresAt: computed.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"); + heartbeat.applyTaskDrain(computed); + // 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. + // Swallow a publish failure so it cannot turn a committed mutation + // into a false 500. + publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.started"); + return computed; + }); res.json(drain); }, ); @@ -400,13 +406,21 @@ export function instanceSettingsRoutes(db: Db) { 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 { + // See the POST handler above for why the whole read-audit-apply + // sequence runs inside withTaskDrainTransition: it queues this stop + // behind any transition already in flight, so it cannot read a status + // an overlapping request is about to make stale, and its audit row and + // its live-state mutation always land in the same order as every other + // queued transition. + const wasActive = await withTaskDrainTransition(async () => { + const priorStatus = heartbeat.getTaskDrainStatus(); + // Read wasActive once, here, and use this same value for the audit + // detail and the response body below. A TTL that expires between two + // separate reads would otherwise make the two values disagree. + const wasActive = priorStatus.draining; + // See the POST handler above for why this is one transaction, and why + // the drain mutation runs only after it commits. + const postCommitActivityPublications: ActivityPublication[] = []; await db.transaction((tx) => Promise.all( companyIds.map((companyId) => @@ -421,33 +435,22 @@ export function instanceSettingsRoutes(db: Db) { entityType: "instance_settings", entityId: "default", details: { - wasActive: result.wasActive, + 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); + heartbeat.stopTaskDrain(); + // See the POST handler above for 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"); + return wasActive; + }); + res.json({ wasActive }); }); return router; diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index 879b41afd8..cacb3fa560 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -4206,7 +4206,7 @@ registry.registerPath({ method: "get", path: "/api/instance/task-drain", tags: ["instance"], - summary: "Get the task-drain status", + summary: "Get the task-drain status for this process only; quiescent counts in-process work, and a process restart clears it even when the database still holds running rows", responses: { 200: r.ok(), 401: r.unauthorized }, }); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index dfcc80ffe3..5a8ace8b63 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -11,7 +11,6 @@ import { CONNECTION_RUNTIME_TOOL_NAMES, ISSUE_CONTINUATION_SUMMARY_DOCUMENT_KEY, ISSUE_DISPOSITION_REPAIR_RETRY_REASON, - MAX_TASK_DRAIN_TTL_MS, MODEL_PROFILE_KEYS, PROVIDER_QUOTA_MONITOR_SERVICE_NAME, envBindingSchema, @@ -866,38 +865,6 @@ 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 @@ -905,13 +872,6 @@ class RunClaimReleaseUnresolvedFailure extends Error { // 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()) { @@ -920,45 +880,37 @@ function readTaskDrain(now: Date): { startedAt: Date; expiresAt: Date | null } | return taskDrainState; } -export function startTaskDrain(opts: { ttlMs?: number | null } = {}): { startedAt: Date; expiresAt: Date | null } { +/** Compute the drain a start call would apply, without changing state. */ +export function computeTaskDrain(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; + const expiresAt = ttlMs === null ? null : new Date(startedAt.getTime() + ttlMs); + return { startedAt, expiresAt }; +} + +/** Assign the given drain as the current task-drain state. */ +export function applyTaskDrain(drain: { startedAt: Date; expiresAt: Date | null }): void { + taskDrainState = drain; +} + +export function startTaskDrain(opts: { ttlMs?: number | null } = {}): { startedAt: Date; expiresAt: Date | null } { + const drain = computeTaskDrain(opts); + applyTaskDrain(drain); + return drain; } 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. + * Report the task-drain state for this process only. `activeRuns` and + * `pendingWakes` count in-process work. A process restart clears both + * counters, even when the database still holds `running` rows for runs + * this process did not finish. */ -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; @@ -968,10 +920,7 @@ export function getTaskDrainStatus(): { 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 activeRuns = activeRunExecutionPromises.size; const pendingWakes = activeWakeupPromises.size; return { draining: state !== null, @@ -13114,91 +13063,6 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) }); } - // 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, @@ -14238,19 +14102,6 @@ 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", @@ -14521,14 +14372,6 @@ 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 @@ -14537,11 +14380,8 @@ 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. + // drainActiveRunExecutions loops on activeRunExecutionPromises.size, + // so an entry that never clears here would hang it forever. activeRunExecutionPromises.delete(execution); }); } @@ -14593,17 +14433,8 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) } catch (err) { logger.error( { err, runId }, - "failed to release run claimed just before task-drain suppression; failing the run instead", + "failed to release run claimed just before task-drain suppression; the run row stays running, and the orphan reaper finalizes it and releases the issue lock on its next cycle", ); - 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; } @@ -20343,8 +20174,8 @@ export function heartbeatService(db: Db, options: HeartbeatServiceOptions = {}) startTaskDrain, stopTaskDrain, getTaskDrainStatus, - getTaskDrainGeneration, - restoreTaskDrainIfCurrent, + computeTaskDrain, + applyTaskDrain, promoteDueScheduledRetries, retryScheduledRetryNow,