diff --git a/server/src/__tests__/instance-settings-routes.test.ts b/server/src/__tests__/instance-settings-routes.test.ts index 61a3a5b3644..70da7c9f6bd 100644 --- a/server/src/__tests__/instance-settings-routes.test.ts +++ b/server/src/__tests__/instance-settings-routes.test.ts @@ -70,6 +70,24 @@ function defaultSelectChain() { }), }; } +// Read the literal string value a drizzle-orm `eq()`/`and()` condition +// binds as a query parameter. A test that must prove which dedup key the +// route actually queried with — not just whether a row came back — walks +// the condition object for its bound `Param` values, since the mock +// `where()` above normally ignores its argument. +function collectBoundParamValues(node: unknown, found: string[]): void { + if (node === null || typeof node !== "object") return; + const record = node as Record; + if (typeof record.value === "string") found.push(record.value); + if (Array.isArray(record.value)) { + for (const entry of record.value) { + if (typeof entry === "string") found.push(entry); + } + } + if (Array.isArray(record.queryChunks)) { + for (const chunk of record.queryChunks) collectBoundParamValues(chunk, found); + } +} const mockDb = { transaction: vi.fn(defaultTransactionImplementation), select: vi.fn(defaultSelectChain), @@ -1977,6 +1995,71 @@ describe("instance settings routes", () => { expect(mockInstanceSettingsService.listCompanyIds).not.toHaveBeenCalled(); }); + it("does_not_treat_a_pre_restart_activity_row_as_a_duplicate_of_a_same_numbered_generation", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z")); + const terminateAt = new Date("2026-01-01T00:00:30.000Z"); + mockHeartbeatService.computeTaskDrain.mockReturnValue({ + startedAt: new Date("2026-01-01T00:00:00.000Z"), + expiresAt: new Date("2026-01-01T00:01:00.000Z"), + terminateActiveTasks: true, + terminateAt, + }); + // The in-memory generation counter starts at 0 on every process + // boot, so a fresh process can hand out generation 1 again. Model + // that here: this run's own generation is 1, standing in for a + // process restart after an earlier, unrelated termination also + // held generation 1. + mockHeartbeatService.applyTaskDrain.mockReturnValue(1); + mockHeartbeatService.isTaskDrainGenerationLive.mockReturnValue(true); + mockHeartbeatService.terminateActiveRunsForTaskDrain.mockResolvedValue( + new Map([["company-1", { attemptedRunIds: ["run-1"], cancelledRunIds: ["run-1"], failedRunIds: [] }]]), + ); + + // A row already sits in the table under the dedup key a + // generation-only scheme would have used for generation 1. It + // belongs to that earlier, unrelated termination — a real + // termination from before the restart — not to this run. + const preRestartEntityId = "task-drain-termination:1"; + let queriedEntityId: string | null = null; + mockDb.select.mockImplementation(() => ({ + from: () => ({ + where: (condition: unknown) => ({ + limit: () => { + const boundValues: string[] = []; + collectBoundParamValues(condition, boundValues); + queriedEntityId = + boundValues.find((value) => value.startsWith("task-drain-termination:")) ?? null; + const matchesPreRestartRow = queriedEntityId === preRestartEntityId; + return Promise.resolve( + matchesPreRestartRow ? [{ id: "activity-from-before-restart" }] : [], + ); + }, + }), + }), + })); + + const app = await createApp(adminActor); + await request(app).post("/api/instance/task-drain").send({ ttlMs: 60_000, terminateActiveTasks: true }); + mockLogActivity.mockClear(); + mockDb.transaction.mockClear(); + + await vi.advanceTimersByTimeAsync(30_000); + + // The dedup key this run queried with must not be the plain + // generation number, and it must not match the pre-restart row — + // otherwise this run's own outcome audit would be silently + // skipped as an apparent duplicate of a termination it never ran. + expect(queriedEntityId).not.toBeNull(); + expect(queriedEntityId).not.toBe(preRestartEntityId); + expect(mockDb.transaction).toHaveBeenCalledTimes(1); + expect(mockLogActivity).toHaveBeenCalledWith( + TX_SENTINEL, + expect.objectContaining({ action: "instance.task_drain.active_tasks_terminated" }), + expect.any(Array), + ); + }); + it("logs_the_full_termination_outcome_when_the_audit_write_cannot_be_retried_into_success", async () => { vi.useFakeTimers(); vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z")); diff --git a/server/src/routes/instance-settings.ts b/server/src/routes/instance-settings.ts index b1d5a8b7d29..19092ef91a4 100644 --- a/server/src/routes/instance-settings.ts +++ b/server/src/routes/instance-settings.ts @@ -1,3 +1,4 @@ +import { randomUUID } from "node:crypto"; import { Router, type Request } from "express"; import { and, eq } from "drizzle-orm"; import { activityLog, type Db } from "@paperclipai/db"; @@ -143,8 +144,19 @@ type TaskDrainActorSnapshot = { // will succeed. If a later step in that same request then fails, the // request re-arms the deadline from this snapshot, so a failed // replacement or a failed stop never leaves a live drain with no timer. +// +// `generation` and `terminationId` answer two different questions. +// `generation` comes from the process-local counter in the heartbeat +// service, resets to 0 on every restart, and only proves whether this +// timer's callback still names the live in-memory drain. `terminationId` +// is a UUID minted once, when this termination is armed, and stays +// unique across process restarts, so the audit dedup key below can use +// it safely: an old activity row from before a restart can carry the +// same `generation` number as a new termination, but never the same +// `terminationId`. type ArmedTaskDrainTermination = { generation: number; + terminationId: string; initiatingActor: TaskDrainActorSnapshot; terminateAt: Date; }; @@ -190,6 +202,7 @@ export function instanceSettingsRoutes(db: Db) { // obsolete after the event loop had already queued it). async function runTaskDrainTermination( generation: number, + terminationId: string, initiatingActor: TaskDrainActorSnapshot, terminateAt: Date, ) { @@ -206,16 +219,20 @@ export function instanceSettingsRoutes(db: Db) { // // A dropped connection can happen after the server commits the // transaction but before the driver reports success. So a retry - // cannot assume the first attempt never landed. The termination - // generation names this exact execution and never repeats, so each + // cannot assume the first attempt never landed. `terminationId` + // names this exact execution and never repeats — not even across a + // process restart, unlike the in-memory generation number — so each // retry uses it as a key: it checks for an existing row first, and // writes only when that check finds none. This stops a retry after - // an ambiguous failure from duplicating the row or its publication. + // an ambiguous failure from duplicating the row or its publication, + // and it stops a termination after a restart from matching an old + // row left by an unrelated termination that reused the same + // generation number. // // If every retry still fails, log the full outcome. An operator can // then recover it from the application log even though the audit // table write did not land. - const terminationEntityId = `task-drain-termination:${generation}`; + const terminationEntityId = `task-drain-termination:${terminationId}`; try { await retryOnTransientDbConnectionError(async () => { const alreadyRecorded = await db @@ -274,6 +291,7 @@ export function instanceSettingsRoutes(db: Db) { { err, generation, + terminationId, terminateAt, executedAt, initiatingActor, @@ -304,7 +322,12 @@ export function instanceSettingsRoutes(db: Db) { taskDrainTerminationTimer = null; armedTaskDrainTermination = null; } - runTaskDrainTermination(armed.generation, armed.initiatingActor, armed.terminateAt).catch((err) => { + runTaskDrainTermination( + armed.generation, + armed.terminationId, + armed.initiatingActor, + armed.terminateAt, + ).catch((err) => { logger.error({ err }, "task drain termination failed"); }); }, delayMs); @@ -555,7 +578,12 @@ export function instanceSettingsRoutes(db: Db) { publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.started"); const terminateAt = computed.terminateAt; if (terminateAt) { - scheduleTaskDrainTermination({ generation, initiatingActor, terminateAt }); + scheduleTaskDrainTermination({ + generation, + terminationId: randomUUID(), + initiatingActor, + terminateAt, + }); } return computed; } catch (err) {