From 7dcdd6ae1bbc7d0329bcc8002a141a00db4eafe2 Mon Sep 17 00:00:00 2001 From: nickyleach <331803+nickyleach@users.noreply.github.com> Date: Wed, 16 Sep 2026 08:20:27 +0000 Subject: [PATCH] fix(server): make the task-drain termination audit write idempotent A connection can drop after the database commits the outcome transaction but before the driver reports success. The retry around that transaction could not tell the two cases apart, so it could replay a committed write and create a duplicate outcome row and a duplicate publication. Give each termination a stable key from its generation number, and check for an existing row under that key before every write attempt, including a retry. The write proceeds only when the check finds none. Co-authored-by: Paperclip --- .../cloud-control-task-drain.test.ts | 16 ++++- .../instance-settings-routes.test.ts | 59 +++++++++++++++++++ server/src/routes/instance-settings.ts | 41 +++++++++---- 3 files changed, 105 insertions(+), 11 deletions(-) diff --git a/server/src/__tests__/cloud-control-task-drain.test.ts b/server/src/__tests__/cloud-control-task-drain.test.ts index d4b0fcf9247..36789c46792 100644 --- a/server/src/__tests__/cloud-control-task-drain.test.ts +++ b/server/src/__tests__/cloud-control-task-drain.test.ts @@ -363,7 +363,21 @@ describe("the termination option through the production chain", () => { // pass it through the real cloudControlMiddleware and the real task-drain // route, so they pin the board's decision that a task-drain:start assertion // may also set terminateActiveTasks, with no extra permission. - const mockDb = { transaction: vi.fn((fn: (tx: unknown) => Promise) => fn({})) }; + // The termination deadline handler checks for an already-committed + // outcome row before it writes one (see instance-settings.ts), so this + // mock db needs a chainable select().from().where().limit() alongside + // transaction(). Every test here wants the write to proceed, so the + // check always reports no prior row. + const mockDb = { + transaction: vi.fn((fn: (tx: unknown) => Promise) => fn({})), + select: vi.fn(() => ({ + from: () => ({ + where: () => ({ + limit: () => Promise.resolve([] as { id: string }[]), + }), + }), + })), + }; beforeEach(() => { resetCloudControlReplayFenceForTests(); diff --git a/server/src/__tests__/instance-settings-routes.test.ts b/server/src/__tests__/instance-settings-routes.test.ts index 9923bf002bd..61a3a5b3644 100644 --- a/server/src/__tests__/instance-settings-routes.test.ts +++ b/server/src/__tests__/instance-settings-routes.test.ts @@ -55,8 +55,24 @@ function defaultTransactionImplementation(fn: (tx: unknown) => Promise) // Module-scoped (not rebuilt per createApp call) so a test can assert how // many times a request opened a transaction — the task-drain audit writes // for every company must share ONE transaction, not one each. +// The termination deadline handler checks for an already-committed outcome +// row before it writes one (see instance-settings.ts), so the mock db needs +// a chainable select().from().where().limit() alongside transaction(). A +// test that wants to model "a prior attempt already committed" overrides +// this default resolved value; every other test keeps the default empty +// result, so the write proceeds exactly as it did before that check existed. +function defaultSelectChain() { + return { + from: () => ({ + where: () => ({ + limit: () => Promise.resolve([] as { id: string }[]), + }), + }), + }; +} const mockDb = { transaction: vi.fn(defaultTransactionImplementation), + select: vi.fn(defaultSelectChain), }; describe("instance settings routes", () => { @@ -89,6 +105,8 @@ describe("instance settings routes", () => { // next one. mockDb.transaction.mockReset(); mockDb.transaction.mockImplementation(defaultTransactionImplementation); + mockDb.select.mockReset(); + mockDb.select.mockImplementation(defaultSelectChain); mockInstanceSettingsService.get.mockReset(); mockInstanceSettingsService.getGeneral.mockReset(); mockInstanceSettingsService.getExperimental.mockReset(); @@ -1918,6 +1936,47 @@ describe("instance settings routes", () => { expect(outcomeCalls).toHaveLength(2); }); + it("does_not_duplicate_the_outcome_row_when_a_retry_finds_a_prior_attempt_already_committed", 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, + }); + mockHeartbeatService.applyTaskDrain.mockReturnValue(1); + mockHeartbeatService.isTaskDrainGenerationLive.mockReturnValue(true); + mockHeartbeatService.terminateActiveRunsForTaskDrain.mockResolvedValue( + new Map([["company-1", { attemptedRunIds: ["run-1"], cancelledRunIds: ["run-1"], failedRunIds: [] }]]), + ); + const app = await createApp(adminActor); + await request(app).post("/api/instance/task-drain").send({ ttlMs: 60_000, terminateActiveTasks: true }); + mockLogActivity.mockClear(); + mockDb.transaction.mockClear(); + mockInstanceSettingsService.listCompanyIds.mockClear(); + + // A connection can drop after the server commits the outcome + // transaction but before the driver reports success back. Model + // that: an existing row for this exact termination is already in + // the table when the deadline handler runs, exactly as it would be + // after that kind of ambiguous failure. The write must not repeat. + mockDb.select.mockImplementation(() => ({ + from: () => ({ + where: () => ({ + limit: () => Promise.resolve([{ id: "activity-already-written" }]), + }), + }), + })); + + await vi.advanceTimersByTimeAsync(30_000); + + expect(mockDb.transaction).not.toHaveBeenCalled(); + expect(mockLogActivity).not.toHaveBeenCalled(); + expect(mockInstanceSettingsService.listCompanyIds).not.toHaveBeenCalled(); + }); + 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 42599db333d..b1d5a8b7d29 100644 --- a/server/src/routes/instance-settings.ts +++ b/server/src/routes/instance-settings.ts @@ -1,5 +1,6 @@ import { Router, type Request } from "express"; -import type { Db } from "@paperclipai/db"; +import { and, eq } from "drizzle-orm"; +import { activityLog, type Db } from "@paperclipai/db"; import { patchInstanceSettingsSchema, patchInstanceExperimentalSettingsSchema, @@ -198,16 +199,36 @@ export function instanceSettingsRoutes(db: Db) { const outcomesByCompany = await heartbeat.terminateActiveRunsForTaskDrain( "Stopped by the task drain termination deadline", ); - // The runs above are already cancelled; there is no undo for that - // step. A transient database failure here must not lose the outcome - // of a destructive action, so retry the company list read and the - // audit write together on a transient connection drop (an idempotent - // pair: no row commits unless every company's row commits). If the - // retries are still not enough, log the full outcome so an operator - // can recover it from the application log even though the audit - // table write itself did not land. + // The runs above are already cancelled. There is no undo for that + // step, so a transient database failure here must not lose the + // outcome. Retry the company list read and the audit write together + // on a transient connection drop. + // + // 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 + // 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. + // + // 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}`; try { await retryOnTransientDbConnectionError(async () => { + const alreadyRecorded = await db + .select({ id: activityLog.id }) + .from(activityLog) + .where(and( + eq(activityLog.action, "instance.task_drain.active_tasks_terminated"), + eq(activityLog.entityType, "instance_settings"), + eq(activityLog.entityId, terminationEntityId), + )) + .limit(1); + if (alreadyRecorded.length > 0) return; + const companyIds = await svc.listCompanyIds(); const postCommitActivityPublications: ActivityPublication[] = []; await db.transaction((tx) => @@ -227,7 +248,7 @@ export function instanceSettingsRoutes(db: Db) { agentApiKeyId: initiatingActor.agentApiKeyId, action: "instance.task_drain.active_tasks_terminated", entityType: "instance_settings", - entityId: "default", + entityId: terminationEntityId, details: { terminateAt, executedAt,