diff --git a/server/src/__tests__/instance-settings-routes.test.ts b/server/src/__tests__/instance-settings-routes.test.ts index ad35cffa563..9923bf002bd 100644 --- a/server/src/__tests__/instance-settings-routes.test.ts +++ b/server/src/__tests__/instance-settings-routes.test.ts @@ -2,6 +2,7 @@ import express from "express"; import request from "supertest"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { hoistModuleGraph } from "./helpers/hoist-module-graph.js"; +import { logger } from "../middleware/logger.js"; const mockInstanceSettingsService = vi.hoisted(() => ({ get: vi.fn(), @@ -1797,6 +1798,167 @@ describe("instance settings routes", () => { expect(company1Details?.attemptedRunIds).toEqual(["run-a"]); expect(company2Details?.attemptedRunIds).toEqual(["run-b"]); }); + + // These two tests use the real clock, not vi.useFakeTimers(): a real + // POST/DELETE HTTP round trip followed by an advanced fake clock is + // unreliable in this suite (the fake timer queue does not resolve a + // second supertest request the same way it resolves the first). A + // short real delay proves the same behavior without that coupling. + it("re_arms_the_prior_deadline_when_a_replacing_post_fails_after_clearing_the_timer", async () => { + mockHeartbeatService.isTaskDrainGenerationLive.mockReturnValue(true); + const app = await createApp(adminActor); + + const firstTerminateAt = new Date(Date.now() + 100); + mockHeartbeatService.computeTaskDrain.mockReturnValueOnce({ + startedAt: new Date(), + expiresAt: new Date(Date.now() + 1000), + terminateActiveTasks: true, + terminateAt: firstTerminateAt, + }); + mockHeartbeatService.applyTaskDrain.mockReturnValueOnce(1); + const firstRes = await request(app) + .post("/api/instance/task-drain") + .send({ ttlMs: 1000, terminateActiveTasks: true }); + expect(firstRes.status).toBe(200); + + // The replacing POST clears the first timer, then its own audit + // write fails, so the request as a whole fails. + mockHeartbeatService.computeTaskDrain.mockReturnValueOnce({ + startedAt: new Date(), + expiresAt: new Date(Date.now() + 100_000), + terminateActiveTasks: true, + terminateAt: new Date(Date.now() + 100_000), + }); + mockLogActivity.mockRejectedValueOnce(new Error("activity insert failed")); + const secondRes = await request(app) + .post("/api/instance/task-drain") + .send({ ttlMs: 100_000, terminateActiveTasks: true }); + expect(secondRes.status).toBeGreaterThanOrEqual(500); + // The second drain never applied, so the first one is still live. + expect(mockHeartbeatService.applyTaskDrain).toHaveBeenCalledTimes(1); + + // The first drain's original deadline still fires: the failed + // replacement re-armed the timer it had cleared. + await new Promise((resolve) => setTimeout(resolve, 500)); + + expect(mockHeartbeatService.isTaskDrainGenerationLive).toHaveBeenCalledWith(1); + expect(mockHeartbeatService.terminateActiveRunsForTaskDrain).toHaveBeenCalledTimes(1); + }); + + it("re_arms_the_deadline_when_a_delete_fails_after_clearing_the_timer", async () => { + mockHeartbeatService.isTaskDrainGenerationLive.mockReturnValue(true); + const terminateAt = new Date(Date.now() + 100); + mockHeartbeatService.computeTaskDrain.mockReturnValue({ + startedAt: new Date(), + expiresAt: new Date(Date.now() + 1000), + terminateActiveTasks: true, + terminateAt, + }); + mockHeartbeatService.applyTaskDrain.mockReturnValue(1); + mockHeartbeatService.getTaskDrainStatus.mockReturnValue({ + draining: true, + startedAt: new Date(), + expiresAt: new Date(Date.now() + 1000), + terminateActiveTasks: true, + terminateAt, + activeRuns: 0, + pendingWakes: 0, + quiescent: true, + }); + const app = await createApp(adminActor); + await request(app).post("/api/instance/task-drain").send({ ttlMs: 1000, terminateActiveTasks: true }); + + // The DELETE clears the timer, then its own audit write fails, so + // the request as a whole fails and the drain never stops. + mockLogActivity.mockRejectedValueOnce(new Error("activity insert failed")); + const deleteRes = await request(app).delete("/api/instance/task-drain"); + expect(deleteRes.status).toBeGreaterThanOrEqual(500); + expect(mockHeartbeatService.stopTaskDrain).not.toHaveBeenCalled(); + + // The deadline still fires: the failed stop re-armed the timer it + // had cleared. + await new Promise((resolve) => setTimeout(resolve, 500)); + + expect(mockHeartbeatService.isTaskDrainGenerationLive).toHaveBeenCalledWith(1); + expect(mockHeartbeatService.terminateActiveRunsForTaskDrain).toHaveBeenCalledTimes(1); + }); + + it("retries_the_outcome_audit_write_on_a_transient_connection_drop_and_still_records_it", 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(); + + // The first company-list read after the runs are already cancelled + // hits a transient closed-connection error; the retry then + // succeeds, so the outcome still gets written. + const transientErr = Object.assign(new Error("write CONNECTION_CLOSED db.example.internal:5432"), { + code: "CONNECTION_CLOSED", + }); + mockInstanceSettingsService.listCompanyIds.mockRejectedValueOnce(transientErr); + + await vi.advanceTimersByTimeAsync(30_200); + + const outcomeCalls = mockLogActivity.mock.calls.filter( + ([, input]: [unknown, { action: string }]) => input.action === "instance.task_drain.active_tasks_terminated", + ); + expect(outcomeCalls).toHaveLength(2); + }); + + 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")); + 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(); + + const loggerErrorSpy = vi.spyOn(logger, "error").mockImplementation(() => logger); + // A non-transient failure (no closed-connection code) is not + // retried; the runs are already cancelled, so this must not throw + // out of the deadline handler unlogged. + mockInstanceSettingsService.listCompanyIds.mockRejectedValue(new Error("permanently unavailable")); + + await vi.advanceTimersByTimeAsync(30_000); + + const recoveryCall = loggerErrorSpy.mock.calls.find( + ([, message]) => + message === "task drain termination ran but its audit record failed to write; recover the outcome from this log entry", + ); + expect(recoveryCall).toBeDefined(); + expect(recoveryCall?.[0]).toMatchObject({ + generation: 1, + outcomesByCompany: { + "company-1": { attemptedRunIds: ["run-1"], cancelledRunIds: ["run-1"], failedRunIds: [] }, + }, + }); + expect(mockLogActivity).not.toHaveBeenCalled(); + loggerErrorSpy.mockRestore(); + }); }); }); }); diff --git a/server/src/routes/instance-settings.ts b/server/src/routes/instance-settings.ts index f7df1b9f2fa..42599db333d 100644 --- a/server/src/routes/instance-settings.ts +++ b/server/src/routes/instance-settings.ts @@ -7,6 +7,7 @@ import { startTaskDrainRequestSchema, } from "@paperclipai/shared"; import { forbidden } from "../errors.js"; +import { retryOnTransientDbConnectionError } from "../middleware/auth.js"; import { isCloudManagedInstance } from "../services/cloud-instance.js"; import { getHiddenSettings } from "../services/settings-visibility.js"; import { validate } from "../middleware/validate.js"; @@ -121,13 +122,6 @@ function withTaskDrainTransition(run: () => Promise): Promise { // that already fired. let taskDrainTerminationTimer: ReturnType | null = null; -function clearTaskDrainTerminationTimer() { - if (taskDrainTerminationTimer !== null) { - clearTimeout(taskDrainTerminationTimer); - taskDrainTerminationTimer = null; - } -} - // The initiating actor of an armed termination, held immutable from arm // time through to the outcome record. The operator who armed the // termination is the actor of the outcome, even though the outcome record @@ -142,6 +136,27 @@ type TaskDrainActorSnapshot = { cloudControlRequestId: string | null; }; +// The termination the current timer is armed for, mirrored next to the +// timer handle above. A drain replacement (POST) or a drain deletion +// (DELETE) clears the timer before it knows whether its own replacement +// 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. +type ArmedTaskDrainTermination = { + generation: number; + initiatingActor: TaskDrainActorSnapshot; + terminateAt: Date; +}; +let armedTaskDrainTermination: ArmedTaskDrainTermination | null = null; + +function clearTaskDrainTerminationTimer() { + if (taskDrainTerminationTimer !== null) { + clearTimeout(taskDrainTerminationTimer); + taskDrainTerminationTimer = null; + } + armedTaskDrainTermination = null; +} + function buildTaskDrainActorSnapshot(req: Request): TaskDrainActorSnapshot { const actor = getActorInfo(req); return { @@ -183,48 +198,97 @@ export function instanceSettingsRoutes(db: Db) { const outcomesByCompany = await heartbeat.terminateActiveRunsForTaskDrain( "Stopped by the task drain termination deadline", ); - const companyIds = await svc.listCompanyIds(); - const postCommitActivityPublications: ActivityPublication[] = []; - await db.transaction((tx) => - Promise.all( - companyIds.map((companyId) => { - const outcome = outcomesByCompany.get(companyId) ?? { - attemptedRunIds: [] as string[], - cancelledRunIds: [] as string[], - failedRunIds: [] as string[], - }; - return logActivity(tx as unknown as Db, { - companyId, - actorType: initiatingActor.actorType, - actorId: initiatingActor.actorId, - agentId: initiatingActor.agentId, - runId: initiatingActor.runId, - agentApiKeyId: initiatingActor.agentApiKeyId, - action: "instance.task_drain.active_tasks_terminated", - entityType: "instance_settings", - entityId: "default", - details: { - terminateAt, - executedAt, - initiatingActor, - attemptedRunIds: outcome.attemptedRunIds, - attemptedRunCount: outcome.attemptedRunIds.length, - cancelledRunIds: outcome.cancelledRunIds, - cancelledRunCount: outcome.cancelledRunIds.length, - failedRunIds: outcome.failedRunIds, - failedRunCount: outcome.failedRunIds.length, - }, - }, postCommitActivityPublications); - }), - ), - ); - publishActivitiesBestEffort( - postCommitActivityPublications, - "instance.task_drain.active_tasks_terminated", - ); + // 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. + try { + await retryOnTransientDbConnectionError(async () => { + const companyIds = await svc.listCompanyIds(); + const postCommitActivityPublications: ActivityPublication[] = []; + await db.transaction((tx) => + Promise.all( + companyIds.map((companyId) => { + const outcome = outcomesByCompany.get(companyId) ?? { + attemptedRunIds: [] as string[], + cancelledRunIds: [] as string[], + failedRunIds: [] as string[], + }; + return logActivity(tx as unknown as Db, { + companyId, + actorType: initiatingActor.actorType, + actorId: initiatingActor.actorId, + agentId: initiatingActor.agentId, + runId: initiatingActor.runId, + agentApiKeyId: initiatingActor.agentApiKeyId, + action: "instance.task_drain.active_tasks_terminated", + entityType: "instance_settings", + entityId: "default", + details: { + terminateAt, + executedAt, + initiatingActor, + attemptedRunIds: outcome.attemptedRunIds, + attemptedRunCount: outcome.attemptedRunIds.length, + cancelledRunIds: outcome.cancelledRunIds, + cancelledRunCount: outcome.cancelledRunIds.length, + failedRunIds: outcome.failedRunIds, + failedRunCount: outcome.failedRunIds.length, + }, + }, postCommitActivityPublications); + }), + ), + ); + publishActivitiesBestEffort( + postCommitActivityPublications, + "instance.task_drain.active_tasks_terminated", + ); + }); + } catch (err) { + logger.error( + { + err, + generation, + terminateAt, + executedAt, + initiatingActor, + outcomesByCompany: Object.fromEntries(outcomesByCompany), + }, + "task drain termination ran but its audit record failed to write; recover the outcome from this log entry", + ); + } }); } + // Arm the termination timer for the given deadline, and mirror it in + // armedTaskDrainTermination so a later failed replacement or failed stop + // can re-arm the same deadline instead of leaving a live drain with no + // timer. + function scheduleTaskDrainTermination(armed: ArmedTaskDrainTermination) { + armedTaskDrainTermination = armed; + const delayMs = Math.max(0, armed.terminateAt.getTime() - Date.now()); + taskDrainTerminationTimer = setTimeout(() => { + // Only clear the shared snapshot when it still names this exact + // callback. A callback the event loop already queued before a + // later request replaced it (see the comment above + // taskDrainTerminationTimer) must not clear a newer request's + // timer out from under it; isTaskDrainGenerationLive, inside + // runTaskDrainTermination, is what decides whether this callback + // still gets to act. + if (armedTaskDrainTermination === armed) { + taskDrainTerminationTimer = null; + armedTaskDrainTermination = null; + } + runTaskDrainTermination(armed.generation, armed.initiatingActor, armed.terminateAt).catch((err) => { + logger.error({ err }, "task drain termination failed"); + }); + }, delayMs); + } + router.get("/instance/settings", async (req, res) => { assertBoardOrgAccess(req); res.json(await svc.get()); @@ -421,57 +485,64 @@ export function instanceSettingsRoutes(db: Db) { // A repeated start replaces whatever termination the prior drain // armed. Clear it before this turn computes and applies the new // drain, so a stale timer can never fire alongside — or instead - // of — the one this turn arms below. + // of — the one this turn arms below. Keep a snapshot of what this + // clear just disarmed: if a later step in this same request then + // fails, the catch block below re-arms it from the snapshot, so a + // failed replacement never leaves the prior, still-live drain with + // no timer. + const previousArmedTermination = armedTaskDrainTermination; clearTaskDrainTerminationTimer(); - const computed = heartbeat.computeTaskDrain({ ttlMs, terminateActiveTasks }); - // One transaction for every company's audit row, so a write that - // succeeds for one company and fails for another never leaves a - // partial activity history behind — either every company gets the - // record, or none does. The drain mutation below runs only after - // this transaction commits, so a failed write leaves the live - // drain untouched and there is no partial state to roll back. - const postCommitActivityPublications: ActivityPublication[] = []; - await db.transaction((tx) => - Promise.all( - companyIds.map((companyId) => - logActivity(tx as unknown as Db, { - companyId, - actorType: initiatingActor.actorType, - actorId: initiatingActor.actorId, - agentId: initiatingActor.agentId, - runId: initiatingActor.runId, - agentApiKeyId: initiatingActor.agentApiKeyId, - action: "instance.task_drain.started", - entityType: "instance_settings", - entityId: "default", - details: { - terminateActiveTasks: computed.terminateActiveTasks, - startedAt: computed.startedAt, - expiresAt: computed.expiresAt, - terminateAt: computed.terminateAt, - initiatingActor, - }, - }, postCommitActivityPublications), + try { + const computed = heartbeat.computeTaskDrain({ ttlMs, terminateActiveTasks }); + // One transaction for every company's audit row, so a write that + // succeeds for one company and fails for another never leaves a + // partial activity history behind — either every company gets the + // record, or none does. The drain mutation below runs only after + // this transaction commits, so a failed write leaves the live + // drain untouched and there is no partial state to roll back. + const postCommitActivityPublications: ActivityPublication[] = []; + await db.transaction((tx) => + Promise.all( + companyIds.map((companyId) => + logActivity(tx as unknown as Db, { + companyId, + actorType: initiatingActor.actorType, + actorId: initiatingActor.actorId, + agentId: initiatingActor.agentId, + runId: initiatingActor.runId, + agentApiKeyId: initiatingActor.agentApiKeyId, + action: "instance.task_drain.started", + entityType: "instance_settings", + entityId: "default", + details: { + terminateActiveTasks: computed.terminateActiveTasks, + startedAt: computed.startedAt, + expiresAt: computed.expiresAt, + terminateAt: computed.terminateAt, + initiatingActor, + }, + }, postCommitActivityPublications), + ), ), - ), - ); - const generation = heartbeat.applyTaskDrain(computed); - // The audit record already committed, so a failure to publish it - // here is not a reason to undo the drain: reverting the in-memory - // state at this point would desync it from the committed row. - // Swallow a publish failure so it cannot turn a committed mutation - // into a false 500. - publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.started"); - const terminateAt = computed.terminateAt; - if (terminateAt) { - const delayMs = Math.max(0, terminateAt.getTime() - Date.now()); - taskDrainTerminationTimer = setTimeout(() => { - runTaskDrainTermination(generation, initiatingActor, terminateAt).catch((err) => { - logger.error({ err }, "task drain termination failed"); - }); - }, delayMs); + ); + const generation = heartbeat.applyTaskDrain(computed); + // The audit record already committed, so a failure to publish it + // here is not a reason to undo the drain: reverting the in-memory + // state at this point would desync it from the committed row. + // Swallow a publish failure so it cannot turn a committed mutation + // into a false 500. + publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.started"); + const terminateAt = computed.terminateAt; + if (terminateAt) { + scheduleTaskDrainTermination({ generation, initiatingActor, terminateAt }); + } + return computed; + } catch (err) { + if (previousArmedTermination) { + scheduleTaskDrainTermination(previousArmedTermination); + } + throw err; } - return computed; }); res.json(drain); }, @@ -490,43 +561,54 @@ export function instanceSettingsRoutes(db: Db) { const wasActive = await withTaskDrainTransition(async () => { // A stop must cancel a termination the current drain armed, before // the timer's own generation check would otherwise have to catch it. + // Keep a snapshot of what this clear just disarmed: see the POST + // handler above for why a later step's failure re-arms it from the + // snapshot instead of leaving the still-live drain with no timer. + const previousArmedTermination = armedTaskDrainTermination; clearTaskDrainTerminationTimer(); - const priorStatus = heartbeat.getTaskDrainStatus(); - // Read wasActive once, here, and use this same value for the audit - // detail and the response body below. A TTL that expires between two - // separate reads would otherwise make the two values disagree. - const wasActive = priorStatus.draining; - // See the POST handler above for why this is one transaction, and why - // the drain mutation runs only after it commits. - const postCommitActivityPublications: ActivityPublication[] = []; - await db.transaction((tx) => - Promise.all( - companyIds.map((companyId) => - logActivity(tx as unknown as Db, { - companyId, - actorType: actor.actorType, - actorId: actor.actorId, - agentId: actor.agentId, - runId: actor.runId, - agentApiKeyId: actor.agentApiKeyId, - action: "instance.task_drain.stopped", - entityType: "instance_settings", - entityId: "default", - details: { - wasActive, - }, - }, postCommitActivityPublications), + try { + const priorStatus = heartbeat.getTaskDrainStatus(); + // Read wasActive once, here, and use this same value for the audit + // detail and the response body below. A TTL that expires between two + // separate reads would otherwise make the two values disagree. + const wasActive = priorStatus.draining; + // See the POST handler above for why this is one transaction, and why + // the drain mutation runs only after it commits. + const postCommitActivityPublications: ActivityPublication[] = []; + await db.transaction((tx) => + Promise.all( + companyIds.map((companyId) => + logActivity(tx as unknown as Db, { + companyId, + actorType: actor.actorType, + actorId: actor.actorId, + agentId: actor.agentId, + runId: actor.runId, + agentApiKeyId: actor.agentApiKeyId, + action: "instance.task_drain.stopped", + entityType: "instance_settings", + entityId: "default", + details: { + wasActive, + }, + }, postCommitActivityPublications), + ), ), - ), - ); - heartbeat.stopTaskDrain(); - // See the POST handler above for why a publish failure here is - // swallowed instead of failing the route: the audit record already - // committed, so a publish failure here must not undo a drain-stop - // that is already correct in the database, and must not report the - // stop as failed when it succeeded. - publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.stopped"); - return wasActive; + ); + heartbeat.stopTaskDrain(); + // See the POST handler above for why a publish failure here is + // swallowed instead of failing the route: the audit record already + // committed, so a publish failure here must not undo a drain-stop + // that is already correct in the database, and must not report the + // stop as failed when it succeeded. + publishActivitiesBestEffort(postCommitActivityPublications, "instance.task_drain.stopped"); + return wasActive; + } catch (err) { + if (previousArmedTermination) { + scheduleTaskDrainTermination(previousArmedTermination); + } + throw err; + } }); res.json({ wasActive }); });