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,