@rivetkit/workflow-engine 2.3.0-rc.9 → 2.3.1

package/src/index.ts

@@ -13,6 +13,7 @@ export {
 } from "./context.js";
 // Driver
 export type { EngineDriver, KVEntry, KVWrite } from "./driver.js";
+export { extractErrorInfo } from "./error-utils.js";
 // Errors
 export {
 	CancelledError,
@@ -29,7 +30,6 @@ export {
 	StepExhaustedError,
 	StepFailedError,
 } from "./errors.js";
-export { extractErrorInfo } from "./error-utils.js";
 
 // Location utilities
 export {
@@ -82,9 +82,6 @@ export type {
 	PathSegment,
 	RaceEntry,
 	RemovedEntry,
-	WorkflowEntryMetadataSnapshot,
-	WorkflowHistoryEntry,
-	WorkflowHistorySnapshot,
 	RollbackCheckpointEntry,
 	RollbackContextInterface,
 	RunWorkflowOptions,
@@ -102,20 +99,23 @@ export type {
 	TryStepFailure,
 	TryStepResult,
 	WorkflowContextInterface,
+	WorkflowEntryMetadataSnapshot,
 	WorkflowError,
 	WorkflowErrorEvent,
 	WorkflowErrorHandler,
 	WorkflowFunction,
 	WorkflowHandle,
-	WorkflowRollbackErrorEvent,
+	WorkflowHistoryEntry,
+	WorkflowHistorySnapshot,
+	WorkflowMessageDriver,
 	WorkflowQueue,
 	WorkflowQueueMessage,
 	WorkflowQueueNextBatchOptions,
 	WorkflowQueueNextOptions,
-	WorkflowMessageDriver,
 	WorkflowResult,
-	WorkflowRunMode,
+	WorkflowRollbackErrorEvent,
 	WorkflowRunErrorEvent,
+	WorkflowRunMode,
 	WorkflowState,
 	WorkflowStepErrorEvent,
 } from "./types.js";
@@ -185,9 +185,9 @@ import type {
 	RunWorkflowOptions,
 	Storage,
 	WorkflowErrorEvent,
-	WorkflowHistorySnapshot,
 	WorkflowFunction,
 	WorkflowHandle,
+	WorkflowHistorySnapshot,
 	WorkflowMessageDriver,
 	WorkflowResult,
 	WorkflowRunMode,
@@ -532,26 +532,46 @@ async function executeLiveWorkflow<TInput, TOutput>(
 		const hasDeadline = result.sleepUntil !== undefined;
 
 		if (hasMessages && hasDeadline) {
-			// Wait for EITHER a message OR the deadline (for queue.next timeout)
+			// Wait for EITHER a message OR the deadline (for queue.next timeout).
+			// Scope both waiters to a per-iteration controller chained to the run
+			// signal so that whichever loses the race is cancelled immediately.
+			// Otherwise the loser (a message waiter or an armed sleep timer) stays
+			// registered on the long-lived run signal and leaks once per cycle.
+			const iterationAbort = new AbortController();
+			const onRunAbort = () => iterationAbort.abort();
+			if (abortController.signal.aborted) {
+				iterationAbort.abort();
+			} else {
+				abortController.signal.addEventListener("abort", onRunAbort, {
+					once: true,
+				});
+			}
+			const messagePromise = awaitWithEviction(
+				driver.waitForMessages(
+					result.waitingForMessages!,
+					iterationAbort.signal,
+				),
+				iterationAbort.signal,
+			);
+			const sleepPromise = waitForSleep(
+				runtime,
+				result.sleepUntil!,
+				iterationAbort.signal,
+			);
+			// Swallow the loser's rejection once it is aborted below so it does
+			// not surface as an unhandled rejection.
+			messagePromise.catch(() => {});
+			sleepPromise.catch(() => {});
 			try {
-				const messagePromise = awaitWithEviction(
-					driver.waitForMessages(
-						result.waitingForMessages!,
-						abortController.signal,
-					),
-					abortController.signal,
-				);
-				const sleepPromise = waitForSleep(
-					runtime,
-					result.sleepUntil!,
-					abortController.signal,
-				);
 				await Promise.race([messagePromise, sleepPromise]);
 			} catch (error) {
 				if (error instanceof EvictedError) {
 					return lastResult;
 				}
 				throw error;
+			} finally {
+				iterationAbort.abort();
+				abortController.signal.removeEventListener("abort", onRunAbort);
 			}
 			continue;
 		}

package/src/location.ts

@@ -159,7 +159,7 @@ export function getChildEntries(
 		const isChild =
 			parentKey === ""
 				? true
-				: key.startsWith(parentKey + "/") || key === parentKey;
+				: key.startsWith(`${parentKey}/`) || key === parentKey;
 
 		if (isChild) {
 			// Return the actual entry's location, not the parent location

package/src/types.ts

@@ -399,6 +399,8 @@ export interface StepConfig<T> {
 	retryBackoffMax?: number;
 	/** Timeout in ms for step execution (default: 30000). Set to 0 to disable. */
 	timeout?: number;
+	/** If true, step timeouts retry like any other error instead of failing immediately as critical. Default: false. */
+	retryOnTimeout?: boolean;
 }
 
 export type TryStepCatchKind =
@@ -455,7 +457,7 @@ export type LoopResult<S, T> =
  * `Loop.continue(undefined)`.
  */
 export type LoopIterationResult<S, T> = Promise<
-	LoopResult<S, T> | (S extends undefined ? void : never)
+	LoopResult<S, T> | (S extends undefined ? undefined : never)
 >;
 
 /**