fix(api): preserve terminal scheduled automation dispatch

apps/api/src/app/api/automations/dispatch/[id]/route.ts

@@ -1,12 +1,16 @@
 import { dbWs } from "@superset/db/client";
-import { automations } from "@superset/db/schema";
+import { automations, type SelectAutomation } from "@superset/db/schema";
 import { dispatchAutomation } from "@superset/trpc/automation-dispatch";
 import { Receiver } from "@upstash/qstash";
-import { eq } from "drizzle-orm";
+import { and, eq, sql } from "drizzle-orm";
 import { z } from "zod";
 
 import { env } from "@/env";
 import { getRelayUrl } from "@/lib/relay-url";
+import {
+	matchesTerminalOccurrence,
+	matchesTerminalReservation,
+} from "../../terminal-occurrence";
 
 export const maxDuration = 60;
 export const dynamic = "force-dynamic";
@@ -19,6 +23,13 @@ const receiver = new Receiver({
 const payloadSchema = z.object({
 	automationId: z.string().uuid(),
 	scheduledFor: z.string().datetime(),
+	// The token identifies the evaluator's terminal reservation. The previous
+	// value lets dispatch claim an enabled row only if no edit won the race.
+	terminal: z.boolean().default(false),
+	terminalDispatchToken: z.string().datetime().optional(),
+	terminalPreviousUpdatedAt: z.string().datetime().optional(),
+	// Accept messages created before the updatedAt reservation was introduced.
+	terminalPendingNextRunAt: z.string().datetime().optional(),
 });
 
 export async function POST(
@@ -56,17 +67,140 @@ export async function POST(
 	if (!automation) {
 		return Response.json({ ok: true, skipped: "deleted" });
 	}
-	if (!automation.enabled) {
+
+	const scheduledFor = new Date(parsed.data.scheduledFor);
+	if (parsed.data.terminal) {
+		const terminalReservation = matchesTerminalReservation({
+			updatedAt: automation.updatedAt,
+			terminalDispatchToken: parsed.data.terminalDispatchToken,
+		});
+		if (!automation.enabled && !terminalReservation) {
+			return Response.json({ ok: true, skipped: "disabled" });
+		}
+		if (
+			!matchesTerminalOccurrence({
+				nextRunAt: automation.nextRunAt,
+				scheduledFor,
+				legacyPendingNextRunAt: parsed.data.terminalPendingNextRunAt,
+			})
+		) {
+			return Response.json({ ok: true, skipped: "stale" });
+		}
+
+		const claimed = await claimTerminalOccurrence({
+			automation,
+			automationId: parsed.data.automationId,
+			terminalDispatchToken: parsed.data.terminalDispatchToken,
+			terminalPreviousUpdatedAt: parsed.data.terminalPreviousUpdatedAt,
+		});
+		if (!claimed) {
+			return Response.json({
+				ok: true,
+				skipped: automation.enabled ? "stale" : "disabled",
+			});
+		}
+	} else if (!automation.enabled) {
 		return Response.json({ ok: true, skipped: "disabled" });
 	}
 
 	const outcome = await dispatchAutomation({
 		automation,
-		scheduledFor: new Date(parsed.data.scheduledFor),
+		scheduledFor,
 		// The owner's host may be on an overridden relay (relay-url-override);
 		// env.RELAY_URL alone reaches only hosts still on the default relay.
 		relayUrl: await getRelayUrl(automation.ownerUserId),
 	});
 
+	if (parsed.data.terminal && outcome.status === "conflict") {
+		// A conflict can mean that another worker owns a still-dispatching run.
+		// Do not acknowledge the message or close the recurrence: QStash retries,
+		// then its failure callback records dispatch_failed if the owner never
+		// reaches a terminal run state.
+		return Response.json(
+			{
+				ok: false,
+				error: "Terminal automation dispatch is already in progress",
+			},
+			{ status: 409 },
+		);
+	}
+
 	return Response.json({ ok: true, outcome });
 }
+
+async function claimTerminalOccurrence({
+	automation,
+	automationId,
+	terminalDispatchToken,
+	terminalPreviousUpdatedAt,
+}: {
+	automation: SelectAutomation;
+	automationId: string;
+	terminalDispatchToken?: string;
+	terminalPreviousUpdatedAt?: string;
+}): Promise<boolean> {
+	if (terminalDispatchToken !== undefined) {
+		if (terminalPreviousUpdatedAt === undefined) {
+			return false;
+		}
+
+		const token = new Date(terminalDispatchToken);
+		if (automation.enabled) {
+			const [claimed] = await dbWs
+				.update(automations)
+				.set({ enabled: false, updatedAt: token })
+				.where(
+					and(
+						eq(automations.id, automationId),
+						eq(automations.enabled, true),
+						eq(automations.nextRunAt, automation.nextRunAt),
+						// Compare the original database value without converting it to a
+						// millisecond-only JavaScript Date.
+						sql`${automations.updatedAt} = ${terminalPreviousUpdatedAt}::timestamptz`,
+					),
+				)
+				.returning({ id: automations.id });
+
+			if (claimed !== undefined) return true;
+		}
+
+		// Evaluation may have completed the reservation after the SELECT above
+		// but before this claim. Claim that exact reservation instead of
+		// acknowledging the queued occurrence without a run.
+		const [reserved] = await dbWs
+			.update(automations)
+			.set({ updatedAt: token })
+			.where(
+				and(
+					eq(automations.id, automationId),
+					eq(automations.enabled, false),
+					eq(automations.nextRunAt, automation.nextRunAt),
+					sql`${automations.updatedAt} = ${terminalDispatchToken}::timestamptz`,
+				),
+			)
+			.returning({ id: automations.id });
+
+		return reserved !== undefined;
+	}
+
+	// Rolling-deployment compatibility for messages that carry only the old
+	// nextRunAt reservation. They may claim an enabled row, but never bypass a
+	// user-disabled row.
+	if (!automation.enabled) {
+		return false;
+	}
+
+	const [claimed] = await dbWs
+		.update(automations)
+		.set({ enabled: false })
+		.where(
+			and(
+				eq(automations.id, automationId),
+				eq(automations.enabled, true),
+				eq(automations.nextRunAt, automation.nextRunAt),
+			),
+		)
+		.returning({ id: automations.id });
+
+	return claimed !== undefined;
+}

apps/api/src/app/api/automations/evaluate/route.test.ts

@@ -0,0 +1,198 @@
+import { describe, expect, mock, test } from "bun:test";
+import { nextOccurrenceAfter } from "@superset/shared/rrule";
+
+const terminalOccurrence = new Date("2026-08-06T18:46:30.000Z");
+const nonTerminalOccurrence = new Date("2026-08-05T18:46:30.000Z");
+const terminalUpdatedAt = new Date("2026-08-01T00:00:00.123Z");
+const terminalUpdatedAtToken = "2026-08-01T00:00:00.123456Z";
+const nonTerminalUpdatedAt = new Date("2026-08-01T00:00:00.000Z");
+const nonTerminalUpdatedAtToken = "2026-08-01T00:00:00.000000Z";
+const automationId = "75c82d06-77af-454c-9f0c-e6c617ea702b";
+const terminalDispatchToken = new Date(terminalOccurrence.getTime() + 1);
+
+let dueAutomations: Array<{
+	id: string;
+	nextRunAt: Date;
+	rrule: string;
+	dtstart: Date;
+	timezone: string;
+	updatedAt: Date;
+	updatedAtToken: string;
+}> = [];
+const updateValues: unknown[] = [];
+const updateWhereValues: unknown[] = [];
+const batchJSON = mock(async (_messages: unknown[]) => undefined);
+
+mock.module("@/env", () => ({
+	env: {
+		NEXT_PUBLIC_API_URL: "http://localhost:3001",
+		QSTASH_TOKEN: "test-token",
+		QSTASH_URL: "https://qstash.example.com",
+		QSTASH_CURRENT_SIGNING_KEY: "current-key",
+		QSTASH_NEXT_SIGNING_KEY: "next-key",
+	},
+}));
+
+mock.module("@upstash/qstash", () => ({
+	Client: class {
+		batchJSON = batchJSON;
+	},
+	Receiver: class {
+		verify = mock(async () => true);
+	},
+}));
+
+mock.module("@superset/db/schema", () => ({
+	automations: {
+		id: "id",
+		enabled: "enabled",
+		nextRunAt: "nextRunAt",
+		updatedAt: "updatedAt",
+	},
+	automationRuns: {
+		automationId: "automationId",
+		organizationId: "organizationId",
+		scheduledFor: "scheduledFor",
+	},
+}));
+
+mock.module("@superset/db/client", () => ({
+	dbWs: {
+		select: () => ({
+			from: () => ({
+				where: () => ({
+					orderBy: () => ({
+						limit: async () => dueAutomations,
+					}),
+				}),
+			}),
+		}),
+		update: () => ({
+			set: (values: unknown) => {
+				updateValues.push(values);
+				return {
+					where: async (condition: unknown) => {
+						updateWhereValues.push(condition);
+						return undefined;
+					},
+				};
+			},
+		}),
+	},
+}));
+
+mock.module("drizzle-orm", () => ({
+	and: (...conditions: unknown[]) => ({ type: "and", conditions }),
+	eq: (field: unknown, value: unknown) => ({ type: "eq", field, value }),
+	getTableColumns: () => ({}),
+	lte: (field: unknown, value: unknown) => ({ type: "lte", field, value }),
+	sql: (strings: TemplateStringsArray, ...values: unknown[]) => ({
+		type: "sql",
+		strings: [...strings],
+		values,
+	}),
+}));
+
+const { POST } = await import("./route");
+
+type EnqueuedMessage = {
+	body: {
+		automationId: string;
+		scheduledFor: string;
+		terminal: boolean;
+	};
+};
+
+function request(): Request {
+	return new Request("http://localhost:3001/api/automations/evaluate", {
+		method: "POST",
+		headers: { "upstash-signature": "valid-signature" },
+		body: "{}",
+	});
+}
+
+describe("automations evaluate route", () => {
+	test("marks a COUNT=1 occurrence as terminal in the queued payload", async () => {
+		dueAutomations = [
+			{
+				id: automationId,
+				nextRunAt: terminalOccurrence,
+				rrule: "FREQ=DAILY;COUNT=1",
+				dtstart: terminalOccurrence,
+				timezone: "UTC",
+				updatedAt: terminalUpdatedAt,
+				updatedAtToken: terminalUpdatedAtToken,
+			},
+		];
+		updateValues.length = 0;
+		updateWhereValues.length = 0;
+		batchJSON.mockClear();
+
+		const response = await POST(request());
+
+		expect(response.status).toBe(200);
+		expect(batchJSON).toHaveBeenCalledTimes(1);
+		const messages = batchJSON.mock.calls[0]?.[0] as EnqueuedMessage[];
+		expect(messages[0]?.body).toEqual({
+			automationId,
+			scheduledFor: new Date(
+				Math.floor(terminalOccurrence.getTime() / 60_000) * 60_000,
+			).toISOString(),
+			terminal: true,
+			terminalDispatchToken: terminalDispatchToken.toISOString(),
+			terminalPreviousUpdatedAt: terminalUpdatedAtToken,
+		});
+		expect(updateValues).toEqual([
+			{
+				enabled: false,
+				updatedAt: terminalDispatchToken,
+			},
+		]);
+		expect(updateWhereValues[0]).toEqual({
+			type: "and",
+			conditions: [
+				{ type: "eq", field: "id", value: automationId },
+				{ type: "eq", field: "enabled", value: true },
+				{ type: "eq", field: "nextRunAt", value: terminalOccurrence },
+				{
+					type: "sql",
+					strings: ["", " = ", "::timestamptz"],
+					values: ["updatedAt", terminalUpdatedAtToken],
+				},
+			],
+		});
+	});
+
+	test("keeps the existing non-terminal advance path", async () => {
+		dueAutomations = [
+			{
+				id: automationId,
+				nextRunAt: nonTerminalOccurrence,
+				rrule: "FREQ=DAILY",
+				dtstart: nonTerminalOccurrence,
+				timezone: "UTC",
+				updatedAt: nonTerminalUpdatedAt,
+				updatedAtToken: nonTerminalUpdatedAtToken,
+			},
+		];
+		updateValues.length = 0;
+		updateWhereValues.length = 0;
+		batchJSON.mockClear();
+
+		const response = await POST(request());
+
+		expect(response.status).toBe(200);
+		const messages = batchJSON.mock.calls[0]?.[0] as EnqueuedMessage[];
+		expect(messages[0]?.body.terminal).toBe(false);
+		expect(updateValues).toEqual([
+			{
+				nextRunAt: nextOccurrenceAfter({
+					rrule: "FREQ=DAILY",
+					dtstart: nonTerminalOccurrence,
+					timezone: "UTC",
+					after: nonTerminalOccurrence,
+				}),
+			},
+		]);
+	});
+});