fix(cron): clear automation after failed scheduler shutdown (#115316)

* fix(cron): unregister stopped scheduler when draining fails

* test(cron): simplify deferred drain regression type

---------

Co-authored-by: Peter Steinberger <steipete@gmail.com>

src/gateway/server-cron.drain.test.ts

@@ -0,0 +1,134 @@
+import { mkdtemp, rm } from "node:fs/promises";
+import os from "node:os";
+import path from "node:path";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { CliDeps } from "../cli/deps.types.js";
+import type { OpenClawConfig } from "../config/types.openclaw.js";
+import { createDeferred } from "../test-utils/deferred.js";
+
+const { getRuntimeConfigMock, stopAllMock } = vi.hoisted(() => ({
+  getRuntimeConfigMock: vi.fn(),
+  stopAllMock: vi.fn<() => Promise<void>>(),
+}));
+
+vi.mock("../config/io.js", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("../config/io.js")>()),
+  getRuntimeConfig: getRuntimeConfigMock,
+}));
+
+vi.mock("./cron-stream-watchers.js", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("./cron-stream-watchers.js")>()),
+  createCronStreamWatchers: () => ({
+    reconcile: vi.fn(async () => {}),
+    resume: vi.fn(),
+    start: vi.fn(async () => {}),
+    stop: vi.fn(async () => {}),
+    stopAll: stopAllMock,
+    activeJobIds: () => [],
+    inspect: () => undefined,
+  }),
+}));
+
+import { buildGatewayCronService } from "./server-cron.js";
+import { sessionHasAutomation } from "./session-automation-index.js";
+
+type StartedGatewayCron = {
+  state: ReturnType<typeof buildGatewayCronService>;
+  cfg: OpenClawConfig;
+  stateDir: string;
+};
+
+async function startGatewayCron(label: string): Promise<StartedGatewayCron> {
+  const stateDir = await mkdtemp(path.join(os.tmpdir(), `openclaw-cron-drain-${label}-`));
+  const cfg: OpenClawConfig = {
+    session: { mainKey: "main" },
+    cron: { triggers: { enabled: true } },
+  };
+  getRuntimeConfigMock.mockReturnValue(cfg);
+  const state = buildGatewayCronService({
+    cfg,
+    deps: {} as CliDeps,
+    broadcast: () => {},
+    env: { ...process.env, OPENCLAW_SKIP_CRON: "0", OPENCLAW_STATE_DIR: stateDir },
+  });
+  await state.cron.start();
+  await state.cron.add({
+    name: `${label} stream source`,
+    enabled: true,
+    schedule: { kind: "stream", command: ["source"] },
+    payload: { kind: "systemEvent", text: "event" },
+    sessionTarget: "main",
+    wakeMode: "next-heartbeat",
+  });
+  return { state, cfg, stateDir };
+}
+
+async function cleanGatewayCron({ state, stateDir }: StartedGatewayCron): Promise<void> {
+  try {
+    await state.cron.stopAndDrain?.();
+  } finally {
+    await rm(stateDir, { recursive: true, force: true });
+  }
+}
+
+describe("gateway cron stop-and-drain automation ownership", () => {
+  beforeEach(() => {
+    getRuntimeConfigMock.mockReset();
+    stopAllMock.mockReset();
+  });
+
+  it("unregisters a stopped scheduler when stream draining fails and permits retry", async () => {
+    stopAllMock.mockRejectedValueOnce(new Error("stream drain failed"));
+    stopAllMock.mockResolvedValue(undefined);
+    const original = await startGatewayCron("failed");
+
+    try {
+      expect(sessionHasAutomation("agent:main:main", original.cfg)).toBe(true);
+
+      await expect(original.state.cron.stopAndDrain?.()).rejects.toThrow("stream drain failed");
+
+      expect(sessionHasAutomation("agent:main:main", original.cfg)).toBe(false);
+      await expect(original.state.cron.stopAndDrain?.()).resolves.toBeUndefined();
+      expect(stopAllMock).toHaveBeenCalledTimes(2);
+      expect(sessionHasAutomation("agent:main:main", original.cfg)).toBe(false);
+    } finally {
+      await cleanGatewayCron(original);
+    }
+  });
+
+  it("does not unregister a replacement scheduler when a stale drain fails", async () => {
+    const pendingDrain = createDeferred();
+    stopAllMock.mockImplementationOnce(() => pendingDrain.promise);
+    stopAllMock.mockResolvedValue(undefined);
+    const original = await startGatewayCron("stale");
+    let replacement: StartedGatewayCron | undefined;
+
+    try {
+      expect(sessionHasAutomation("agent:main:main", original.cfg)).toBe(true);
+
+      const staleDrain = original.state.cron.stopAndDrain?.();
+      if (!staleDrain) {
+        throw new Error("expected cron stop-and-drain");
+      }
+
+      replacement = await startGatewayCron("replacement");
+      expect(sessionHasAutomation("agent:main:main", replacement.cfg)).toBe(true);
+
+      const failedDrain = expect(staleDrain).rejects.toThrow("stream drain failed");
+      pendingDrain.reject(new Error("stream drain failed"));
+      await failedDrain;
+
+      expect(sessionHasAutomation("agent:main:main", replacement.cfg)).toBe(true);
+      await expect(original.state.cron.stopAndDrain?.()).resolves.toBeUndefined();
+      expect(sessionHasAutomation("agent:main:main", replacement.cfg)).toBe(true);
+    } finally {
+      try {
+        if (replacement) {
+          await cleanGatewayCron(replacement);
+        }
+      } finally {
+        await cleanGatewayCron(original);
+      }
+    }
+  });
+});

src/gateway/server-cron.ts

@@ -1360,28 +1360,33 @@ export function buildGatewayCronService(params: {
     unregisterSessionAutomationSource(automationSource);
   };
   cron.stopAndDrain = async () => {
-    stopCron();
-    stopExitWatchers();
-    stopHeartbeatReconcileRetry();
-    const streamWatchersStop = stopStreamWatchers().then(
-      () => ({ ok: true as const }),
-      (error: unknown) => ({ ok: false as const, error }),
-    );
-    const abortedRuns = abortActiveCronTaskRuns("Gateway shutting down.");
-    const [activeRunDrain, streamWatchersResult] = await Promise.all([
-      waitForActiveCronTaskRuns(CRON_ACTIVE_RUN_SHUTDOWN_DRAIN_MS),
-      streamWatchersStop,
-    ]);
-    if (!activeRunDrain.drained) {
-      cronLogger.warn(
-        { abortedRuns, activeRuns: activeRunDrain.active },
-        "cron: active runs did not drain before shutdown timeout",
+    try {
+      stopCron();
+      stopExitWatchers();
+      stopHeartbeatReconcileRetry();
+      const streamWatchersStop = stopStreamWatchers().then(
+        () => ({ ok: true as const }),
+        (error: unknown) => ({ ok: false as const, error }),
       );
+      const abortedRuns = abortActiveCronTaskRuns("Gateway shutting down.");
+      const [activeRunDrain, streamWatchersResult] = await Promise.all([
+        waitForActiveCronTaskRuns(CRON_ACTIVE_RUN_SHUTDOWN_DRAIN_MS),
+        streamWatchersStop,
+      ]);
+      if (!activeRunDrain.drained) {
+        cronLogger.warn(
+          { abortedRuns, activeRuns: activeRunDrain.active },
+          "cron: active runs did not drain before shutdown timeout",
+        );
+      }
+      if (!streamWatchersResult.ok) {
+        throw streamWatchersResult.error;
+      }
+    } finally {
+      // A failed drain still stops this source; owner comparison protects a
+      // replacement that registered while the old watchers were settling.
+      unregisterSessionAutomationSource(automationSource);
     }
-    if (!streamWatchersResult.ok) {
-      throw streamWatchersResult.error;
-    }
-    unregisterSessionAutomationSource(automationSource);
   };
   // Reconciliations serialize on one tail and only the latest requested epoch
   // executes, so an older reload's convergence can never clobber a newer one.