Fix flaky tests

Author: pombosilva

packages/workflows-shared/package.json

@@ -39,13 +39,13 @@
 		"zod": "^3.22.3"
 	},
 	"devDependencies": {
-		"@cloudflare/vitest-pool-workers": "catalog:vitest-3",
+		"@cloudflare/vitest-pool-workers": "^0.13.1",
 		"@cloudflare/workers-tsconfig": "workspace:*",
 		"@cloudflare/workers-types": "catalog:default",
 		"@types/mime": "^3.0.4",
 		"esbuild": "catalog:default",
 		"typescript": "catalog:default",
-		"vitest": "catalog:vitest-3"
+		"vitest": "catalog:default"
 	},
 	"engines": {
 		"node": ">=18.0.0"

packages/workflows-shared/src/binding.ts

@@ -1,7 +1,6 @@
 import { RpcTarget, WorkerEntrypoint } from "cloudflare:workers";
 import { InstanceEvent, instanceStatusName } from "./instance";
 import {
-	isAbortError,
 	isUserTriggeredPause,
 	isUserTriggeredRestart,
 	isUserTriggeredTerminate,
@@ -58,11 +57,9 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> {
 					val[Symbol.dispose]();
 				}
 			})
-			.catch((e) => {
-				// Suppress abort errors since they're expected
-				if (!isAbortError(e)) {
-					throw e;
-				}
+			.catch(() => {
+				// Suppress all rejections: create() should queue and
+				// return immediately
 			});
 
 		this.ctx.waitUntil(initPromise);

packages/workflows-shared/src/context.ts

@@ -72,6 +72,11 @@ export class Context extends RpcTarget {
 		}
 
 		const status = await this.#engine.getStatus();
+
+		if (status === InstanceStatus.Paused) {
+			throw new Error(ABORT_REASONS.USER_PAUSE);
+		}
+
 		if (status === InstanceStatus.WaitingForPause) {
 			await this.#state.storage.put(PAUSE_DATETIME, new Date());
 			const metadata =
@@ -83,7 +88,7 @@ export class Context extends RpcTarget {
 					InstanceStatus.Paused
 				);
 			}
-			await this.#engine.abort(ABORT_REASONS.USER_PAUSE);
+			throw new Error(ABORT_REASONS.USER_PAUSE);
 		}
 	}
 
@@ -511,8 +516,34 @@ export class Context extends RpcTarget {
 						type: "retry",
 					});
 					await this.#engine.timeoutHandler.release(this.#engine);
-					// this may never finish because of the grace period - but waker will take of it
-					await scheduler.wait(durationMs);
+					// Race retry wait against the pause signal so pause
+					// takes effect immediately during retries
+					{
+						const retryPauseSignal = this.#engine.pauseController.signal;
+						let pausedDuringRetry = false;
+						await Promise.race([
+							scheduler.wait(durationMs),
+							new Promise<void>((resolve) => {
+								if (retryPauseSignal.aborted) {
+									resolve();
+									return;
+								}
+								retryPauseSignal.addEventListener("abort", () => resolve(), {
+									once: true,
+								});
+							}),
+						]);
+						const retryStatus = await this.#engine.getStatus();
+						if (
+							retryStatus === InstanceStatus.Paused ||
+							retryStatus === InstanceStatus.WaitingForPause
+						) {
+							pausedDuringRetry = true;
+						}
+						if (pausedDuringRetry) {
+							throw new Error(ABORT_REASONS.USER_PAUSE);
+						}
+					}
 
 					// if it ever reaches here, we can try to remove it from the priority queue since it's no longer useful
 					// @ts-expect-error priorityQueue is initiated in init
@@ -634,8 +665,37 @@ export class Context extends RpcTarget {
 			type: "sleep",
 		});
 
-		// this probably will never finish except if sleep is less than the grace period
-		await scheduler.wait(disableSleep ? 0 : duration);
+		// Race the sleep against the pause signal
+		const pauseSignal = this.#engine.pauseController.signal;
+		const sleepDuration = disableSleep ? 0 : duration;
+
+		let pausedDuringSleep = false;
+		await Promise.race([
+			scheduler.wait(sleepDuration),
+			new Promise<void>((resolve) => {
+				if (pauseSignal.aborted) {
+					resolve();
+					return;
+				}
+				pauseSignal.addEventListener("abort", () => resolve(), {
+					once: true,
+				});
+			}),
+		]);
+
+		// Check if we were paused during the sleep
+		const statusAfterSleep = await this.#engine.getStatus();
+		if (
+			statusAfterSleep === InstanceStatus.Paused ||
+			statusAfterSleep === InstanceStatus.WaitingForPause
+		) {
+			pausedDuringSleep = true;
+		}
+
+		if (pausedDuringSleep) {
+			// Throw pause error
+			throw new Error(ABORT_REASONS.USER_PAUSE);
+		}
 
 		this.#engine.writeLog(
 			InstanceEvent.SLEEP_COMPLETE,
@@ -806,33 +866,49 @@ export class Context extends RpcTarget {
 			this.#engine.waiters.set(options.type, callbacks);
 		});
 
-		const result = await Promise.race([
+		// Race event, timeout, and pause signal. The pause promise resolves
+		// when the race settles via event/timeout before the pause signal fires
+		const pauseSignal = this.#engine.pauseController.signal;
+		const pausePromise = new Promise<void>((resolve) => {
+			if (pauseSignal.aborted) {
+				resolve();
+				return;
+			}
+			pauseSignal.addEventListener("abort", () => resolve(), {
+				once: true,
+			});
+		});
+
+		const raceResult = await Promise.race([
 			eventPromise,
 			timeoutEntryPQ !== undefined
 				? timeoutPromise(timeoutEntryPQ.targetTimestamp - Date.now(), false)
 				: timeoutPromise(ms(options.timeout), true),
-		])
-			.then(async (event) => {
-				this.#engine.writeLog(
-					InstanceEvent.WAIT_COMPLETE,
-					cacheKey,
-					waitForEventNameWithCounter,
-					event as Event
-				);
-				await this.#state.storage.put(waitForEventKey, event);
-				return event;
-			})
-			.catch(async (error) => {
-				this.#engine.writeLog(
-					InstanceEvent.WAIT_TIMED_OUT,
-					cacheKey,
-					waitForEventNameWithCounter,
-					error
-				);
-				await this.#state.storage.put(errorKey, error);
-				throw error;
-			});
+			pausePromise,
+		]).catch(async (error) => {
+			this.#engine.writeLog(
+				InstanceEvent.WAIT_TIMED_OUT,
+				cacheKey,
+				waitForEventNameWithCounter,
+				error
+			);
+			await this.#state.storage.put(errorKey, error);
+			throw error;
+		});
+
+		// Pause signal won the race — throw to stop the workflow
+		if (raceResult === undefined) {
+			throw new Error(ABORT_REASONS.USER_PAUSE);
+		}
+
+		this.#engine.writeLog(
+			InstanceEvent.WAIT_COMPLETE,
+			cacheKey,
+			waitForEventNameWithCounter,
+			raceResult as Event
+		);
+		await this.#state.storage.put(waitForEventKey, raceResult);
 
-		return result as WorkflowStepEvent<T>;
+		return raceResult as WorkflowStepEvent<T>;
 	}
 }

packages/workflows-shared/src/engine.ts

@@ -13,7 +13,6 @@ import {
 	ABORT_REASONS,
 	createWorkflowError,
 	isAbortError,
-	isUserTriggeredPause,
 	WorkflowFatalError,
 } from "./lib/errors";
 import {
@@ -99,6 +98,8 @@ export class Engine extends DurableObject<Env> {
 	timeoutHandler: GracePeriodSemaphore;
 	priorityQueue: TimePriorityQueue | undefined;
 	stepLimit: number;
+	engineAbortController: AbortController = new AbortController();
+	pauseController: AbortController = new AbortController();
 
 	waiters: Map<string, Array<(event: Event | PromiseLike<Event>) => void>> =
 		new Map();
@@ -425,6 +426,11 @@ export class Engine extends DurableObject<Env> {
 
 	async abort(reason: string) {
 		await this.ctx.storage.sync();
+
+		// Clean up pending JS operations
+		this.timeoutHandler.dispose();
+		this.engineAbortController.abort(new Error(reason));
+
 		this.ctx.abort(reason);
 	}
 
@@ -655,14 +661,13 @@ export class Engine extends DurableObject<Env> {
 					metadata.instance.id,
 					InstanceStatus.Paused
 				);
-				await this.abort(ABORT_REASONS.USER_PAUSE);
+				// Signal the pause controller to interrupt any active
+				// scheduler.wait (sleep/waitForEvent). The workflow will
+				// throw a pause error at the next step boundary via
+				// #checkForPendingPause()
+				this.pauseController.abort(ABORT_REASONS.USER_PAUSE);
 			})
-			.catch((e) => {
-				// Expected: abort rejects the promise chain when it kills the DO
-				if (!isAbortError(e)) {
-					throw e;
-				}
-			});
+			.catch(() => {});
 	}
 
 	async userTriggeredRestart() {
@@ -765,15 +770,24 @@ export class Engine extends DurableObject<Env> {
 		const pausedDate = await this.ctx.storage.get<Date>(PAUSE_DATETIME);
 		if (pausedDate !== undefined) {
 			const offset = Date.now() - new Date(pausedDate).valueOf();
-			const pq = new TimePriorityQueue(this.ctx, metadata);
-			pq.offsetAll(offset);
+			if (this.priorityQueue) {
+				this.priorityQueue.offsetAll(offset);
+			} else {
+				const pq = new TimePriorityQueue(this.ctx, metadata);
+				pq.offsetAll(offset);
+			}
 		}
 		await this.ctx.storage.delete(PAUSE_DATETIME);
 
 		const { accountId, workflow, version, instance, event } = metadata;
 
 		await this.ctx.storage.put(ENGINE_STATUS_KEY, InstanceStatus.Queued);
 
+		// Reset pause state for the new run. The DO stays alive across
+		// pause/resume, so we need to clear stale in-memory state from the previous run.
+		this.pauseController = new AbortController();
+		this.waiters = new Map();
+
 		void this.init(accountId, workflow, version, instance, event);
 	}
 
@@ -879,7 +893,7 @@ export class Engine extends DurableObject<Env> {
 			});
 			this.isRunning = false;
 		} catch (err) {
-			if (isUserTriggeredPause(err)) {
+			if (isAbortError(err)) {
 				this.isRunning = false;
 				return;
 			}

packages/workflows-shared/src/lib/gracePeriodSemaphore.ts

@@ -1,5 +1,4 @@
 import { ms } from "itty-time";
-import { ABORT_REASONS } from "./errors";
 import type { Engine } from "../engine";
 import type { WorkflowSleepDuration } from "cloudflare:workers";
 
@@ -121,6 +120,22 @@ export class GracePeriodSemaphore {
 		this.#waitingSteps = [];
 	}
 
+	dispose() {
+		// Reject all waiting step promises so they stop blocking
+		for (const promise of this.#waitingSteps) {
+			promise.rejectCallback();
+		}
+		this.#waitingSteps = [];
+
+		// Reject all waiting promises
+		for (const promise of this.#waitingPromises) {
+			promise.rejectCallback();
+		}
+		this.#waitingPromises = [];
+
+		this.#canInitiateSteps = false;
+	}
+
 	isRunningStep() {
 		return this.#counter > 0;
 	}
@@ -166,7 +181,10 @@ export const startGracePeriod: GracePeriodCallback = async (
 
 		// Ensure next alarm is set before we abort
 		await engine.priorityQueue?.handleNextAlarm();
-		await engine.abort(ABORT_REASONS.GRACE_PERIOD_COMPLETE);
+		// await engine.abort(ABORT_REASONS.GRACE_PERIOD_COMPLETE);
 	};
-	void gracePeriodHandler();
+	void gracePeriodHandler().catch(() => {
+		// Swallow — the engine is shutting down (abort kills the context,
+		// which rejects the scheduler.wait inside gracePeriodHandler)
+	});
 };