Allow list of sleep steps in disableSleeps()

Author: pombosilva

packages/workflows-shared/src/context.ts

@@ -464,7 +464,13 @@ export class Context extends RpcTarget {
 		const sleepLogWrittenKey = `${cacheKey}-log-written`;
 		const maybeResult = await this.#state.storage.get(sleepKey);
 
-		const disableSleeps = await this.#state.storage.get("disableSleeps");
+		const sleepNameCountHash = await computeHash(
+			name + this.#getCount("sleep-" + name)
+		);
+		const disableThisSleep = await this.#state.storage.get(sleepNameCountHash);
+		const disableAllSleeps = await this.#state.storage.get("disableAllSleeps");
+
+		const disableSleep = disableAllSleeps || disableThisSleep;
 
 		if (maybeResult != undefined) {
 			// @ts-expect-error priorityQueue is initiated in init
@@ -474,7 +480,7 @@ export class Context extends RpcTarget {
 			// in case the engine dies while sleeping and wakes up before the retry period
 			if (entryPQ !== undefined) {
 				await scheduler.wait(
-					disableSleeps ? 0 : entryPQ.targetTimestamp - Date.now()
+					disableSleep ? 0 : entryPQ.targetTimestamp - Date.now()
 				);
 				// @ts-expect-error priorityQueue is initiated in init
 				this.#engine.priorityQueue.remove({ hash: cacheKey, type: "sleep" });
@@ -513,13 +519,12 @@ export class Context extends RpcTarget {
 		// @ts-expect-error priorityQueue is initiated in init
 		await this.#engine.priorityQueue.add({
 			hash: cacheKey,
-			targetTimestamp: Date.now() + (disableSleeps ? 0 : duration),
+			targetTimestamp: Date.now() + (disableSleep ? 0 : duration),
 			type: "sleep",
 		});
 
-		console.log("disable sleeps:", disableSleeps);
 		// this probably will never finish except if sleep is less than the grace period
-		await scheduler.wait(disableSleeps ? 0 : duration);
+		await scheduler.wait(disableSleep ? 0 : duration);
 
 		this.#engine.writeLog(
 			InstanceEvent.SLEEP_COMPLETE,

packages/workflows-shared/src/modifier.ts

@@ -48,15 +48,42 @@ export class InstanceModifier extends RpcTarget {
 		return valueKey;
 	}
 
-	async disableSleeps(): Promise<void> {
-		await this.#state.storage.put("disableSleeps", true);
+	async #getSleepStepDisableKey(step: StepSelector): Promise<string> {
+		let count = 1;
+		if (step.index) {
+			count = step.index;
+		}
+		const sleepNameCountHash = await computeHash(step.name + count);
+
+		return sleepNameCountHash;
+	}
+
+	async disableSleeps(steps?: StepSelector[]): Promise<void> {
+		if (!steps) {
+			await this.#state.storage.put("disableAllSleeps", true);
+		} else {
+			for (const step of steps) {
+				const sleepDisableKey = await this.#getSleepStepDisableKey(step);
+				await this.#state.storage.put(sleepDisableKey, true);
+			}
+		}
 	}
 
 	async mockStepResult(step: StepSelector, stepResult: unknown): Promise<void> {
+		console.log("mocking step result");
 		const valueKey = await this.#getStepCacheKey(step);
 		await this.#state.storage.put(`replace-result-${valueKey}`, stepResult);
 	}
 
+	async mockStepImplementation(
+		_step: StepSelector,
+		_implementation: () => Promise<unknown>
+	): Promise<void> {
+		// TODO
+		// can call the new implementation here and replace the step result - meh
+		// ideally pass the implementation to context so it can race with timeout
+	}
+
 	async mockEvent(event: UserEvent): Promise<void> {
 		// could maybe:
 		// WorkflowInstance.sendEvent()