RomanNum3ral
/
paperclip
mirror of https://github.com/paperclipai/paperclip.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(server): remove unreachable task-drain compensation paths (#12511) 

## Thinking Path

> - Paperclip is the open source app people use to manage AI agents for
work.
> - The task-drain service controls when task execution can start and
stop.
> - The service had compensation paths for states that its validators or
recovery process already handle.
> - These paths added rollback state and a stuck-claim marker without
improving normal drain behavior.
> - This pull request removes the unreachable TTL clamp, audit rollback,
generation counter, and double-fault marker.
> - The result keeps input validation, audit ordering, atomic release,
and orphan recovery.

## Linked Issues or Issue Description

**What existing behavior does this improve?**
The task-drain service and its routes manage drain state, audit rows,
and execution locks.

**Subsystem affected**
server/ — REST API and orchestration services.

**Current behavior**
The service clamps a validated TTL value. The routes mutate drain state
before audit writes and then restore state after a failed write. Claim
release also tracks a second durable-write failure with an in-memory
marker.

**Proposed behavior**
The validator remains the single TTL policy. The routes write audit rows
before they mutate drain state. Claim release logs a failed write and
lets the orphan reaper release the issue lock.

**Reason and benefit**
The removed paths cannot handle a valid API request that reaches them.
The rollback can lose the original start time. The marker can keep a
drain non-quiescent until process restart. The simpler flow keeps state
consistent and uses the existing recovery path.

**Breaking changes**
None to the public API. A failed claim release keeps the issue lock
until the next orphan-reaper cycle.

## What Changed

- Remove the service-layer TTL clamp because the shared validator
rejects values above the limit.
- Write task-drain audit rows before drain mutation and remove the
rollback helpers.
- Remove the rollback generation counter and its unused state.
- Remove double-fault stuck-claim tracking and keep the atomic release
path.
- State that the quiescent flag describes work in this process.
- Keep the orphan reaper as the recovery path after a failed claim
release.

## Verification

- Run `pnpm --filter @paperclipai/server test
server/src/__tests__/heartbeat-task-drain-admission-release.test.ts`.
- Run `pnpm --filter @paperclipai/server test
server/src/__tests__/heartbeat-task-drain.test.ts`.
- Run `pnpm --filter @paperclipai/server test
server/src/__tests__/instance-settings-routes.test.ts`.
- Run `pnpm --filter @paperclipai/server test
server/src/__tests__/heartbeat-scheduling-suppression.test.ts`.
- Run `pnpm --filter @paperclipai/server test
server/src/__tests__/execution-lock-orphan-cleanup.test.ts`.
- The five affected test files pass with 70 tests.
- Confirm the full pull request checks pass before merge.

## Risks

The issue lock remains held until the orphan reaper runs after a failed
claim release. This uses the existing recovery path for interrupted
runs. The change does not alter the public API or database schema.

## Model Used

OpenAI Codex, GPT-5, 400K context window, tool use and code execution.

## 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 <noreply@paperclip.ing>
This commit is contained in:
Nicky Leach 2026-08-29 11:43:16 -07:00 committed by GitHub
parent e3eed3a3ae
commit 4310b0c947
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 287 additions and 851 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

473
server/src/__tests__/heartbeat-task-drain-admission-release.test.ts
View File

@ -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<number, unknown>) {
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<unknown>) => {
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<unknown>) =>
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);
});

48
server/src/__tests__/heartbeat-task-drain.test.ts
View File

@ -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();
});
});

243
server/src/__tests__/instance-settings-routes.test.ts
View File

@ -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<unknown>) => {
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();
});
});
});

153
server/src/routes/instance-settings.ts
View File

@ -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<void> = Promise.resolve();
function withTaskDrainTransition<T>(run: () => Promise<T>): Promise<T> {
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;

2
server/src/routes/openapi.ts
View File

@ -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 },
});

219
server/src/services/heartbeat.ts
View File

@ -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<Promise<void>>();
// 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<Promise<unknown>>();
// 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<string>();
// 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,