fix(server): preserve the task-drain termination deadline across a failed swap, and retry its audit write

A repeated drain start, or a drain stop, cleared the termination timer
before it knew whether its own audit write would succeed. A failure
after that clear left the prior, still-live drain with no timer, so an
active run could survive its advertised deadline. Snapshot the armed
termination before the clear and re-arm it if the replacement fails.

The deadline handler also committed run cancellations before it wrote
the outcome audit record. A failure in the company list read or the
audit transaction after that point silently dropped the record for an
already-irreversible action. Retry that step on a transient database
connection drop, and log the full outcome when the retries do not
recover it, so the result stays recoverable.

Co-authored-by: Paperclip <noreply@paperclip.ing>

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();
+      });
     });
   });
 });

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<T>(run: () => Promise<T>): Promise<T> {
 // that already fired.
 let taskDrainTerminationTimer: ReturnType<typeof setTimeout> | 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 });
   });