diff --git a/src/gateway/server-cron.drain.test.ts b/src/gateway/server-cron.drain.test.ts new file mode 100644 index 000000000000..eb43eac08e5e --- /dev/null +++ b/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>(), +})); + +vi.mock("../config/io.js", async (importOriginal) => ({ + ...(await importOriginal()), + getRuntimeConfig: getRuntimeConfigMock, +})); + +vi.mock("./cron-stream-watchers.js", async (importOriginal) => ({ + ...(await importOriginal()), + 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; + cfg: OpenClawConfig; + stateDir: string; +}; + +async function startGatewayCron(label: string): Promise { + 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 { + 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); + } + } + }); +}); diff --git a/src/gateway/server-cron.ts b/src/gateway/server-cron.ts index 3cc3dffbb687..3e5528ce7e02 100644 --- a/src/gateway/server-cron.ts +++ b/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.