From d1038cc04fe3ad00bcae0c88479f9859be567e1a Mon Sep 17 00:00:00 2001 From: om952 Date: Thu, 2 Jul 2026 13:25:14 +0530 Subject: [PATCH] fix(server): drain live heartbeat runs on shutdown so same-agent checkout can reenter On a graceful SIGINT/SIGTERM shutdown, every still-live heartbeat-run row (queued / running / scheduled_retry) is now cancelled and the issue locks they hold (executionRunId, checkoutRunId) are cleared. - New drainStaleHeartbeatRunsOnShutdown service, wired into index.ts shutdown handler - Audit row issue.stale_lock_takeover on adoptStaleCheckoutRun success - Tests for shutdown drain and stale-lock takeover audit behavior Closes #30 --- .../heartbeat-shutdown-drain.test.ts | 216 +++++++++++++++++ ...ssues-checkout-stale-lock-takeover.test.ts | 220 ++++++++++++++++++ server/src/index.ts | 19 +- .../src/services/heartbeat-shutdown-drain.ts | 90 +++++++ server/src/services/issues.ts | 30 +++ 5 files changed, 574 insertions(+), 1 deletion(-) create mode 100644 server/src/__tests__/heartbeat-shutdown-drain.test.ts create mode 100644 server/src/__tests__/issues-checkout-stale-lock-takeover.test.ts create mode 100644 server/src/services/heartbeat-shutdown-drain.ts diff --git a/server/src/__tests__/heartbeat-shutdown-drain.test.ts b/server/src/__tests__/heartbeat-shutdown-drain.test.ts new file mode 100644 index 00000000000..2e1e9557553 --- /dev/null +++ b/server/src/__tests__/heartbeat-shutdown-drain.test.ts @@ -0,0 +1,216 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + agents, + companies, + createDb, + heartbeatRuns, + issues, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { drainStaleHeartbeatRunsOnShutdown } from "../services/heartbeat-shutdown-drain.js"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +if (!embeddedPostgresSupport.supported) { + console.warn( + `Skipping embedded Postgres heartbeat shutdown drain tests on this host: ${ + embeddedPostgresSupport.reason ?? "unsupported environment" + }`, + ); +} + +describeEmbeddedPostgres("drainStaleHeartbeatRunsOnShutdown", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("levi-heartbeat-shutdown-drain-"); + db = createDb(tempDb.connectionString); + }, 20_000); + + afterEach(async () => { + await db.delete(issues); + await db.delete(heartbeatRuns); + await db.delete(agents); + await db.delete(companies); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seed() { + const companyId = randomUUID(); + const agentId = randomUUID(); + await db.insert(companies).values({ + id: companyId, + name: "OpenScanAI", + issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + await db.insert(agents).values({ + id: agentId, + companyId, + name: "Producer", + role: "engineer", + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }); + return { companyId, agentId }; + } + + it("returns a zero result when no live runs exist", async () => { + const result = await drainStaleHeartbeatRunsOnShutdown(db); + expect(result).toEqual({ runsTerminated: 0, issuesUnlocked: 0 }); + }); + + it("cancels every live heartbeat run and clears the issue locks they hold", async () => { + const { companyId, agentId } = await seed(); + const liveRunId = randomUUID(); + const queuedRunId = randomUUID(); + const retryRunId = randomUUID(); + const succeededRunId = randomUUID(); + await db.insert(heartbeatRuns).values([ + { + id: liveRunId, + companyId, + agentId, + status: "running", + invocationSource: "manual", + startedAt: new Date(), + }, + { + id: queuedRunId, + companyId, + agentId, + status: "queued", + invocationSource: "manual", + }, + { + id: retryRunId, + companyId, + agentId, + status: "scheduled_retry", + invocationSource: "manual", + }, + { + id: succeededRunId, + companyId, + agentId, + status: "succeeded", + invocationSource: "manual", + finishedAt: new Date(), + }, + ]); + + const issueWithExecutionLock = randomUUID(); + const issueWithCheckoutOnlyLock = randomUUID(); + const issueOnTerminalRun = randomUUID(); + await db.insert(issues).values([ + { + id: issueWithExecutionLock, + companyId, + title: "Locked by execution", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + checkoutRunId: liveRunId, + executionRunId: liveRunId, + executionAgentNameKey: "producer", + executionLockedAt: new Date(), + }, + { + id: issueWithCheckoutOnlyLock, + companyId, + title: "Locked by checkout only", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + checkoutRunId: queuedRunId, + executionRunId: null, + }, + { + id: issueOnTerminalRun, + companyId, + title: "Already terminal", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + checkoutRunId: succeededRunId, + executionRunId: succeededRunId, + }, + ]); + + const result = await drainStaleHeartbeatRunsOnShutdown(db); + expect(result).toEqual({ runsTerminated: 3, issuesUnlocked: 2 }); + + const liveCount = await db + .select({ id: heartbeatRuns.id }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.status, "running")) + .then((rows) => rows.length); + expect(liveCount).toBe(0); + + const cancelledRow = await db + .select({ + status: heartbeatRuns.status, + errorCode: heartbeatRuns.errorCode, + finishedAt: heartbeatRuns.finishedAt, + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, liveRunId)) + .then((rows) => rows[0]); + expect(cancelledRow?.status).toBe("cancelled"); + expect(cancelledRow?.errorCode).toBe("server_shutdown_stale_lock_cleanup"); + expect(cancelledRow?.finishedAt).not.toBeNull(); + + const executionLockedRow = await db + .select({ + checkoutRunId: issues.checkoutRunId, + executionRunId: issues.executionRunId, + executionAgentNameKey: issues.executionAgentNameKey, + executionLockedAt: issues.executionLockedAt, + }) + .from(issues) + .where(eq(issues.id, issueWithExecutionLock)) + .then((rows) => rows[0]); + expect(executionLockedRow).toEqual({ + checkoutRunId: null, + executionRunId: null, + executionAgentNameKey: null, + executionLockedAt: null, + }); + + const checkoutOnlyRow = await db + .select({ + checkoutRunId: issues.checkoutRunId, + executionRunId: issues.executionRunId, + }) + .from(issues) + .where(eq(issues.id, issueWithCheckoutOnlyLock)) + .then((rows) => rows[0]); + expect(checkoutOnlyRow).toEqual({ checkoutRunId: null, executionRunId: null }); + + const terminalRow = await db + .select({ + checkoutRunId: issues.checkoutRunId, + executionRunId: issues.executionRunId, + }) + .from(issues) + .where(eq(issues.id, issueOnTerminalRun)) + .then((rows) => rows[0]); + expect(terminalRow).toEqual({ + checkoutRunId: succeededRunId, + executionRunId: succeededRunId, + }); + }); +}); diff --git a/server/src/__tests__/issues-checkout-stale-lock-takeover.test.ts b/server/src/__tests__/issues-checkout-stale-lock-takeover.test.ts new file mode 100644 index 00000000000..eb2a2684921 --- /dev/null +++ b/server/src/__tests__/issues-checkout-stale-lock-takeover.test.ts @@ -0,0 +1,220 @@ +import { randomUUID } from "node:crypto"; +import { eq, and } from "drizzle-orm"; +import { afterAll, afterEach, beforeAll, describe, expect, it } from "vitest"; +import { + activityLog, + agents, + companies, + createDb, + heartbeatRuns, + instanceSettings, + issueComments, + issueInboxArchives, + issueRelations, + issues, +} from "@paperclipai/db"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; +import { issueService } from "../services/issues.js"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +if (!embeddedPostgresSupport.supported) { + console.warn( + `Skipping embedded Postgres stale-lock takeover tests on this host: ${ + embeddedPostgresSupport.reason ?? "unsupported environment" + }`, + ); +} + +describeEmbeddedPostgres("issueService.checkout stale-lock takeover", () => { + let db!: ReturnType; + let svc!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("levi-issues-stale-lock-takeover-"); + db = createDb(tempDb.connectionString); + svc = issueService(db); + }, 20_000); + + afterEach(async () => { + await db.delete(issueComments); + await db.delete(issueRelations); + await db.delete(issueInboxArchives); + await db.delete(activityLog); + await db.delete(issues); + await db.delete(heartbeatRuns); + await db.delete(agents); + await db.delete(instanceSettings); + await db.delete(companies); + }); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function seedSameAgentReentryScene() { + const companyId = randomUUID(); + const agentId = randomUUID(); + const priorRunId = randomUUID(); + const newRunId = randomUUID(); + + await db.insert(companies).values({ + id: companyId, + name: "OpenScanAI", + issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`, + requireBoardApprovalForNewAgents: false, + }); + await db.insert(agents).values({ + id: agentId, + companyId, + name: "Producer", + role: "engineer", + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }); + // Prior run already cancelled (e.g. by the graceful shutdown drain) but + // the issue row still references it as the active checkoutRunId. + // The new run is the live one the actor brings to the checkout call. + await db.insert(heartbeatRuns).values([ + { + id: priorRunId, + companyId, + agentId, + status: "cancelled", + invocationSource: "manual", + finishedAt: new Date(), + errorCode: "server_shutdown_stale_lock_cleanup", + }, + { + id: newRunId, + companyId, + agentId, + status: "running", + invocationSource: "manual", + startedAt: new Date(), + }, + ]); + + const issueId = randomUUID(); + await db.insert(issues).values({ + id: issueId, + companyId, + title: "Producer-1 dogfood reentry", + status: "in_progress", + priority: "high", + assigneeAgentId: agentId, + checkoutRunId: priorRunId, + executionRunId: priorRunId, + executionAgentNameKey: "producer", + executionLockedAt: new Date(), + }); + + return { companyId, agentId, priorRunId, newRunId, issueId }; + } + + it("lets the same agent take over a stale lock and writes an audit row", async () => { + const { companyId, agentId, priorRunId, newRunId, issueId } = + await seedSameAgentReentryScene(); + + const result = await svc.checkout(issueId, agentId, ["in_progress"], newRunId); + expect(result?.id).toBe(issueId); + expect(result?.checkoutRunId).toBe(newRunId); + expect(result?.executionRunId).toBe(newRunId); + + const audit = await db + .select({ + action: activityLog.action, + actorType: activityLog.actorType, + actorId: activityLog.actorId, + agentId: activityLog.agentId, + runId: activityLog.runId, + entityType: activityLog.entityType, + entityId: activityLog.entityId, + details: activityLog.details, + }) + .from(activityLog) + .where( + and( + eq(activityLog.companyId, companyId), + eq(activityLog.entityId, issueId), + eq(activityLog.action, "issue.stale_lock_takeover"), + ), + ) + .then((rows) => rows[0]); + + expect(audit).toBeTruthy(); + expect(audit).toMatchObject({ + action: "issue.stale_lock_takeover", + actorType: "agent", + actorId: agentId, + agentId, + runId: newRunId, + entityType: "issue", + entityId: issueId, + details: { + priorCheckoutRunId: priorRunId, + actorRunId: newRunId, + source: "adoptStaleCheckoutRun", + }, + }); + }); + + it("refuses to take over the lock for a different (non-assignee) agent — cross-agent regression", async () => { + const { companyId, priorRunId, newRunId, issueId } = await seedSameAgentReentryScene(); + + const intruderId = randomUUID(); + await db.insert(agents).values({ + id: intruderId, + companyId, + name: "Intruder", + role: "engineer", + status: "active", + adapterType: "codex_local", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }); + + await expect( + svc.checkout(issueId, intruderId, ["in_progress"], newRunId), + ).rejects.toMatchObject({ status: 409 }); + + const row = await db + .select({ + assigneeAgentId: issues.assigneeAgentId, + checkoutRunId: issues.checkoutRunId, + executionRunId: issues.executionRunId, + }) + .from(issues) + .where(eq(issues.id, issueId)) + .then((rows) => rows[0]); + // The intruder must not own the checkout lock; the assignee/lock binding + // is preserved. (`executionRunId` may legitimately be cleared because the + // prior run is in a terminal state — the gate is the `checkoutRunId` + // assignee binding, not the execution lock byproduct.) + expect(row?.assigneeAgentId).not.toBe(intruderId); + expect(row?.checkoutRunId).toBe(priorRunId); + expect(row?.checkoutRunId).not.toBe(newRunId); + + const audit = await db + .select({ id: activityLog.id }) + .from(activityLog) + .where( + and( + eq(activityLog.companyId, companyId), + eq(activityLog.entityId, issueId), + eq(activityLog.action, "issue.stale_lock_takeover"), + ), + ) + .then((rows) => rows[0]); + expect(audit).toBeUndefined(); + }); +}); diff --git a/server/src/index.ts b/server/src/index.ts index 105f23ecaf7..f8a0b4d8907 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -7,6 +7,7 @@ import { stdin, stdout } from "node:process"; import { pathToFileURL } from "node:url"; import type { Request as ExpressRequest, RequestHandler } from "express"; import { and, eq } from "drizzle-orm"; +import type { Db } from "@paperclipai/db"; import { createDb, ensurePostgresDatabase, @@ -36,6 +37,7 @@ import { routineService, } from "./services/index.js"; import { createFeedbackTraceShareClientFromConfig } from "./services/feedback-share-client.js"; +import { drainStaleHeartbeatRunsOnShutdown } from "./services/heartbeat-shutdown-drain.js"; import { buildRuntimeApiCandidateUrls, choosePrimaryRuntimeApiUrl } from "./runtime-api.js"; import { createPluginWorkerManager } from "./services/plugin-worker-manager.js"; import { createStorageServiceFromConfig } from "./storage/index.js"; @@ -260,7 +262,7 @@ export async function startServer(): Promise { } } - let db; + let db: ReturnType; let pluginMigrationDb; let embeddedPostgres: EmbeddedPostgresInstance | null = null; let embeddedPostgresStartedByThisProcess = false; @@ -878,6 +880,21 @@ export async function startServer(): Promise { await telemetryClient.flush(); } + // Mark every still-live heartbeat run as cancelled and clear the + // `executionRunId` / `checkoutRunId` locks they hold on issues. Without + // this, a clean shutdown leaves the run row in "running" status and the + // next checkout from the same agent hits a 409 because the prior lock + // is still considered live. Bounded to 5s so a DB hang cannot wedge + // process exit; best-effort logging only. + try { + await Promise.race([ + drainStaleHeartbeatRunsOnShutdown(db), + new Promise((resolve) => setTimeout(() => resolve(null), 5000)), + ]); + } catch (err) { + logger.warn({ err, signal }, "stale heartbeat-run drain on shutdown threw"); + } + if (embeddedPostgres && embeddedPostgresStartedByThisProcess) { logger.info({ signal }, "Stopping embedded PostgreSQL"); try { diff --git a/server/src/services/heartbeat-shutdown-drain.ts b/server/src/services/heartbeat-shutdown-drain.ts new file mode 100644 index 00000000000..88a53af4e6a --- /dev/null +++ b/server/src/services/heartbeat-shutdown-drain.ts @@ -0,0 +1,90 @@ +import { inArray } from "drizzle-orm"; +import type { Db } from "@paperclipai/db"; +import { heartbeatRuns, issues } from "@paperclipai/db"; +import { logger } from "../middleware/logger.js"; + +const LIVE_HEARTBEAT_RUN_STATUSES = ["queued", "running", "scheduled_retry"] as const; + +const SHUTDOWN_DRAIN_ERROR = "Server shutdown — stale heartbeat-run lock cleanup."; +const SHUTDOWN_DRAIN_ERROR_CODE = "server_shutdown_stale_lock_cleanup"; + +export interface DrainStaleHeartbeatRunsResult { + runsTerminated: number; + issuesUnlocked: number; +} + +/** + * On graceful shutdown (SIGINT/SIGTERM), terminate every still-live + * heartbeat run row (queued / running / scheduled_retry) and clear the + * `executionRunId` / `checkoutRunId` locks pointing at those runs on the + * `issues` table. Without this, a clean shutdown leaves the run row in + * "running" status and the next checkout from the same agent hits a 409 + * because the prior lock is still considered live. + * + * Best-effort: returns `null` on any DB failure and never throws. + */ +export async function drainStaleHeartbeatRunsOnShutdown( + db: Db, +): Promise { + try { + const now = new Date(); + const liveRuns = await db + .select({ id: heartbeatRuns.id }) + .from(heartbeatRuns) + .where(inArray(heartbeatRuns.status, [...LIVE_HEARTBEAT_RUN_STATUSES])); + if (liveRuns.length === 0) { + return { runsTerminated: 0, issuesUnlocked: 0 }; + } + const liveRunIds = liveRuns.map((r) => r.id); + + const unlockedByExecution = await db + .update(issues) + .set({ + checkoutRunId: null, + executionRunId: null, + executionAgentNameKey: null, + executionLockedAt: null, + updatedAt: now, + }) + .where(inArray(issues.executionRunId, liveRunIds)) + .returning({ id: issues.id }); + + // Also clear `checkoutRunId` locks pointing at these runs but whose + // `executionRunId` was already null (so the previous UPDATE skipped them). + const unlockedByCheckoutOnly = await db + .update(issues) + .set({ + checkoutRunId: null, + updatedAt: now, + }) + .where(inArray(issues.checkoutRunId, liveRunIds)) + .returning({ id: issues.id }); + + const terminated = await db + .update(heartbeatRuns) + .set({ + status: "cancelled", + finishedAt: now, + error: SHUTDOWN_DRAIN_ERROR, + errorCode: SHUTDOWN_DRAIN_ERROR_CODE, + updatedAt: now, + }) + .where(inArray(heartbeatRuns.id, liveRunIds)) + .returning({ id: heartbeatRuns.id }); + + const result: DrainStaleHeartbeatRunsResult = { + runsTerminated: terminated.length, + issuesUnlocked: unlockedByExecution.length + unlockedByCheckoutOnly.length, + }; + logger.info(result, "Cleared stale heartbeat-run locks on shutdown"); + return result; + } catch (err) { + logger.warn({ err }, "Stale heartbeat-run drain on shutdown failed"); + return null; + } +} + +export const __INTERNAL_FOR_TESTS = { + LIVE_HEARTBEAT_RUN_STATUSES, + SHUTDOWN_DRAIN_ERROR_CODE, +}; diff --git a/server/src/services/issues.ts b/server/src/services/issues.ts index ae349906a60..5ed8fd6cc21 100644 --- a/server/src/services/issues.ts +++ b/server/src/services/issues.ts @@ -52,6 +52,7 @@ import { } from "@paperclipai/shared"; import { conflict, HttpError, notFound, unprocessable } from "../errors.js"; import { logger } from "../middleware/logger.js"; +import { logActivity } from "./activity-log.js"; import { parseObject } from "../adapters/utils.js"; import { defaultIssueExecutionWorkspaceSettingsForProject, @@ -3297,6 +3298,7 @@ export function issueService(db: Db) { ) .returning({ id: issues.id, + companyId: issues.companyId, status: issues.status, assigneeAgentId: issues.assigneeAgentId, checkoutRunId: issues.checkoutRunId, @@ -3304,6 +3306,34 @@ export function issueService(db: Db) { }) .then((rows) => rows[0] ?? null); + if (adopted) { + // Audit trail for stale-lock takeover so operators can reconstruct what + // happened when an issue silently changed `checkoutRunId` ownership. + // Best-effort: a logging failure must never block checkout adoption. + try { + await logActivity(db, { + companyId: adopted.companyId, + actorType: "agent", + actorId: input.actorAgentId, + agentId: input.actorAgentId, + runId: input.actorRunId, + action: "issue.stale_lock_takeover", + entityType: "issue", + entityId: input.issueId, + details: { + priorCheckoutRunId: input.expectedCheckoutRunId, + actorRunId: input.actorRunId, + source: "adoptStaleCheckoutRun", + }, + }); + } catch (err) { + logger.warn( + { err, issueId: input.issueId }, + "stale lock takeover audit log failed", + ); + } + } + return adopted; }