refactor(heartbeat): extract restart recovery and lease cleanup (#15622)

## Thinking Path

> - Paperclip is the open source app people use to manage AI agents for
work.
> - The heartbeat service coordinates agent runs and their recovery.
> - The earlier refactors separated workspaces, run preparation, run
state, and retries.
> - Restart recovery and lease cleanup still occupy more than 2,000
lines in the main service.
> - These operations share process ownership and database claims that
must stay intact.
> - This pull request moves that group into one recovery module and adds
boundary tests.
> - Reviewers can now inspect recovery separately from queue admission
and run execution.

## Linked Issues or Issue Description

**What existing behavior does this improve?**

The structure and testability of heartbeat restart recovery and lease
cleanup.

**Subsystem affected**

server/ — orchestration services.

**Current behavior**

Hot-restart adoption, native restart recovery, shutdown draining, orphan
reaping, and lease cleanup live inside `heartbeat.ts`.

**Proposed behavior**

Keep the same recovery behavior in
`server/src/services/heartbeat/recovery.ts`. Bind the database and
lifecycle callbacks with `createHeartbeatRecovery`.

**Reason and benefit**

This removes 2,181 lines from the main service. The extracted module is
about 2,400 lines. It keeps related recovery operations together in one
file.

**Breaking changes**

None. The public service methods, database claims, cleanup fences, and
ownership checks stay the same.

**Additional context**

Continues the merged heartbeat extractions in #15568, #15573, #15578,
#15591, #15601, and #15617. A search of open and closed PRs found no
duplicate recovery extraction. This does not add a roadmap feature.

## What Changed

- Move hot-restart snapshots and adoption, native restart recovery,
shutdown draining, orphan reaping, and lease sweeps into
`heartbeat/recovery.ts`.
- Pass lifecycle callbacks and the existing shared execution sets into
the factory. Preserve module-wide cleanup single-flight state.
- Set the service's existing shutdown flag through a callback at the
same point in shutdown preparation.
- Add 14 focused tests for construction, shutdown ordering, empty
selective drains, ownership races, shared state, retry timers, and
background execution tracking.
- Wait for the existing execution-drain barrier in the Stop recovery
test instead of polling for one second.
- Document the module boundary in `doc/DEVELOPING.md`.

## Verification

- New module and additional cleanup coverage: **57 tests passed** across
five files.
  ```sh
pnpm exec vitest run server/src/services/heartbeat/recovery.test.ts
server/src/__tests__/heartbeat-native-cleanup-admission.test.ts
server/src/__tests__/heartbeat-run-terminalize-before-release.test.ts
server/src/__tests__/heartbeat-run-lease-release-terminalization.test.ts
server/src/services/hot-restart.test.ts
  ```
- Final recovery coverage on the rebased branch: **514 tests passed
across 11 files**, including all 14 new module tests.
  ```sh
pnpm exec vitest run server/src/services/heartbeat/recovery.test.ts
server/src/__tests__/heartbeat-process-recovery.test.ts
server/src/__tests__/heartbeat-pending-cleanup-sweep.test.ts
server/src/__tests__/heartbeat-orphaned-active-lease-sweep.test.ts
server/src/__tests__/native-session-resumption.test.ts
server/src/__tests__/heartbeat-task-drain-admission-release.test.ts
server/src/shutdown.test.ts
server/src/__tests__/heartbeat-native-cleanup-admission.test.ts
server/src/__tests__/heartbeat-run-terminalize-before-release.test.ts
server/src/__tests__/heartbeat-run-lease-release-terminalization.test.ts
server/src/services/hot-restart.test.ts
  ```
- Full `pnpm -r typecheck` and `pnpm build` passed, including after
rebasing onto master.
- Full local `pnpm test:run` was started again after the rebase. It was
stopped after the full CI suite passed. It did not finish and is not
counted as a full local pass. An earlier local AI-connection run hit an
HTTP-test timeout; all corresponding CI coverage passed.
- Structural comparison confirms 28 moved functions and 15 moved
declarations. The only function-body adaptation is the shutdown flag
callback. Remaining service functions and all 140 public exports match
master.
- All **54 current-head checks passed**; two unrelated Storybook checks
were skipped. This includes the full general and serialized test matrix,
typechecks, build checks, Runner verification, and canary dry run. [CI
run](https://github.com/paperclipai/paperclip/actions/runs/37843023706).
- Greptile completed on head `01fe39818d483e05f940b9fe842162a08e9b9e8e`
with **5/5 and no actionable findings**. There are no review threads or
merge conflicts.

## Risks

The main risk is breaking shared execution ownership or changing
recovery order while moving closures. The factory receives the original
execution sets and lifecycle callbacks. It does no work during
construction. Database transactions, compare-and-set predicates, cleanup
deadlines, and native ownership gates stay intact. This PR has no schema
or API changes.

## Model Used

OpenAI GPT-6 via Codex. The exact serving model ID and context window
were not exposed in this session. Used reasoning, repository inspection,
shell tools, 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:
DottaandPaperclip authored and GitHub committed 2026-10-08 16:15:01 -05:00
1 parent 428b617fc9
commit 04de3132c2
5 files changed
+2767 -2247

No files matched your search

+9
View File
@@ -1263,6 +1263,15 @@ run events, issue-lock release, plan-resume reporting, and worktree cutoffs. It
re-exports the existing retry helpers and workspace-busy error class. Keep retry
policy changes separate from this extraction.
Restart recovery and lease cleanup are in `server/src/services/heartbeat/recovery.ts`.
It owns hot-restart snapshots and adoption, native restart recovery, shutdown
draining, orphaned-run reaping, and active/pending-cleanup lease sweeps.
`createHeartbeatRecovery` binds the service database and explicit lifecycle
callbacks without starting work. The service supplies its shutdown flag callback
and shared execution sets so separate service instances keep the same ownership
and shutdown barriers. Cleanup single-flight state stays at module scope.
Keep recovery policy changes separate from retry scheduling and execution changes.
## Wake Context Delivery
Built-in adapters deliver wake context through the run prompt, including structured
@@ -7974,7 +7974,8 @@ describeEmbeddedPostgres("heartbeat orphaned process recovery", () => {
});
expect(next?.contextSnapshot?.wakeCommentIds).toEqual([pending!.id, go!.id]);
expect((await db.select().from(agentWakeupRequests).where(eq(agentWakeupRequests.id, deferred!.id)))[0]).toMatchObject({ status: "coalesced", runId: next!.id });
await vi.waitFor(async () => expect((await heartbeat.getRun(next!.id))?.status).not.toBe("running"));
await heartbeat.drainActiveRunExecutions();
expect((await heartbeat.getRun(next!.id))?.status).not.toBe("running");
});
it.each(["dedicated deferred donor", "non-coalescing recipient", "persistent agent conversation"] as const)(
File diff suppressed because it is too large. Load diff
@@ -0,0 +1,268 @@
import { randomUUID } from "node:crypto";
import type { Db } from "@paperclipai/db";
import { agents, companies, createDb, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS, heartbeatRuns, issues, nativeRunFinalizations } from "@paperclipai/db";
import { eq } from "drizzle-orm";
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../../__tests__/helpers/embedded-postgres.js";
import { environmentRuntimeService } from "../environment-runtime.js";
import { environmentService } from "../environments.js";
import { createHeartbeatRecovery, type HeartbeatRecoveryDependencies } from "./recovery.js";
import type { HotRestartIntent } from "../hot-restart.js";
import type { NativeRestartRecoveryClaim, NativeRestartRecoveryDisposition } from "../native-runtime/index.js";
import { NATIVE_OWNERSHIP_UNVERIFIED_ERROR_CODE } from "../native-runtime/native-runner-ownership.js";
type Run = typeof heartbeatRuns.$inferSelect;
const native = vi.hoisted(() => ({
closeIdle: vi.fn(async () => ({ closed: 0, failed: 0 })),
detach: vi.fn(async () => 0),
finalizations: vi.fn(async (_db: Db, _ids?: string[], _options?: unknown) => undefined),
cleanups: vi.fn(async () => undefined),
drainMaintenance: vi.fn<() => Promise<void>>(async () => undefined),
claims: vi.fn<() => Promise<NativeRestartRecoveryDisposition[]>>(async () => []),
}));
const hotRestart = vi.hoisted(() => ({
read: vi.fn<() => Promise<HotRestartIntent | null>>(async () => null),
}));
vi.mock("../hot-restart.js", async (importOriginal) => ({
...await importOriginal<typeof import("../hot-restart.js")>(),
readHotRestartIntent: hotRestart.read,
}));
vi.mock("../native-runtime/index.js", async (importOriginal) => ({
...await importOriginal<typeof import("../native-runtime/index.js")>(),
closeIdleWarmNativeSessionsForRestart: native.closeIdle,
detachNativeSessionsForRestart: native.detach,
reconcileNativeFinalizations: native.finalizations,
reconcileRetainedNativeSessionCleanups: native.cleanups,
claimNativeRestartRecoveries: native.claims,
}));
vi.mock("../../vendor/paperclip-runner/index.js", async (importOriginal) => ({
...await importOriginal<typeof import("../../vendor/paperclip-runner/index.js")>(),
drainRetainedRunnerdMaintenanceOperations: native.drainMaintenance,
}));
function callbacks(db: Db) {
return {
enterShutdown: vi.fn(),
getRun: vi.fn<HeartbeatRecoveryDependencies["getRun"]>(async () => null),
getAgent: vi.fn<HeartbeatRecoveryDependencies["getAgent"]>(async () => null),
appendRunEvent: vi.fn<HeartbeatRecoveryDependencies["appendRunEvent"]>(async () => undefined),
setRunStatusIfRunning: vi.fn<HeartbeatRecoveryDependencies["setRunStatusIfRunning"]>(async () => ({ updated: false, run: null })),
setRunStatusFromLive: vi.fn<HeartbeatRecoveryDependencies["setRunStatusFromLive"]>(async () => ({ updated: false, run: null })),
setRunStatus: vi.fn<HeartbeatRecoveryDependencies["setRunStatus"]>(async () => null),
setWakeupStatus: vi.fn(async () => undefined),
releaseIssueExecutionAndPromote: vi.fn(async () => undefined),
finalizeAgentStatus: vi.fn(async () => undefined),
environmentRuntime: environmentRuntimeService(db),
environmentsSvc: environmentService(db),
instructionCopies: {
recoverStopped: vi.fn(async () => 0),
recoverCaptured: vi.fn(async () => 0),
},
runtimeEnv: {},
activeRunExecutions: new Set<string>(),
activeRunExecutionPromises: new Set<Promise<void>>(),
scheduleBoundedRetryForRun: vi.fn<HeartbeatRecoveryDependencies["scheduleBoundedRetryForRun"]>(async () => ({ outcome: "retry_exhausted", retryReason: "transient_failure", attempt: 1, maxAttempts: 0 })),
scheduleInteractionContinuationInfrastructureRetryIfEligible: vi.fn<HeartbeatRecoveryDependencies["scheduleInteractionContinuationInfrastructureRetryIfEligible"]>(async () => null),
timerClaimWasFirstHeartbeat: vi.fn(() => undefined),
executeRun: vi.fn<HeartbeatRecoveryDependencies["executeRun"]>(async () => undefined),
scheduleNativeSessionResumeDispatch: vi.fn(),
cancelHeartbeatNativeRun: vi.fn(async () => undefined),
terminateHeartbeatRunProcess: vi.fn(async () => undefined),
mergeRunStopMetadataForAgent: vi.fn<HeartbeatRecoveryDependencies["mergeRunStopMetadataForAgent"]>((_agent, _outcome, options) => options?.resultJson ?? null),
classifyAndPersistRunLiveness: vi.fn<HeartbeatRecoveryDependencies["classifyAndPersistRunLiveness"]>(async run => run),
releaseEnvironmentLeasesForRun: vi.fn(async () => undefined),
acknowledgeRemoteStop: vi.fn(async () => undefined),
resumeRemoteStopComments: vi.fn(async () => undefined),
dispatchPendingNativeStatusWakeups: vi.fn(async () => undefined),
cancelRunInternal: vi.fn(async () => undefined),
startNextQueuedRunForAgent: vi.fn(async () => undefined),
} satisfies HeartbeatRecoveryDependencies;
}
function guardedDatabase() {
const access = vi.fn(() => { throw new Error("Unexpected recovery database access"); });
return { db: new Proxy({}, { get: access }) as Db, access };
}
beforeEach(() => {
vi.clearAllMocks();
hotRestart.read.mockResolvedValue(null);
native.closeIdle.mockResolvedValue({ closed: 0, failed: 0 });
native.finalizations.mockResolvedValue(undefined);
native.cleanups.mockResolvedValue(undefined);
native.drainMaintenance.mockResolvedValue(undefined);
native.claims.mockResolvedValue([]);
});
describe("heartbeat recovery module callbacks", () => {
it("constructs without database access, shutdown, or background work", () => {
const database = guardedDatabase();
const deps = callbacks(database.db);
createHeartbeatRecovery(database.db, deps);
expect(database.access).not.toHaveBeenCalled();
for (const callback of Object.values(deps)) if (typeof callback === "function") expect(callback).not.toHaveBeenCalled();
for (const callback of Object.values(native)) expect(callback).not.toHaveBeenCalled();
});
it.each(["not_requested", "read_error"] as const)("enters shutdown before checkpointing for %s", async (mode) => {
const database = guardedDatabase();
const deps = callbacks(database.db);
if (mode === "read_error") hotRestart.read.mockRejectedValueOnce(new Error("unreadable intent"));
const recovery = createHeartbeatRecovery(database.db, deps);
expect(await recovery.prepareHotRestartShutdown("SIGTERM")).toEqual({ mode, skipDrain: false, activeRunIds: [] });
expect(native.closeIdle).toHaveBeenCalledAfter(deps.enterShutdown);
expect(hotRestart.read).toHaveBeenCalledAfter(native.closeIdle);
expect(database.access).not.toHaveBeenCalled();
});
it("does not adopt or query runs without a restart intent", async () => {
const database = guardedDatabase();
expect(await createHeartbeatRecovery(database.db, callbacks(database.db)).reconcileHotRestartAdoption()).toEqual({
mode: "not_requested", adoptedRunIds: [], finalizedWhileDownRunIds: [], lostRunIds: [], skippedRunIds: [],
});
expect(database.access).not.toHaveBeenCalled();
});
it("does not broaden an empty selective shutdown drain to every run", async () => {
const database = guardedDatabase();
const deps = callbacks(database.db);
expect(await createHeartbeatRecovery(database.db, deps).drainRunningRunsForShutdown("SIGINT", new Date(), [])).toEqual({
interrupted: 0, interruptedRunIds: [], retryRunIds: [], restartSuspendedRunIds: [],
});
expect(database.access).not.toHaveBeenCalled();
expect(deps.cancelHeartbeatNativeRun).not.toHaveBeenCalled();
expect(deps.terminateHeartbeatRunProcess).not.toHaveBeenCalled();
});
it.each(["running", "failed"] as const)("keeps a %s ownership hold when a status race wins", async (status) => {
const database = guardedDatabase();
const deps = callbacks(database.db);
const run = { id: "run", runtimeMode: "native", status } as Run;
const winner = { ...run, status: "cancelled" };
deps.setRunStatusFromLive.mockResolvedValue({ updated: false, run: winner });
expect(await createHeartbeatRecovery(database.db, deps).markNativeOwnershipUnverified(run, { reason: "live_process_identifier" })).toBe(winner);
expect(deps.setRunStatusFromLive).toHaveBeenCalledWith("run", status, [status], expect.objectContaining({ errorCode: NATIVE_OWNERSHIP_UNVERIFIED_ERROR_CODE }));
expect(deps.appendRunEvent).not.toHaveBeenCalled();
expect(database.access).not.toHaveBeenCalled();
});
it("persists an authentication ownership hold before its event without terminating or retrying", async () => {
const database = guardedDatabase();
const deps = callbacks(database.db);
const run = { id: "run", runtimeMode: "native", status: "running" } as Run;
const blocked = { ...run, nativePhase: "terminal_failure", errorCode: NATIVE_OWNERSHIP_UNVERIFIED_ERROR_CODE };
deps.setRunStatusFromLive.mockResolvedValue({ updated: true, run: blocked });
expect(await createHeartbeatRecovery(database.db, deps).markNativeOwnershipUnverified(run, { reason: "adopted_runner_authentication_timeout" })).toBe(blocked);
expect(deps.setRunStatusFromLive).toHaveBeenCalledWith("run", "running", ["running"], expect.objectContaining({ nativePhase: "terminal_failure" }));
expect(deps.appendRunEvent).toHaveBeenCalledAfter(deps.setRunStatusFromLive);
expect(deps.cancelHeartbeatNativeRun).not.toHaveBeenCalled();
expect(deps.terminateHeartbeatRunProcess).not.toHaveBeenCalled();
expect(deps.scheduleBoundedRetryForRun).not.toHaveBeenCalled();
expect(database.access).not.toHaveBeenCalled();
});
});
const support = await getEmbeddedPostgresTestSupport();
const describePostgres = support.supported ? describe : describe.skip;
if (!support.supported) console.warn(`Skipping recovery module database tests: ${support.reason}`);
describePostgres("heartbeat recovery module ownership", () => {
let db: Db;
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>>;
beforeAll(async () => {
database = await startEmbeddedPostgresTestDatabase("paperclip-heartbeat-recovery-module-");
db = createDb(database.connectionString);
}, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS);
afterAll(async () => { await database?.cleanup(); });
async function seedRun() {
const companyId = randomUUID(), agentId = randomUUID(), runId = randomUUID(), issueId = randomUUID();
await db.insert(companies).values({ id: companyId, name: "Recovery", issuePrefix: `T${companyId.slice(0, 6)}` });
await db.insert(agents).values({ id: agentId, companyId, name: "Recovery", role: "engineer", adapterType: "process" });
await db.insert(issues).values({ id: issueId, companyId, title: "Recover this run", status: "in_progress", assigneeAgentId: agentId });
return (await db.insert(heartbeatRuns).values({ id: runId, companyId, agentId, issueId, invocationSource: "automation", status: "running", runtimeMode: "legacy" }).returning())[0];
}
it("uses the shared live-execution set after construction across two factories", async () => {
const run = await seedRun();
const deps = callbacks(db);
const first = createHeartbeatRecovery(db, deps);
const second = createHeartbeatRecovery(db, deps);
deps.activeRunExecutions.add(run.id);
expect(await first.reapOrphanedRuns()).toEqual({ reaped: 0, runIds: [] });
expect(await second.reapOrphanedRuns()).toEqual({ reaped: 0, runIds: [] });
expect(deps.setRunStatusFromLive).not.toHaveBeenCalled();
expect((await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)))[0].status).toBe("running");
await Promise.all([...deps.activeRunExecutionPromises]);
await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
});
it("settles persisted native results before claiming restart authority and tracks the exact execution", async () => {
const deps = callbacks(db);
const claim: NativeRestartRecoveryClaim = { kind: "bootstrap_incomplete", runId: randomUUID(), leaseOwner: "claimed-owner", controllerGeneration: 3, providerAttempt: 2, restartKind: "hard", recoveryRequestId: null };
native.claims.mockResolvedValue([claim]);
let finish!: () => void;
const execution = new Promise<void>(resolve => { finish = resolve; });
deps.executeRun.mockReturnValue(execution);
const result = await createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart();
expect(result.claims).toEqual([claim]);
expect(native.claims).toHaveBeenCalledAfter(native.finalizations);
expect(deps.executeRun).toHaveBeenCalledWith(claim.runId, { nativeLeaseOwner: claim.leaseOwner, nativeRestartRecovery: claim });
expect(deps.activeRunExecutionPromises.size).toBeGreaterThan(0);
finish();
await Promise.all([...deps.activeRunExecutionPromises]);
await vi.waitFor(() => expect(deps.activeRunExecutionPromises.size).toBe(0));
});
it("does not claim or execute recovery when persisted finalization fails", async () => {
const deps = callbacks(db);
native.finalizations.mockRejectedValueOnce(new Error("result settlement failed"));
await expect(createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart()).rejects.toThrow("result settlement failed");
expect(native.claims).not.toHaveBeenCalled();
expect(deps.executeRun).not.toHaveBeenCalled();
expect(native.cleanups).not.toHaveBeenCalled();
});
it("keeps blocked and ambiguous ownership out of execution", async () => {
const deps = callbacks(db);
const dispositions: NativeRestartRecoveryDisposition[] = [
{ kind: "blocked", runId: randomUUID(), reason: "conflicting_owner" },
{ kind: "awaiting_evidence", runId: randomUUID(), reason: "missing_stop_evidence" },
];
native.claims.mockResolvedValue(dispositions);
const result = await createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart();
expect(result.blockedRunIds).toEqual([dispositions[0].runId]);
expect(result.awaitingEvidenceRunIds).toEqual([dispositions[1].runId]);
expect(result.claims).toEqual([]);
expect(deps.executeRun).not.toHaveBeenCalled();
await Promise.all([...deps.activeRunExecutionPromises]);
});
it("rearms the persisted native retry deadline through the service callback", async () => {
const run = await seedRun();
const nextAttemptAt = new Date(Date.now() + 60_000);
await db.update(heartbeatRuns).set({ runtimeMode: "native", nativeIssueId: run.issueId, status: "failed" }).where(eq(heartbeatRuns.id, run.id));
await db.insert(nativeRunFinalizations).values({ runId: run.id, companyId: run.companyId, issueId: run.issueId!, phase: "retryable_failure", nextAttemptAt });
const deps = callbacks(db);
const result = await createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart();
expect(result.scheduledRetryRunIds).toEqual([run.id]);
expect(deps.scheduleNativeSessionResumeDispatch).toHaveBeenCalledWith(run.id, nextAttemptAt);
expect(deps.executeRun).not.toHaveBeenCalled();
await Promise.all([...deps.activeRunExecutionPromises]);
await db.delete(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id));
await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, run.id));
});
it("retains shutdown ownership until background native cleanup physically settles", async () => {
const deps = callbacks(db);
let settle!: () => void;
native.drainMaintenance.mockReturnValueOnce(new Promise<void>(resolve => { settle = resolve; }));
await createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart();
await vi.waitFor(() => expect(native.drainMaintenance).toHaveBeenCalled());
expect(deps.activeRunExecutionPromises.size).toBe(1);
settle();
await Promise.all([...deps.activeRunExecutionPromises]);
await vi.waitFor(() => expect(deps.activeRunExecutionPromises.size).toBe(0));
});
});
File diff suppressed because it is too large. Load diff