From dc53f527ec542ddecaa0523f68ca7e6bd884bdf2 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 2 Jan 2026 11:58:18 -0800 Subject: [PATCH] feat(rivetkit): workflows --- pnpm-lock.yaml | 45 + .../packages/workflow-engine/QUICKSTART.md | 470 +++++ .../packages/workflow-engine/TODO.md | 59 + .../packages/workflow-engine/agents.md | 9 + .../packages/workflow-engine/architecture.md | 466 +++++ .../packages/workflow-engine/docs/advanced.md | 27 + .../workflow-engine/docs/cancellation.md | 31 + .../workflow-engine/docs/control-flow.md | 95 + .../docs/high-performance-workflows.md | 46 + .../docs/long-running-workflows.md | 47 + .../docs/migrating-workflows.md | 60 + .../packages/workflow-engine/docs/retries.md | 63 + .../workflow-engine/docs/rollback-behavior.md | 54 + .../packages/workflow-engine/docs/state.md | 45 + ...aiting-for-events-and-human-in-the-loop.md | 66 + .../packages/workflow-engine/package.json | 69 + .../packages/workflow-engine/schemas/serde.ts | 655 ++++++ .../packages/workflow-engine/schemas/v1.bare | 203 ++ .../workflow-engine/schemas/versioned.ts | 131 ++ .../workflow-engine/scripts/compile-bare.ts | 127 ++ .../packages/workflow-engine/src/context.ts | 1821 +++++++++++++++++ .../packages/workflow-engine/src/driver.ts | 88 + .../packages/workflow-engine/src/errors.ts | 162 ++ .../packages/workflow-engine/src/index.ts | 862 ++++++++ .../packages/workflow-engine/src/keys.ts | 311 +++ .../packages/workflow-engine/src/location.ts | 168 ++ .../packages/workflow-engine/src/storage.ts | 436 ++++ .../packages/workflow-engine/src/testing.ts | 134 ++ .../packages/workflow-engine/src/types.ts | 442 ++++ .../packages/workflow-engine/src/utils.ts | 48 + .../workflow-engine/tests/driver-kv.test.ts | 120 ++ .../tests/driver-scheduling.test.ts | 42 + .../tests/eviction-cancel.test.ts | 82 + .../workflow-engine/tests/handle.test.ts | 86 + .../workflow-engine/tests/join.test.ts | 149 ++ .../workflow-engine/tests/loops.test.ts | 174 ++ .../workflow-engine/tests/messages.test.ts | 316 +++ .../workflow-engine/tests/race.test.ts | 155 ++ .../workflow-engine/tests/removals.test.ts | 93 + .../workflow-engine/tests/rollback.test.ts | 135 ++ .../workflow-engine/tests/sleep.test.ts | 206 ++ .../workflow-engine/tests/steps.test.ts | 372 ++++ .../workflow-engine/tests/storage.test.ts | 66 + .../packages/workflow-engine/tsconfig.json | 8 + .../packages/workflow-engine/tsup.config.ts | 7 + .../packages/workflow-engine/vitest.config.ts | 6 + 46 files changed, 9257 insertions(+) create mode 100644 rivetkit-typescript/packages/workflow-engine/QUICKSTART.md create mode 100644 rivetkit-typescript/packages/workflow-engine/TODO.md create mode 100644 rivetkit-typescript/packages/workflow-engine/agents.md create mode 100644 rivetkit-typescript/packages/workflow-engine/architecture.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/advanced.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/cancellation.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/control-flow.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/high-performance-workflows.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/migrating-workflows.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/retries.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/rollback-behavior.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/state.md create mode 100644 rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md create mode 100644 rivetkit-typescript/packages/workflow-engine/package.json create mode 100644 rivetkit-typescript/packages/workflow-engine/schemas/serde.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/schemas/v1.bare create mode 100644 rivetkit-typescript/packages/workflow-engine/schemas/versioned.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/context.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/driver.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/errors.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/index.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/keys.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/location.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/storage.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/testing.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/types.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/src/utils.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/driver-kv.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/driver-scheduling.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/eviction-cancel.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/join.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/loops.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/race.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/rollback.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/sleep.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/steps.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/storage.test.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/tsconfig.json create mode 100644 rivetkit-typescript/packages/workflow-engine/tsup.config.ts create mode 100644 rivetkit-typescript/packages/workflow-engine/vitest.config.ts diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4c5b41a89b..1437b75e4d 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3787,6 +3787,46 @@ importers: specifier: ^8.5.0 version: 8.5.0(@microsoft/api-extractor@7.53.2(@types/node@22.18.1))(jiti@1.21.7)(postcss@8.5.6)(tsx@4.20.6)(typescript@5.9.3)(yaml@2.8.2) + rivetkit-typescript/packages/workflow-engine: + dependencies: + '@rivetkit/bare-ts': + specifier: ^0.6.2 + version: 0.6.2 + cbor-x: + specifier: ^1.6.0 + version: 1.6.0 + fdb-tuple: + specifier: ^1.0.0 + version: 1.0.0 + vbare: + specifier: ^0.0.4 + version: 0.0.4 + devDependencies: + '@bare-ts/tools': + specifier: ^0.13.0 + version: 0.13.0(@bare-ts/lib@0.6.0) + '@biomejs/biome': + specifier: ^2.2.3 + version: 2.3.11 + '@types/node': + specifier: ^22.13.1 + version: 22.19.5 + commander: + specifier: ^12.0.0 + version: 12.1.0 + tsup: + specifier: ^8.4.0 + version: 8.5.1(@microsoft/api-extractor@7.53.2(@types/node@22.19.5))(jiti@1.21.7)(postcss@8.5.6)(tsx@4.20.6)(typescript@5.9.3)(yaml@2.8.2) + tsx: + specifier: ^4.7.0 + version: 4.20.6 + typescript: + specifier: ^5.7.3 + version: 5.9.3 + vitest: + specifier: ^3.1.1 + version: 3.2.4(@types/debug@4.1.12)(@types/node@22.19.5)(less@4.4.1)(lightningcss@1.30.2)(sass@1.93.2)(stylus@0.62.0)(terser@5.44.1) + scripts/release: dependencies: commander: @@ -11535,6 +11575,9 @@ packages: fd-slicer@1.1.0: resolution: {integrity: sha512-cE1qsB/VwyQozZ+q1dGxR8LBYNZeofhEdUNGSMbQD3Gw2lAzX9Zb3uIU6Ebc/Fmyjo9AWWfnn0AUCHqtevs/8g==} + fdb-tuple@1.0.0: + resolution: {integrity: sha512-8jSvKPCYCgTpi9Pt87qlfTk6griyMx4Gk3Xv31Dp72Qp8b6XgIyFsMm8KzPmFJ9iJ8K4pGvRxvOS8D0XGnrkjw==} + fdir@6.5.0: resolution: {integrity: sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==} engines: {node: '>=12.0.0'} @@ -25342,6 +25385,8 @@ snapshots: dependencies: pend: 1.2.0 + fdb-tuple@1.0.0: {} + fdir@6.5.0(picomatch@4.0.3): optionalDependencies: picomatch: 4.0.3 diff --git a/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md b/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md new file mode 100644 index 0000000000..5875d2f870 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/QUICKSTART.md @@ -0,0 +1,470 @@ +# Workflow Engine Quickstart + +A durable execution engine for TypeScript. Write long-running, fault-tolerant workflows as regular async functions that can survive process restarts, crashes, and deployments. + +## Overview + +The workflow engine enables reliable execution through: + +1. **History Tracking** - Every operation is recorded in persistent storage +2. **Replay** - On restart, operations replay from history instead of re-executing +3. **Deterministic Execution** - Same inputs produce the same execution path + +## Installation + +```bash +npm install @rivetkit/workflow-engine +``` + +## Quick Example + +```typescript +import { runWorkflow, Loop, type WorkflowContextInterface } from "@rivetkit/workflow-engine"; + +async function orderWorkflow(ctx: WorkflowContextInterface, orderId: string) { + // Steps are durable - if this crashes after payment, it won't charge twice on restart + const payment = await ctx.step("process-payment", async () => { + return await chargeCard(orderId); + }); + + // Wait for shipping confirmation (external message) + const tracking = await ctx.listen("wait-shipping", "shipment-confirmed"); + + // Send notification + await ctx.step("notify-customer", async () => { + await sendEmail(orderId, tracking); + }); + + return { orderId, payment, tracking }; +} + +// Run the workflow +const handle = runWorkflow("order-123", orderWorkflow, "order-123", driver); + +// Send a message from external system +await handle.message("shipment-confirmed", "TRACK-456"); + +// Wait for completion +const result = await handle.result; +``` + +## Core Concepts + +### The Driver + +The workflow engine requires an `EngineDriver` implementation for persistence and scheduling. Each workflow instance operates on an isolated KV namespace. + +```typescript +interface EngineDriver { + // KV Operations + get(key: Uint8Array): Promise; + set(key: Uint8Array, value: Uint8Array): Promise; + delete(key: Uint8Array): Promise; + deletePrefix(prefix: Uint8Array): Promise; + list(prefix: Uint8Array): Promise; // MUST be sorted + batch(writes: KVWrite[]): Promise; // Should be atomic + + // Scheduling + setAlarm(workflowId: string, wakeAt: number): Promise; + clearAlarm(workflowId: string): Promise; + readonly workerPollInterval: number; +} +``` + +### Running Workflows + +```typescript +import { runWorkflow } from "@rivetkit/workflow-engine"; + +const handle = runWorkflow( + "workflow-id", // Unique ID for this workflow instance + myWorkflow, // Your workflow function + { input: "data" }, // Input passed to workflow + driver // Your EngineDriver implementation +); + +// The handle provides methods to interact with the running workflow +await handle.result; // Wait for completion/yield +await handle.message("name", data); // Send a message +await handle.wake(); // Wake immediately +await handle.recover(); // Reset exhausted retries +handle.evict(); // Request graceful shutdown +await handle.cancel(); // Cancel permanently +await handle.getState(); // Get current state +await handle.getOutput(); // Get output if completed +``` + +## Features + +### Steps + +Steps execute arbitrary async code. Results are persisted and replayed on restart. + +```typescript +// Simple form +const result = await ctx.step("fetch-user", async () => { + return await fetchUser(userId); +}); + +// With configuration +const result = await ctx.step({ + name: "external-api", + maxRetries: 5, // Default: 3 + retryBackoffBase: 200, // Default: 100ms + retryBackoffMax: 60000, // Default: 30000ms + timeout: 10000, // Default: 30000ms (0 to disable) + ephemeral: false, // Default: false - batch writes + run: async () => { + return await callExternalApi(); + }, +}); +``` + +**Retry Behavior:** +- Regular errors trigger automatic retry with exponential backoff +- `CriticalError` and `RollbackError` bypass retry logic for unrecoverable errors +- After exhausting retries, `StepExhaustedError` is thrown + +```typescript +import { CriticalError, RollbackError } from "@rivetkit/workflow-engine"; + +await ctx.step("validate", async () => { + if (!isValid(input)) { + throw new CriticalError("Invalid input - no point retrying"); + } + return processInput(input); +}); + +await ctx.step("halt", async () => { + throw new RollbackError("Stop and roll back"); +}); +``` + +### Rollback + +Rollback handlers run only for steps that completed successfully and registered +`rollback`. They run in reverse completion order, persist across restarts, and +skip handlers that already finished. Rollback is disabled until you set a +checkpoint. + +```typescript +await ctx.rollbackCheckpoint("billing"); + +await ctx.step({ + name: "charge", + run: async () => chargeCard(orderId), + rollback: async (_ctx, receipt) => { + await refundCharge(receipt); + }, +}); +``` + +### Loops + +Loops maintain durable state across iterations with periodic checkpointing and bounded history. + +```typescript +import { Loop } from "@rivetkit/workflow-engine"; + +const total = await ctx.loop({ + name: "process-batches", + state: { cursor: null, count: 0 }, // Initial state + commitInterval: 10, // Checkpoint every 10 iterations + historyEvery: 10, // Trim history every 10 iterations + historyKeep: 10, // Keep last 10 iterations of history + run: async (ctx, state) => { + const batch = await ctx.step("fetch", () => fetchBatch(state.cursor)); + + if (!batch.items.length) { + return Loop.break(state.count); // Exit with final value + } + + await ctx.step("process", () => processBatch(batch.items)); + + return Loop.continue({ // Continue with new state + cursor: batch.nextCursor, + count: state.count + batch.items.length, + }); + }, +}); +``` + +**Simple loops** (no state needed): + +```typescript +const result = await ctx.loop("my-loop", async (ctx) => { + // ... do work + if (done) return Loop.break(finalValue); + return Loop.continue(undefined); +}); +``` + +### Sleep + +Pause workflow execution for a duration or until a specific time. + +```typescript +// Sleep for duration +await ctx.sleep("wait-5-min", 5 * 60 * 1000); + +// Sleep until timestamp +await ctx.sleepUntil("wait-midnight", midnightTimestamp); +``` + +Short sleeps (< `driver.workerPollInterval`) wait in memory. Longer sleeps yield to the scheduler and set an alarm for wake-up. + +### Messages (Listen) + +Wait for external events delivered via `handle.message()`. + +```typescript +// Wait for a single message +const data = await ctx.listen("payment", "payment-completed"); + +// Wait for N messages +const items = await ctx.listenN("batch", "item-added", 10); + +// Wait with timeout (returns null on timeout) +const result = await ctx.listenWithTimeout( + "api-response", + "response-received", + 30000 +); + +// Wait until timestamp (returns null on timeout) +const result = await ctx.listenUntil( + "api-response", + "response-received", + deadline +); + +// Wait for up to N messages with timeout +const items = await ctx.listenNWithTimeout( + "batch", + "item-added", + 10, // max items + 60000 // timeout ms +); + +// Wait for up to N messages until timestamp +const items = await ctx.listenNUntil( + "batch", + "item-added", + 10, // max items + deadline // timestamp +); +``` + +**Message delivery:** Messages are loaded once at workflow start. If a message is sent during execution, the workflow yields and picks it up on the next run. + +### Join (Parallel - Wait All) + +Execute multiple branches in parallel and wait for all to complete. + +```typescript +const results = await ctx.join("fetch-all", { + user: { + run: async (ctx) => { + return await ctx.step("get-user", () => fetchUser(userId)); + } + }, + posts: { + run: async (ctx) => { + return await ctx.step("get-posts", () => fetchPosts(userId)); + } + }, + notifications: { + run: async (ctx) => { + return await ctx.step("get-notifs", () => fetchNotifications(userId)); + } + }, +}); + +// results.user, results.posts, results.notifications are all available +// Type-safe: each branch output type is preserved +``` + +If any branch fails, all errors are collected into a `JoinError`: + +```typescript +import { JoinError } from "@rivetkit/workflow-engine"; + +try { + await ctx.join("risky", { /* branches */ }); +} catch (error) { + if (error instanceof JoinError) { + console.log("Failed branches:", Object.keys(error.errors)); + } +} +``` + +### Race (Parallel - First Wins) + +Execute multiple branches and return when the first completes. + +```typescript +const { winner, value } = await ctx.race("timeout-race", [ + { + name: "work", + run: async (ctx) => { + return await ctx.step("do-work", () => doExpensiveWork()); + } + }, + { + name: "timeout", + run: async (ctx) => { + await ctx.sleep("wait", 30000); + return null; // Timeout value + } + }, +]); + +if (winner === "work") { + console.log("Work completed:", value); +} else { + console.log("Timed out"); +} +``` + +- Other branches are cancelled via `AbortSignal` when a winner is determined +- If all branches fail, throws `RaceError` with all errors + +### Eviction and Cancellation + +**Eviction** - Graceful shutdown (workflow can be resumed elsewhere): + +```typescript +handle.evict(); // Request shutdown + +// In workflow, check eviction status: +if (ctx.isEvicted()) { + // Clean up and return +} + +// Or use the abort signal directly: +await fetch(url, { signal: ctx.abortSignal }); +``` + +**Cancellation** - Permanent stop: + +```typescript +await handle.cancel(); // Sets state to "cancelled", clears alarms +``` + +### Recovery + +Reset exhausted retries after a workflow fails. This clears step retry metadata, +removes the workflow error, and schedules the workflow to run again. + +```typescript +await handle.recover(); // Resets retry attempts for failed steps +``` + +### Workflow Migrations (Removed) + +When removing steps from workflow code, use `ctx.removed()` to maintain history compatibility: + +```typescript +async function myWorkflow(ctx: WorkflowContextInterface) { + // This step was removed in v2 + await ctx.removed("old-validation", "step"); + + // New code continues here + await ctx.step("new-logic", async () => { /* ... */ }); +} +``` + +This creates a placeholder entry that satisfies history validation without executing anything. + +## Configuration Constants + +Default values are exported and can be referenced when overriding: + +```typescript +import { + DEFAULT_MAX_RETRIES, // 3 + DEFAULT_RETRY_BACKOFF_BASE, // 100ms + DEFAULT_RETRY_BACKOFF_MAX, // 30000ms + DEFAULT_LOOP_COMMIT_INTERVAL, // 20 iterations + DEFAULT_STEP_TIMEOUT, // 30000ms +} from "@rivetkit/workflow-engine"; +``` + +## Error Types + +```typescript +import { + // User-facing errors + CriticalError, // Throw to skip retries + RollbackError, // Throw to force rollback + RollbackCheckpointError, // Rollback used without checkpoint + StepExhaustedError, // Step failed after all retries + JoinError, // One or more join branches failed + RaceError, // All race branches failed + HistoryDivergedError, // Workflow code changed incompatibly + + // Internal yield errors (caught by runtime) + SleepError, // Workflow sleeping + MessageWaitError, // Waiting for messages + EvictedError, // Workflow evicted + + // User errors + EntryInProgressError, // Forgot to await a step + CancelledError, // Branch cancelled (race) +} from "@rivetkit/workflow-engine"; +``` + +## Workflow States + +```typescript +type WorkflowState = + | "pending" // Not yet started + | "running" // Currently executing + | "sleeping" // Waiting for deadline or message + | "completed" // Finished successfully + | "failed" // Unrecoverable error + | "rolling_back" // Running rollback handlers + | "cancelled"; // Permanently stopped +``` + +## Best Practices + +1. **Unique step names** - Each step/loop/sleep/listen within a scope must have a unique name + +2. **Deterministic code** - Workflow code outside of steps must be deterministic. Don't use `Math.random()`, `Date.now()`, or read external state outside steps. + +3. **Use steps for side effects** - All I/O and non-deterministic operations should be inside steps + +4. **Use CriticalError or RollbackError for permanent failures** - When an error is unrecoverable, throw one to avoid wasting retries + +5. **Check eviction in long operations** - Use `ctx.isEvicted()` or `ctx.abortSignal` to handle graceful shutdown + +6. **Pass AbortSignal to cancellable operations**: + ```typescript + await ctx.step("fetch", async () => { + return fetch(url, { signal: ctx.abortSignal }); + }); + ``` + +7. **Ephemeral steps for batching** - Use `ephemeral: true` for steps where you want to batch writes: + ```typescript + // These batch their writes + await ctx.step({ name: "a", ephemeral: true, run: async () => { ... }}); + await ctx.step({ name: "b", ephemeral: true, run: async () => { ... }}); + // This flushes all pending writes + await ctx.step({ name: "c", run: async () => { ... }}); + ``` + +8. do not try catch in side of a workflow + +9. do not use native loops + +## Further Reading + +See [architecture.md](./architecture.md) for detailed implementation information including: + +- Storage schema and key encoding +- Location system and NameIndex optimization +- Driver requirements +- Error handling internals +- Loop state management and history forgetting diff --git a/rivetkit-typescript/packages/workflow-engine/TODO.md b/rivetkit-typescript/packages/workflow-engine/TODO.md new file mode 100644 index 0000000000..89255f5a75 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/TODO.md @@ -0,0 +1,59 @@ +## loops as checkpoints + +- todo + +## test migrating workflows + +- run workfklow code a +- run workflow code b + +## remove the internal signal queue + +## observe workflow status + +- add support for monitoring workflow state by emitting events from the workflow +- subscribe to events +- update all tests to listen for specific events + - update tests to use snapshot testing based on these events + +## ephemeral-by-default steps + +- make steps ephemeral by default +- add helper fns for things like fetch, clients, etc that auto-flag a step as required to be durable +- can also opt-in to flag a step as durable +- tests: + - default steps are ephemeral without config + - opt-in durable steps flush immediately + - helper wrappers mark steps durable + - mixed ephemeral/durable sequences flush as expected + +## rollback + +- support rollback steps +- tests: + - rollback executes in reverse order + - rollback persists across restart + - rollback skips completed steps when resumed + - rollback respects abort/eviction + +## misc + +- remove workflow state in favor of actor state +- tests: + - workflow state mirrors actor state transitions + - cancelled workflows report actor state + - storage no longer writes workflow state key + +## types + +- generic messages +- tests: + - typed messages enforce payload shape + - message serialization round-trips typed payloads + - listen helpers preserve generic type inference + +## otel + +- tbd on what the trace id represents +- tbd on actions + diff --git a/rivetkit-typescript/packages/workflow-engine/agents.md b/rivetkit-typescript/packages/workflow-engine/agents.md new file mode 100644 index 0000000000..cd920e763d --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/agents.md @@ -0,0 +1,9 @@ +# Workflow Engine Notes + +## Dirty State Requirements + +- History entries must set `entry.dirty = true` whenever the entry is created or mutated. `flush()` persists dirty entries and clears the flag. +- Entry metadata must set `metadata.dirty = true` whenever metadata fields change. `flush()` persists dirty metadata and clears the flag. +- Name registry writes are tracked by `storage.flushedNameCount`. New names must be registered with `registerName()` before flushing. +- Workflow state/output/error are tracked via `storage.flushedState`, `storage.flushedOutput`, and `storage.flushedError`. Update the fields and call `flush()`; it will write if the value changed. +- `flush()` does not clear workflow output/error keys when values are unset. If you need to clear them, explicitly `driver.delete(buildWorkflowOutputKey())` or `driver.delete(buildWorkflowErrorKey())`. diff --git a/rivetkit-typescript/packages/workflow-engine/architecture.md b/rivetkit-typescript/packages/workflow-engine/architecture.md new file mode 100644 index 0000000000..59e2f6f9a6 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/architecture.md @@ -0,0 +1,466 @@ +# Workflow Engine Architecture + +This document describes the architecture of the workflow engine, a durable execution system for TypeScript. + +## Overview + +The workflow engine enables writing long-running, fault-tolerant workflows as regular async functions. Workflows can be interrupted at any point (process crash, eviction, sleep) and resume from where they left off. This is achieved through: + +1. **History tracking** - Every operation is recorded in persistent storage +2. **Replay** - On restart, operations replay from history instead of re-executing +3. **Deterministic execution** - Same inputs produce same execution path + +## Isolation Model + +**Critical architectural assumption**: Each workflow instance operates on an isolated KV namespace. + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Host System │ +│ ┌─────────────────┐ ┌─────────────────┐ │ +│ │ Workflow A │ │ Workflow B │ │ +│ │ ┌───────────┐ │ │ ┌───────────┐ │ │ +│ │ │ Engine │ │ │ │ Engine │ │ │ +│ │ └─────┬─────┘ │ │ └─────┬─────┘ │ │ +│ │ │ │ │ │ │ │ +│ │ ┌─────▼─────┐ │ │ ┌─────▼─────┐ │ │ +│ │ │ Driver │ │ │ │ Driver │ │ │ +│ │ │(isolated) │ │ │ │(isolated) │ │ │ +│ │ └─────┬─────┘ │ │ └─────┬─────┘ │ │ +│ │ │ │ │ │ │ │ +│ │ ┌─────▼─────┐ │ │ ┌─────▼─────┐ │ │ +│ │ │ KV A │ │ │ │ KV B │ │ │ +│ │ └───────────┘ │ │ └───────────┘ │ │ +│ └─────────────────┘ └─────────────────┘ │ +└─────────────────────────────────────────────────────────────┘ +``` + +Key guarantees: + +1. **KV Isolation** - Each workflow's `EngineDriver` operates on a completely separate KV namespace. The workflow engine does not include workflow IDs in KV keys because isolation is provided by the driver implementation. + +2. **Single Writer** - A workflow instance is the **only** reader and writer of its KV namespace during execution. There is no concurrent access from other workflow instances. + +3. **Message Delivery** - Messages are written to the workflow's isolated KV by external systems (via `WorkflowHandle.message()`), then read by the workflow on its next execution. Since each workflow has its own KV, messages are inherently workflow-scoped. + +4. **External Mutation** - The only external mutations to a workflow's KV are: + - Message delivery (appending to message queue) + - Eviction markers (not yet implemented) + + These are coordinated through the host system's scheduling to avoid conflicts. + +This isolation model means: +- The `EngineDriver` interface has no workflow ID parameters for KV operations +- Keys like `messages/0`, `history/...` are relative to each workflow's namespace +- The host system (e.g., Cloudflare Durable Objects, dedicated actor processes) provides the isolation boundary +- Alarms use workflow IDs because they may be managed by a shared scheduler + +## Driver Requirements + +The `EngineDriver` implementation must satisfy these requirements: + +1. **Sorted list results** - `list()` MUST return entries sorted by key in lexicographic byte order. The workflow engine relies on this for: + - Message FIFO ordering (messages consumed in order received) + - Name registry reconstruction (names at correct indices) + - Deterministic replay behavior + +2. **Atomic batch writes** - `batch()` SHOULD be atomic (all-or-nothing). If atomicity is not possible, partial writes may cause inconsistent state on crash. + +3. **Prefix isolation** - `list(prefix)` and `deletePrefix(prefix)` must only affect keys that start with the exact prefix bytes. + +4. **No concurrent modification** - The driver may assume no other writer modifies the KV during workflow execution (see Isolation Model). + +## Core Concepts + +### Workflow + +A workflow is an async function that receives a `WorkflowContext` and optional input: + +```typescript +async function myWorkflow(ctx: WorkflowContext, input: MyInput): Promise { + const result = await ctx.step("fetch-data", async () => { + return await fetchData(input.id); + }); + return result; +} +``` + +### Entries + +Every operation in a workflow creates an **entry** in the history. Entry types: + +| Type | Purpose | +|------|---------| +| `step` | Execute arbitrary async code | +| `loop` | Iterate with durable state | +| `sleep` | Wait for a duration or timestamp | +| `message` | Wait for external events | +| `join` | Execute branches in parallel, wait for all | +| `race` | Execute branches in parallel, first wins | +| `removed` | Placeholder for migrated-away entries | + +### Location System + +Each entry is identified by a **location** - a path through the workflow's execution tree: + +``` +step("a") -> [0] -> "a" +step("b") -> [1] -> "b" +loop("outer") { + step("inner") -> [2, ~0, 3] -> "outer/~0/inner" +} +join("parallel") { + branch "x" { + step("work") -> [4, 5, 6] -> "parallel/x/work" + } +} +``` + +#### NameIndex Optimization + +Locations use numeric indices into a **name registry** rather than storing strings directly: + +```typescript +type NameIndex = number; +type PathSegment = NameIndex | LoopIterationMarker; +type Location = PathSegment[]; + +// Name registry: ["a", "b", "outer", "inner", "parallel", "x", "work"] +// Location [4, 5, 6] resolves to "parallel/x/work" +``` + +This optimization reduces storage size when the same names appear many times. + +## Component Architecture + +``` +┌─────────────────────────────────────────────────────────────┐ +│ runWorkflow() │ +│ Entry point that orchestrates workflow execution │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ WorkflowContextImpl │ +│ Implements WorkflowContext interface │ +│ - step(), loop(), sleep(), listen(), join(), race() │ +│ - Manages current location │ +│ - Creates branch contexts for parallel execution │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Storage │ +│ In-memory representation of workflow state │ +│ - nameRegistry: string[] │ +│ - history: Map │ +│ - entryMetadata: Map │ +│ - messages: Message[] │ +│ - state: WorkflowState │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ EngineDriver │ +│ Interface for persistent storage (per-workflow isolated) │ +│ - get/set/delete/list for KV operations │ +│ - batch for atomic writes │ +│ - setAlarm/clearAlarm for scheduled wake-ups │ +│ See "Isolation Model" above for KV scoping │ +└─────────────────────────────────────────────────────────────┘ +``` + +## Data Flow + +### First Execution + +``` +1. runWorkflow() called +2. loadStorage() loads empty state from driver +3. Workflow function executes +4. Each ctx.step() call: + a. Check history for existing entry (none found) + b. Create new entry + c. Execute the step callback + d. Save output to entry + e. flush() writes to driver +5. Workflow completes, final flush() +``` + +### Replay Execution + +``` +1. runWorkflow() called +2. loadStorage() loads previous state from driver +3. Workflow function executes +4. Each ctx.step() call: + a. Check history for existing entry (found!) + b. Entry has output -> return immediately (no callback execution) +5. Workflow continues from where it left off +``` + +### Sleep/Message Yielding + +``` +1. ctx.sleep() or ctx.listen() called +2. Check if deadline passed or message available (in memory) +3. If not ready: + a. Throw SleepError or MessageWaitError + b. runWorkflow() catches error + c. flush() saves current state + d. setAlarm() schedules wake-up (for sleep) + e. Return { state: "sleeping", ... } +4. External scheduler calls runWorkflow() again when ready +5. loadStorage() loads any new messages added while yielded +6. Replay proceeds, sleep/message now succeeds +``` + +### Message Delivery Model + +**Important**: The workflow's KV is only mutated by the workflow engine itself during execution. Messages follow a specific delivery pattern: + +``` +1. External system sends message to workflow +2. Message written directly to KV: messages/{index} +3. External system triggers workflow wake-up (scheduler/notify) +4. runWorkflow() called +5. loadStorage() loads ALL messages from KV into memory +6. Workflow checks storage.messages (in-memory, no polling) +7. If message found: consume and continue +8. If not found: yield with MessageWaitError +``` + +There is no polling for messages during execution. Messages must be present in KV before `runWorkflow()` is called for them to be available. + +## Storage Schema + +Data is stored using binary key encoding with fdb-tuple for proper byte ordering: + +``` +Key Format (binary tuples): Value Format: +[1, index] -> BARE (name string) +[2, ...locationSegments] -> BARE+versioning (Entry) +[3, messageId] -> BARE+versioning (Message) +[4, 1] -> text (WorkflowState) +[4, 2] -> CBOR (workflow output) +[4, 3] -> CBOR (WorkflowError) +[4, 4] -> text (version) +[4, 5] -> CBOR (workflow input) +[5, entryId] -> BARE+versioning (EntryMetadata) + +Key prefixes: +1 = NAMES - Name registry +2 = HISTORY - History entries +3 = SIGNALS - Message queue +4 = WORKFLOW - Workflow metadata +5 = ENTRY_METADATA - Entry metadata + +Location segments in keys: +- NameIndex (number) -> encoded directly +- LoopIterationMarker -> [loopIdx, iteration] nested tuple +``` + +The fdb-tuple encoding ensures: +- Proper lexicographic byte ordering for `list()` operations +- Compact representation for numeric indices +- Nested tuples for complex segments (loop iterations) + +### Entry Structure + +```typescript +interface Entry { + id: string; // UUID + location: Location; // Path in execution tree + kind: EntryKind; // Type-specific data + dirty: boolean; // Needs flushing? +} + +interface StepEntry { + type: "step"; + data: { + output?: unknown; // Successful result + error?: string; // Error message if failed + }; +} +``` + +### Metadata Structure + +```typescript +interface EntryMetadata { + status: "pending" | "running" | "completed" | "failed" | "exhausted"; + attempts: number; + lastAttemptAt: number; + createdAt: number; + completedAt?: number; + dirty: boolean; +} +``` + +## Error Handling + +### Retryable Errors + +Regular errors thrown from step callbacks trigger retry logic: + +1. Error caught, saved to entry +2. `StepFailedError` thrown +3. On next run, metadata.attempts checked +4. If attempts < maxRetries, apply backoff and retry +5. If attempts >= maxRetries, throw `StepExhaustedError` + +### Critical Errors + +`CriticalError` and `RollbackError` bypass retry logic: + +```typescript +await ctx.step("validate", async () => { + if (!isValid(input)) { + throw new CriticalError("Invalid input"); + } +}); + +await ctx.step("halt", async () => { + throw new RollbackError("Stop and roll back"); +}); +``` + +### Yield Errors + +These message the workflow should pause: + +| Error | Meaning | +|-------|---------| +| `SleepError` | Waiting for a deadline | +| `MessageWaitError` | Waiting for messages | +| `EvictedError` | Workflow being moved to another worker | + +## Parallel Execution + +### Join (All) + +```typescript +const results = await ctx.join("fetch-all", { + user: { run: async (ctx) => await ctx.step("user", fetchUser) }, + posts: { run: async (ctx) => await ctx.step("posts", fetchPosts) }, +}); +// results.user and results.posts available +``` + +- All branches execute concurrently +- Waits for ALL branches to complete +- If any branch fails, collects all errors into `JoinError` +- Branch state tracked: pending -> running -> completed/failed + +### Race (First) + +```typescript +const { winner, value } = await ctx.race("timeout", [ + { name: "work", run: async (ctx) => await doWork(ctx) }, + { name: "timeout", run: async (ctx) => { await ctx.sleep("wait", 5000); return null; } }, +]); +``` + +- All branches execute concurrently +- Returns when FIRST branch completes +- Other branches are cancelled via AbortMessage +- Winner tracked in history for replay + +## Loop State Management + +Loops maintain durable state across iterations: + +```typescript +await ctx.loop({ + name: "process-items", + state: { cursor: null, processed: 0 }, + commitInterval: 10, + historyEvery: 10, + historyKeep: 10, + run: async (ctx, state) => { + const batch = await ctx.step("fetch", () => fetchBatch(state.cursor)); + if (!batch.items.length) { + return Loop.break(state.processed); + } + await ctx.step("process", () => processBatch(batch.items)); + return Loop.continue({ + cursor: batch.nextCursor, + processed: state.processed + batch.items.length, + }); + }, +}); +``` + +### Commit Interval + +- State is persisted every `commitInterval` iterations +- On crash, replay resumes from last committed state + +### History Retention + +Loop history is trimmed every `historyEvery` iterations, keeping the most recent `historyKeep` iterations. Rollback only replays the last retained iteration. + +After trimming (historyEvery=10, historyKeep=10): + +``` +Before trim at iteration 20: + process-items/~0/fetch, ~0/process + process-items/~1/fetch, ~1/process + ... + process-items/~19/fetch, ~19/process + +After trim: + process-items/~10/fetch, ~10/process (kept) + ... + process-items/~19/fetch, ~19/process (kept) + // Iterations 0-9 deleted +``` + +## Dirty Tracking + +To minimize writes, entries track a `dirty` flag: + +1. New entries created with `dirty: true` +2. Modified entries set `dirty = true` +3. `flush()` only writes entries where `dirty === true` +4. After write, `dirty = false` + +This means replay operations that don't modify state don't trigger writes. + +## Design Decisions + +### Why Path-Based Locations? + +Alternative: Coordinate-based (index into flat array) + +Path-based advantages: +- Human-readable keys for debugging +- Natural hierarchy for nested structures +- Prefix-based queries for loop cleanup +- Stable across code changes (names vs positions) + +### Why NameIndex? + +Locations could store strings directly, but: +- Same names repeat frequently (e.g., "step-1" in every loop iteration) +- Numeric indices compress better +- Registry loaded once, indices resolved in memory + +### Why Dirty Tracking? + +Could flush everything on every operation, but: +- Replay would write identical data +- Batch operations would have redundant writes +- Dirty tracking makes replay essentially read-only + +### Why Sequential Test Execution? + +Tests share a module-level `driver` variable via `beforeEach`. While each test gets a fresh driver, Vitest's parallel execution caused race conditions. Sequential execution ensures isolation. + +## Future Considerations + +1. **Version checking** - Detect workflow code changes +2. **Compaction** - Merge history entries to reduce size +3. **Sharding** - Distribute workflow state across multiple keys +4. **Observability** - Structured logging, metrics, tracing +5. **Workflow composition** - Child workflows, messages between workflows diff --git a/rivetkit-typescript/packages/workflow-engine/docs/advanced.md b/rivetkit-typescript/packages/workflow-engine/docs/advanced.md new file mode 100644 index 0000000000..7fc18be508 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/advanced.md @@ -0,0 +1,27 @@ +# Advanced + +This document covers advanced behavior that impacts long-term workflow operation. + +## History Retention + +Workflow history grows as entries are added. The engine provides a few built-in mechanisms to limit growth: + +- Loop iterations are compacted using `historyEvery` and `historyKeep`. Older iterations are deleted after each retention window, so rollback only replays the last retained iteration. +- `ctx.race()` removes history for losing branches after a winner is chosen. +- `ctx.removed()` lets you keep history compatible while removing old entries. + +If you need to delete an entire workflow’s history, remove its driver namespace or use a new workflow ID. + +## OpenTelemetry Observability + +The workflow engine does not ship with built-in OpenTelemetry tracing. To add observability: + +- Wrap step callbacks with your tracing spans. +- Use step, join, and race names as span identifiers. +- Instrument your driver implementation to emit metrics for storage and scheduling operations. + +This approach keeps tracing deterministic while still giving you visibility into workflow execution. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/architecture.md:90` for history entry details. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/cancellation.md b/rivetkit-typescript/packages/workflow-engine/docs/cancellation.md new file mode 100644 index 0000000000..f446b8dd22 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/cancellation.md @@ -0,0 +1,31 @@ +# Cancellation + +Workflows can be stopped in two ways: eviction (graceful) and cancellation (permanent). Both are initiated through the workflow handle. + +## Eviction (Graceful Stop) + +`handle.evict()` requests the workflow to stop gracefully. The workflow receives an aborted `AbortSignal` and should exit at the next yield point. + +```ts +const handle = runWorkflow("wf-1", workflow, input, driver); +handle.evict(); +``` + +Use `ctx.abortSignal` or `ctx.isEvicted()` to detect eviction inside steps. + +## Cancellation (Permanent Stop) + +`handle.cancel()` marks the workflow as cancelled and clears any alarms. Future runs with the same workflow ID will throw `EvictedError`. + +```ts +await handle.cancel(); +const state = await handle.getState(); // "cancelled" +``` + +## Race Branch Cancellation + +In `ctx.race()`, losing branches are cancelled via `AbortSignal` and typically surface as `CancelledError`. Use `ctx.abortSignal` to exit promptly. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:323` for eviction and cancellation details. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md b/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md new file mode 100644 index 0000000000..66625d5012 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/control-flow.md @@ -0,0 +1,95 @@ +# Control Flow + +Workflow control flow is expressed through deterministic helpers on `WorkflowContextInterface`. These helpers create durable history entries, so the workflow can resume after crashes and restarts without re-running completed work. + +## Overview + +- Every control-flow operation has a durable name. Names must be unique within the current scope. +- Code outside of `ctx.step()` must be deterministic. Use steps for I/O and side effects. +- Avoid native loops; use `ctx.loop()` so iterations are replayable and checkpointed. + +## Loops + +Use `ctx.loop()` for repeatable logic with durable state, periodic checkpoints, and bounded history. + +```ts +import { Loop, type WorkflowContextInterface } from "@rivetkit/workflow-engine"; + +async function processBatches(ctx: WorkflowContextInterface) { + return await ctx.loop({ + name: "process-batches", + state: { cursor: null as string | null, processed: 0 }, + commitInterval: 10, + historyEvery: 10, + historyKeep: 10, + run: async (loopCtx, state) => { + const batch = await loopCtx.step("fetch", () => fetchBatch(state.cursor)); + + if (batch.items.length === 0) { + return Loop.break(state.processed); + } + + await loopCtx.step("process", () => processBatch(batch.items)); + + return Loop.continue({ + cursor: batch.nextCursor, + processed: state.processed + batch.items.length, + }); + }, + }); +} +``` + +Loop state is persisted every `commitInterval` iterations. Loop history is trimmed every `historyEvery` iterations, keeping only the most recent `historyKeep` iterations, so rollback only replays the last retained iteration. + +## Join (Wait for All) + +`ctx.join()` runs named branches in parallel and waits for all of them to complete. Branch errors are collected into a `JoinError`. + +```ts +const results = await ctx.join("fetch-all", { + user: { run: async (ctx) => ctx.step("user", () => fetchUser(id)) }, + posts: { run: async (ctx) => ctx.step("posts", () => fetchPosts(id)) }, +}); +``` + +- Branches run concurrently on the same workflow instance. +- All branches settle before `join` returns. +- Failures raise `JoinError` with per-branch errors. + +## Race (First Wins) + +`ctx.race()` runs branches in parallel and returns the first successful result. Losing branches are cancelled via `AbortSignal`. + +```ts +const { winner, value } = await ctx.race("timeout", [ + { name: "work", run: async (ctx) => ctx.step("do-work", doWork) }, + { name: "timeout", run: async (ctx) => { await ctx.sleep("wait", 30000); return null; } }, +]); +``` + +- The winner is persisted for replay. +- Losing branches receive `ctx.abortSignal` and typically throw `CancelledError`. +- If all branches fail, `RaceError` is thrown. + +## Messages in Control Flow + +`ctx.listen()` and its variants pause the workflow until a message arrives. Message names are part of history, so keep them stable and unique. + +```ts +const approval = await ctx.listen("approval", "approval-granted"); +``` + +Messages are loaded at workflow start. If a message arrives during execution, the workflow yields and picks it up on the next run. + +## Best Practices + +- Use stable names for steps, loops, joins, races, and listens. +- Keep all nondeterministic work inside steps. +- Use loop state to avoid native `while`/`for` loops. +- Handle cancellation via `ctx.abortSignal` in long-running branches. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:155` for loop usage. +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for message waits. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/high-performance-workflows.md b/rivetkit-typescript/packages/workflow-engine/docs/high-performance-workflows.md new file mode 100644 index 0000000000..1aa1c2d6c0 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/high-performance-workflows.md @@ -0,0 +1,46 @@ +# High Performance Workflows + +The workflow engine is designed to minimize writes and replay costs, but you can still tune performance for high-throughput workflows. This document covers the user-facing knobs that reduce storage churn and runtime overhead. + +## Ephemeral Steps + +Set `ephemeral: true` on a step to delay flushing its entry to storage. The step still records history, but the flush is deferred until the next non-ephemeral operation. + +```ts +await ctx.step({ + name: "prepare", + ephemeral: true, + run: async () => preparePayload(), +}); + +await ctx.step({ + name: "send", + ephemeral: true, + run: async () => sendPayload(), +}); + +// This step flushes all pending entries +await ctx.step("finalize", async () => finalizeBatch()); +``` + +Use ephemeral steps for idempotent work where you want to batch writes. Do not use them for critical side effects that must be persisted immediately. + +## Batch Writes Intentionally + +- Group short-lived steps between durable checkpoints. +- Use non-ephemeral steps to ensure important state changes are flushed to storage. +- Prefer `ctx.loop()` with `commitInterval`, `historyEvery`, and `historyKeep` to control persistence and history retention. + +## Parallelism Without Extra Workers + +`ctx.join()` and `ctx.race()` let you run async work in parallel inside a single workflow instance. This keeps history consistent while taking advantage of concurrency. + +## Deterministic Hot Paths + +- Move nondeterministic work into `ctx.step()` callbacks so replay is fast. +- Avoid CPU-heavy work outside of steps; replay executes that code again. +- Use stable names to prevent `HistoryDivergedError` on replays. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:100` for ephemeral step usage. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md b/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md new file mode 100644 index 0000000000..f52ee553a6 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/long-running-workflows.md @@ -0,0 +1,47 @@ +# Long-Running Workflows + +Long-running workflows can pause, sleep, and resume across process restarts. This behavior is powered by durable history and the workflow driver scheduler. + +## Yielding Execution + +Use sleep and listen helpers to yield control while waiting: + +```ts +await ctx.sleep("wait-5-min", 5 * 60 * 1000); +const message = await ctx.listen("wait-approval", "approval"); +``` + +When a workflow yields, `runWorkflow` returns a `WorkflowResult` with `state: "sleeping"`. The driver alarm or a message wake-up triggers the next run. + +## Short vs. Long Sleeps + +- Short sleeps (< `driver.workerPollInterval`) wait in memory. +- Longer sleeps set an alarm via the driver and return control to the scheduler. + +This allows workflows to pause for hours or days without tying up worker memory. + +## Checkpointing Loop State + +`ctx.loop()` persists loop state every `commitInterval` iterations. Loop history is trimmed every `historyEvery` iterations, keeping the most recent `historyKeep` iterations. Rollback only replays the last retained iteration, and long-running loops do not accumulate unbounded history. + +## Handling Eviction + +Workers can be evicted for scaling or deployments. Use `ctx.abortSignal` or `ctx.isEvicted()` to stop work safely: + +```ts +await ctx.step("long-task", async () => { + while (!ctx.isEvicted()) { + await doChunkOfWork(); + } +}); +``` + +Evictions save state and return control to the scheduler, allowing the workflow to resume elsewhere. + +## Driver Considerations + +The `EngineDriver` provides scheduling via `setAlarm` and `clearAlarm`. For long-running workflows, ensure your driver implementation persists alarms reliably and returns due alarms to the runner. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:193` for sleep behavior. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/migrating-workflows.md b/rivetkit-typescript/packages/workflow-engine/docs/migrating-workflows.md new file mode 100644 index 0000000000..f1f2ec3e37 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/migrating-workflows.md @@ -0,0 +1,60 @@ +# Migrating Workflows + +Workflow history is durable, which means the workflow engine replays past entries on restart. When you change workflow code, you have to keep history compatibility in mind so replays still match. This document covers user-facing migration rules and the helpers available. + +## Compatibility Rules + +- Step, loop, join, race, sleep, and message names are part of the workflow history. Renaming or removing them without a migration will cause a `HistoryDivergedError` on replay. +- Code inside `ctx.step()` callbacks can change freely because the step result is replayed instead of re-executed. +- The order of entries matters. Moving a step before or after another entry can break replay unless you migrate the old location. +- Adding new entries is safe as long as they are after existing ones or gated behind new logic that only runs for new workflows. + +## Removing or Renaming Entries + +Use `ctx.removed()` to preserve compatibility when you remove or rename an entry. This writes a placeholder entry into history so the replay tree stays aligned. + +```ts +import { WorkflowContextInterface } from "@rivetkit/workflow-engine"; + +async function checkoutWorkflow(ctx: WorkflowContextInterface, orderId: string) { + // Step removed in v2, keep name/location compatible + await ctx.removed("validate-cart", "step"); + + // New step name replaces the old one + await ctx.step("validate-order", async () => validateOrder(orderId)); +} +``` + +`ctx.removed()` accepts the original name and entry type (`"step"`, `"loop"`, `"join"`, `"race"`, `"sleep"`, or `"message"`). + +## Strategy for Safe Migrations + +- Prefer additive changes. Add new steps while leaving old ones in place until you can safely migrate. +- If you need to move logic, keep the old step name and add a new step name for the new logic. +- For renames, remove the old entry with `ctx.removed()` and introduce the new entry at the new location. +- Avoid branching on `Math.random()` or `Date.now()` outside of steps, because non-determinism will diverge history. + +## Versioning Inputs and Outputs + +Workflow input is persisted on first run. If you evolve the input schema, consider versioning the input or handling both shapes: + +```ts +type CheckoutInput = + | { version: 1; orderId: string } + | { version: 2; orderId: string; coupons?: string[] }; + +async function checkoutWorkflow(ctx: WorkflowContextInterface, input: CheckoutInput) { + const orderId = input.orderId; + const coupons = input.version === 2 ? input.coupons ?? [] : []; + await ctx.step("charge", () => charge(orderId, coupons)); +} +``` + +## When to Start Fresh + +If the workflow history can be discarded, create a new workflow ID or a new workflow type. This lets you avoid complex migrations at the cost of losing the old history for those instances. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:54` for an overview of `ctx.removed()`. +- `rivetkit-typescript/packages/workflow-engine/architecture.md:88` for history and location details. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/retries.md b/rivetkit-typescript/packages/workflow-engine/docs/retries.md new file mode 100644 index 0000000000..cc6024400b --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/retries.md @@ -0,0 +1,63 @@ +# Retries + +Step retries provide fault tolerance for transient failures. The workflow engine tracks attempts per step and retries with deterministic backoff. + +## Automatic Retry Behavior + +- Errors thrown from a step callback are treated as retryable by default. +- Each failure increments the step’s attempt counter. +- The engine computes deterministic exponential backoff and yields until the retry time. +- Once attempts exceed `maxRetries`, a `StepExhaustedError` is thrown. + +## Configuring Retries + +```ts +await ctx.step({ + name: "external-api", + maxRetries: 5, + retryBackoffBase: 200, + retryBackoffMax: 60000, + timeout: 10000, + run: async () => callExternalApi(), +}); +``` + +- `maxRetries` defaults to 3. +- `retryBackoffBase` and `retryBackoffMax` control exponential delay. +- `timeout` limits how long the step can run; timeouts are treated as critical failures. + +## Unrecoverable Errors + +Use `CriticalError` or `RollbackError` for failures that should not retry. Rollback requires a checkpoint: + +```ts +import { CriticalError, RollbackError } from "@rivetkit/workflow-engine"; + +await ctx.step("validate", async () => { + if (!isValid(input)) { + throw new CriticalError("Invalid input"); + } +}); + +await ctx.rollbackCheckpoint("rollback"); + +await ctx.step("halt", async () => { + throw new RollbackError("Stop and roll back"); +}); +``` + +`StepTimeoutError` is also treated as critical, so timeouts bypass retries. + +## Exhaustion and Recovery + +When a step exhausts retries, the workflow fails with `StepExhaustedError`. You can reset exhausted steps using the workflow handle: + +```ts +await handle.recover(); +``` + +`recover()` clears retry metadata, removes the workflow error, and schedules the workflow to run again. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:101` for step configuration defaults. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/rollback-behavior.md b/rivetkit-typescript/packages/workflow-engine/docs/rollback-behavior.md new file mode 100644 index 0000000000..1ca135d650 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/rollback-behavior.md @@ -0,0 +1,54 @@ +# Rollback Behavior + +Rollback handlers let you undo completed steps when a workflow fails with an unrecoverable error. Throwing `RollbackError` or `CriticalError` from a step forces rollback immediately. + +## When Rollback Runs + +Rollback runs only when the workflow encounters an unrecoverable error (any error that is not a retryable step failure). It does not run on evictions or cancellations. Rollback is disabled until you call `ctx.rollbackCheckpoint()`. + +## Ordering and Persistence + +- Rollback handlers execute in reverse order of step completion. +- Each rollback is recorded, so restarts resume unfinished rollbacks instead of repeating completed ones. +- Steps without a rollback handler are skipped. +- Rollback handlers require a prior `rollbackCheckpoint`. +- Loop rollbacks only replay the last retained iteration because loop history is trimmed. + +## Rollback Context + +Rollback handlers receive a `RollbackContextInterface` with `abortSignal` and `isEvicted()` for cooperative shutdown. A checkpoint must be registered before any rollback handlers. + +```ts +await ctx.rollbackCheckpoint("billing"); + +await ctx.step({ + name: "charge", + run: async () => chargeCard(orderId), + rollback: async (rollbackCtx, receipt) => { + if (rollbackCtx.abortSignal.aborted) { + return; + } + await refundCharge(receipt); + }, +}); +``` + +## Failure Behavior + +If a rollback handler throws, the workflow fails and preserves the rollback error in metadata. Subsequent retries or reruns will attempt remaining rollback steps again. If a rollback handler is configured without a checkpoint, the workflow fails with `RollbackCheckpointError`. + +```ts +await ctx.step({ + name: "charge", + run: async () => chargeCard(orderId), + rollback: async () => { + await refundCharge(orderId); + }, +}); +// Throws RollbackCheckpointError because no checkpoint was set. +``` + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:139` for rollback examples. +- `rivetkit-typescript/packages/workflow-engine/tests/rollback.test.ts:19` for rollback ordering. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/state.md b/rivetkit-typescript/packages/workflow-engine/docs/state.md new file mode 100644 index 0000000000..3a9f504e27 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/state.md @@ -0,0 +1,45 @@ +# State + +Workflow state is persisted through the engine driver so workflows can resume deterministically after restarts. This document covers the user-facing state model and how to access it. + +## Workflow State Machine + +Workflows move through these states: + +- `pending`: not started yet +- `running`: currently executing +- `sleeping`: waiting for a deadline or message +- `rolling_back`: executing rollback handlers +- `failed`: failed after rollback +- `completed`: finished successfully +- `cancelled`: permanently stopped + +## Stored Workflow Data + +The engine persists: + +- Workflow input (captured on first run) +- Workflow output (when completed) +- Workflow error metadata (when failed) +- History entries for steps, loops, sleeps, joins, races, messages, and rollback checkpoints + +## Reading State + +Use the workflow handle to query current state or output: + +```ts +const state = await handle.getState(); +const output = await handle.getOutput(); +``` + +## Durable Workflow Data + +There is no mutable shared state inside a workflow. To make data durable: + +- Return values from `ctx.step()` and pass them forward. +- Use `ctx.loop()` state for iterative workflows. +- Avoid storing critical data in module-level variables, which are not persisted. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:406` for workflow state enum. diff --git a/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md b/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md new file mode 100644 index 0000000000..4e8b751c5c --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/docs/waiting-for-events-and-human-in-the-loop.md @@ -0,0 +1,66 @@ +# Waiting for Events & Human-in-the-Loop + +Workflows can pause until external events arrive. This enables human approvals, webhook-driven workflows, and long-lived processes that respond to external signals. + +## Message Delivery Model + +- Messages are persisted via `handle.message()`. +- Messages are loaded at workflow start. +- If a message arrives while the workflow is running, the workflow yields and picks it up on the next run. + +In live mode (`runWorkflow(..., { mode: "live" })`), incoming messages can also wake a workflow waiting on `ctx.listen()`. + +## Listening for Messages + +```ts +const approval = await ctx.listen("wait-approval", "approval-granted"); +``` + +Use the `listen*` variants to wait for multiple messages or apply timeouts: + +```ts +const items = await ctx.listenN("batch", "item-added", 10); +const result = await ctx.listenWithTimeout("approval", "approval-granted", 60000); +``` + +## Deadlines and Timeouts + +Use `listenUntil` or `listenWithTimeout` to model approval windows: + +```ts +const approval = await ctx.listenUntil( + "approval-window", + "approval-granted", + Date.now() + 24 * 60 * 60 * 1000, +); +``` + +If the deadline passes, the method returns `null` instead of throwing. + +## Human-in-the-Loop Example + +```ts +const approval = await ctx.listenWithTimeout( + "manual-approval", + "approval-granted", + 30 * 60 * 1000, +); + +if (!approval) { + await ctx.step("notify-timeout", () => sendTimeoutNotice()); + return "timed-out"; +} + +await ctx.step("proceed", () => runApprovedWork()); +``` + +## Best Practices + +- Keep message names stable and unique per scope. +- Store any state you need for follow-up steps inside step outputs. +- Use `handle.wake()` or send another message if you need to resume a yielded workflow. + +## Related + +- `rivetkit-typescript/packages/workflow-engine/QUICKSTART.md:207` for listen helpers. +- `rivetkit-typescript/packages/workflow-engine/architecture.md:218` for message delivery details. diff --git a/rivetkit-typescript/packages/workflow-engine/package.json b/rivetkit-typescript/packages/workflow-engine/package.json new file mode 100644 index 0000000000..f96b4bf6dd --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/package.json @@ -0,0 +1,69 @@ +{ + "name": "@rivetkit/workflow-engine", + "version": "0.0.1", + "description": "Durable workflow engine with reentrant execution", + "license": "Apache-2.0", + "keywords": [ + "workflow", + "durable", + "reentrant", + "stateful" + ], + "files": [ + "dist", + "src", + "schemas", + "package.json" + ], + "type": "module", + "exports": { + ".": { + "import": { + "types": "./dist/tsup/index.d.ts", + "default": "./dist/tsup/index.js" + }, + "require": { + "types": "./dist/tsup/index.d.cts", + "default": "./dist/tsup/index.cjs" + } + }, + "./testing": { + "import": { + "types": "./dist/tsup/testing.d.ts", + "default": "./dist/tsup/testing.js" + }, + "require": { + "types": "./dist/tsup/testing.d.cts", + "default": "./dist/tsup/testing.cjs" + } + } + }, + "engines": { + "node": ">=18.0.0" + }, + "scripts": { + "build": "pnpm run compile:bare && tsup src/index.ts src/testing.ts", + "compile:bare": "tsx scripts/compile-bare.ts compile schemas/v1.bare -o dist/schemas/v1.ts", + "check-types": "tsc --noEmit", + "lint": "biome check .", + "lint:fix": "biome check --write .", + "test": "vitest run", + "test:watch": "vitest" + }, + "dependencies": { + "@rivetkit/bare-ts": "^0.6.2", + "cbor-x": "^1.6.0", + "fdb-tuple": "^1.0.0", + "vbare": "^0.0.4" + }, + "devDependencies": { + "@bare-ts/tools": "^0.13.0", + "commander": "^12.0.0", + "tsx": "^4.7.0", + "@biomejs/biome": "^2.2.3", + "@types/node": "^22.13.1", + "tsup": "^8.4.0", + "typescript": "^5.7.3", + "vitest": "^3.1.1" + } +} diff --git a/rivetkit-typescript/packages/workflow-engine/schemas/serde.ts b/rivetkit-typescript/packages/workflow-engine/schemas/serde.ts new file mode 100644 index 0000000000..ac7ffbfef7 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/schemas/serde.ts @@ -0,0 +1,655 @@ +/** + * Serialization/deserialization utilities for converting between + * internal TypeScript types and BARE schema types. + */ + +import * as cbor from "cbor-x"; +import type * as v1 from "../dist/schemas/v1.js"; +import { + BranchStatusType as BareBranchStatusType, + EntryStatus as BareEntryStatus, + SleepState as BareSleepState, +} from "../dist/schemas/v1.js"; +import type { + BranchStatus as InternalBranchStatus, + BranchStatusType as InternalBranchStatusType, + Entry as InternalEntry, + EntryKind as InternalEntryKind, + EntryMetadata as InternalEntryMetadata, + EntryStatus as InternalEntryStatus, + Location as InternalLocation, + LoopIterationMarker as InternalLoopIterationMarker, + Message as InternalMessage, + PathSegment as InternalPathSegment, + SleepState as InternalSleepState, + WorkflowState as InternalWorkflowState, +} from "../src/types.js"; +import { + CURRENT_VERSION, + ENTRY_METADATA_VERSIONED, + ENTRY_VERSIONED, + MESSAGE_VERSIONED, + WORKFLOW_METADATA_VERSIONED, +} from "./versioned.js"; + +// === Helper: ArrayBuffer to/from utilities === + +function bufferToArrayBuffer(buf: Uint8Array): ArrayBuffer { + // Create a new ArrayBuffer and copy the data to ensure it's not a SharedArrayBuffer + const arrayBuffer = new ArrayBuffer(buf.byteLength); + new Uint8Array(arrayBuffer).set(buf); + return arrayBuffer; +} + +function encodeCbor(value: unknown): ArrayBuffer { + return bufferToArrayBuffer(cbor.encode(value)); +} + +function decodeCbor(data: ArrayBuffer): T { + return cbor.decode(new Uint8Array(data)) as T; +} + +/** + * Validate that a value is a non-null object. + */ +function assertObject( + value: unknown, + context: string, +): asserts value is Record { + if (typeof value !== "object" || value === null) { + throw new Error(`${context}: expected object, got ${typeof value}`); + } +} + +/** + * Validate that a value is a string. + */ +function assertString( + value: unknown, + context: string, +): asserts value is string { + if (typeof value !== "string") { + throw new Error(`${context}: expected string, got ${typeof value}`); + } +} + +/** + * Validate that a value is a number. + */ +function assertNumber( + value: unknown, + context: string, +): asserts value is number { + if (typeof value !== "number") { + throw new Error(`${context}: expected number, got ${typeof value}`); + } +} + +// === Entry Status Conversion === + +function entryStatusToBare(status: InternalEntryStatus): BareEntryStatus { + switch (status) { + case "pending": + return BareEntryStatus.PENDING; + case "running": + return BareEntryStatus.RUNNING; + case "completed": + return BareEntryStatus.COMPLETED; + case "failed": + return BareEntryStatus.FAILED; + case "exhausted": + return BareEntryStatus.EXHAUSTED; + } +} + +function entryStatusFromBare(status: BareEntryStatus): InternalEntryStatus { + switch (status) { + case BareEntryStatus.PENDING: + return "pending"; + case BareEntryStatus.RUNNING: + return "running"; + case BareEntryStatus.COMPLETED: + return "completed"; + case BareEntryStatus.FAILED: + return "failed"; + case BareEntryStatus.EXHAUSTED: + return "exhausted"; + } +} + +// === Sleep State Conversion === + +function sleepStateToBare(state: InternalSleepState): BareSleepState { + switch (state) { + case "pending": + return BareSleepState.PENDING; + case "completed": + return BareSleepState.COMPLETED; + case "interrupted": + return BareSleepState.INTERRUPTED; + } +} + +function sleepStateFromBare(state: BareSleepState): InternalSleepState { + switch (state) { + case BareSleepState.PENDING: + return "pending"; + case BareSleepState.COMPLETED: + return "completed"; + case BareSleepState.INTERRUPTED: + return "interrupted"; + } +} + +// === Branch Status Type Conversion === + +function branchStatusTypeToBare( + status: InternalBranchStatusType, +): BareBranchStatusType { + switch (status) { + case "pending": + return BareBranchStatusType.PENDING; + case "running": + return BareBranchStatusType.RUNNING; + case "completed": + return BareBranchStatusType.COMPLETED; + case "failed": + return BareBranchStatusType.FAILED; + case "cancelled": + return BareBranchStatusType.CANCELLED; + } +} + +function branchStatusTypeFromBare( + status: BareBranchStatusType, +): InternalBranchStatusType { + switch (status) { + case BareBranchStatusType.PENDING: + return "pending"; + case BareBranchStatusType.RUNNING: + return "running"; + case BareBranchStatusType.COMPLETED: + return "completed"; + case BareBranchStatusType.FAILED: + return "failed"; + case BareBranchStatusType.CANCELLED: + return "cancelled"; + } +} + +// === Location Conversion === + +function locationToBare(location: InternalLocation): v1.Location { + return location.map((segment): v1.PathSegment => { + if (typeof segment === "number") { + return { tag: "NameIndex", val: segment }; + } + return { + tag: "LoopIterationMarker", + val: { + loop: segment.loop, + iteration: segment.iteration, + }, + }; + }); +} + +function locationFromBare(location: v1.Location): InternalLocation { + return location.map((segment): InternalPathSegment => { + if (segment.tag === "NameIndex") { + return segment.val; + } + return { + loop: segment.val.loop, + iteration: segment.val.iteration, + }; + }); +} + +// === Branch Status Conversion === + +function branchStatusToBare(status: InternalBranchStatus): v1.BranchStatus { + return { + status: branchStatusTypeToBare(status.status), + output: status.output !== undefined ? encodeCbor(status.output) : null, + error: status.error ?? null, + }; +} + +function branchStatusFromBare(status: v1.BranchStatus): InternalBranchStatus { + return { + status: branchStatusTypeFromBare(status.status), + output: status.output !== null ? decodeCbor(status.output) : undefined, + error: status.error ?? undefined, + }; +} + +// === Entry Kind Conversion === + +function entryKindToBare(kind: InternalEntryKind): v1.EntryKind { + switch (kind.type) { + case "step": + return { + tag: "StepEntry", + val: { + output: + kind.data.output !== undefined + ? encodeCbor(kind.data.output) + : null, + error: kind.data.error ?? null, + }, + }; + case "loop": + return { + tag: "LoopEntry", + val: { + state: encodeCbor(kind.data.state), + iteration: kind.data.iteration, + output: + kind.data.output !== undefined + ? encodeCbor(kind.data.output) + : null, + }, + }; + case "sleep": + return { + tag: "SleepEntry", + val: { + deadline: BigInt(kind.data.deadline), + state: sleepStateToBare(kind.data.state), + }, + }; + case "message": + return { + tag: "MessageEntry", + val: { + name: kind.data.name, + messageData: encodeCbor(kind.data.data), + }, + }; + case "rollback_checkpoint": + return { + tag: "RollbackCheckpointEntry", + val: { + name: kind.data.name, + }, + }; + case "join": + return { + tag: "JoinEntry", + val: { + branches: new Map( + Object.entries(kind.data.branches).map( + ([name, status]) => [ + name, + branchStatusToBare(status), + ], + ), + ), + }, + }; + case "race": + return { + tag: "RaceEntry", + val: { + winner: kind.data.winner, + branches: new Map( + Object.entries(kind.data.branches).map( + ([name, status]) => [ + name, + branchStatusToBare(status), + ], + ), + ), + }, + }; + case "removed": + return { + tag: "RemovedEntry", + val: { + originalType: kind.data.originalType, + originalName: kind.data.originalName ?? null, + }, + }; + } +} + +function entryKindFromBare(kind: v1.EntryKind): InternalEntryKind { + switch (kind.tag) { + case "StepEntry": + return { + type: "step", + data: { + output: + kind.val.output !== null + ? decodeCbor(kind.val.output) + : undefined, + error: kind.val.error ?? undefined, + }, + }; + case "LoopEntry": + return { + type: "loop", + data: { + state: decodeCbor(kind.val.state), + iteration: kind.val.iteration, + output: + kind.val.output !== null + ? decodeCbor(kind.val.output) + : undefined, + }, + }; + case "SleepEntry": + return { + type: "sleep", + data: { + deadline: Number(kind.val.deadline), + state: sleepStateFromBare(kind.val.state), + }, + }; + case "MessageEntry": + return { + type: "message", + data: { + name: kind.val.name, + data: decodeCbor(kind.val.messageData), + }, + }; + case "RollbackCheckpointEntry": + return { + type: "rollback_checkpoint", + data: { + name: kind.val.name, + }, + }; + case "JoinEntry": + return { + type: "join", + data: { + branches: Object.fromEntries( + Array.from(kind.val.branches.entries()).map( + ([name, status]) => [ + name, + branchStatusFromBare(status), + ], + ), + ), + }, + }; + case "RaceEntry": + return { + type: "race", + data: { + winner: kind.val.winner, + branches: Object.fromEntries( + Array.from(kind.val.branches.entries()).map( + ([name, status]) => [ + name, + branchStatusFromBare(status), + ], + ), + ), + }, + }; + case "RemovedEntry": + return { + type: "removed", + data: { + originalType: kind.val + .originalType as InternalEntryKind["type"], + originalName: kind.val.originalName ?? undefined, + }, + }; + default: + throw new Error( + `Unknown entry kind: ${(kind as { tag: string }).tag}`, + ); + } +} + +// === Entry Conversion & Serialization === + +function entryToBare(entry: InternalEntry): v1.Entry { + return { + id: entry.id, + location: locationToBare(entry.location), + kind: entryKindToBare(entry.kind), + }; +} + +function entryFromBare(bareEntry: v1.Entry): InternalEntry { + return { + id: bareEntry.id, + location: locationFromBare(bareEntry.location), + kind: entryKindFromBare(bareEntry.kind), + dirty: false, + }; +} + +export function serializeEntry(entry: InternalEntry): Uint8Array { + const bareEntry = entryToBare(entry); + return ENTRY_VERSIONED.serializeWithEmbeddedVersion( + bareEntry, + CURRENT_VERSION, + ); +} + +export function deserializeEntry(bytes: Uint8Array): InternalEntry { + const bareEntry = ENTRY_VERSIONED.deserializeWithEmbeddedVersion(bytes); + return entryFromBare(bareEntry); +} + +// === Entry Metadata Conversion & Serialization === + +function entryMetadataToBare( + metadata: InternalEntryMetadata, +): v1.EntryMetadata { + return { + status: entryStatusToBare(metadata.status), + error: metadata.error ?? null, + attempts: metadata.attempts, + lastAttemptAt: BigInt(metadata.lastAttemptAt), + createdAt: BigInt(metadata.createdAt), + completedAt: + metadata.completedAt !== undefined + ? BigInt(metadata.completedAt) + : null, + rollbackCompletedAt: + metadata.rollbackCompletedAt !== undefined + ? BigInt(metadata.rollbackCompletedAt) + : null, + rollbackError: metadata.rollbackError ?? null, + }; +} + +function entryMetadataFromBare( + bareMetadata: v1.EntryMetadata, +): InternalEntryMetadata { + return { + status: entryStatusFromBare(bareMetadata.status), + error: bareMetadata.error ?? undefined, + attempts: bareMetadata.attempts, + lastAttemptAt: Number(bareMetadata.lastAttemptAt), + createdAt: Number(bareMetadata.createdAt), + completedAt: + bareMetadata.completedAt !== null + ? Number(bareMetadata.completedAt) + : undefined, + rollbackCompletedAt: + bareMetadata.rollbackCompletedAt !== null + ? Number(bareMetadata.rollbackCompletedAt) + : undefined, + rollbackError: bareMetadata.rollbackError ?? undefined, + dirty: false, + }; +} + +export function serializeEntryMetadata( + metadata: InternalEntryMetadata, +): Uint8Array { + const bareMetadata = entryMetadataToBare(metadata); + return ENTRY_METADATA_VERSIONED.serializeWithEmbeddedVersion( + bareMetadata, + CURRENT_VERSION, + ); +} + +export function deserializeEntryMetadata( + bytes: Uint8Array, +): InternalEntryMetadata { + const bareMetadata = + ENTRY_METADATA_VERSIONED.deserializeWithEmbeddedVersion(bytes); + return entryMetadataFromBare(bareMetadata); +} + +// === Message Conversion & Serialization === + +function messageToBare(message: InternalMessage): v1.Message { + return { + id: message.id, + name: message.name, + messageData: encodeCbor(message.data), + sentAt: BigInt(message.sentAt), + }; +} + +function messageFromBare(bareMessage: v1.Message): InternalMessage { + return { + id: bareMessage.id, + name: bareMessage.name, + data: decodeCbor(bareMessage.messageData), + sentAt: Number(bareMessage.sentAt), + }; +} + +export function serializeMessage(message: InternalMessage): Uint8Array { + const bareMessage = messageToBare(message); + return MESSAGE_VERSIONED.serializeWithEmbeddedVersion( + bareMessage, + CURRENT_VERSION, + ); +} + +export function deserializeMessage(bytes: Uint8Array): InternalMessage { + const bareMessage = MESSAGE_VERSIONED.deserializeWithEmbeddedVersion(bytes); + return messageFromBare(bareMessage); +} + +// === Workflow Metadata Serialization === +// Note: These are used for reading/writing individual workflow fields + +export function serializeWorkflowState( + state: InternalWorkflowState, +): Uint8Array { + // For simple values, we can encode them directly without the full metadata struct + // Using a single byte for efficiency + const encoder = new TextEncoder(); + return encoder.encode(state); +} + +export function deserializeWorkflowState( + bytes: Uint8Array, +): InternalWorkflowState { + const decoder = new TextDecoder(); + const state = decoder.decode(bytes) as InternalWorkflowState; + const validStates: InternalWorkflowState[] = [ + "pending", + "running", + "sleeping", + "failed", + "completed", + "cancelled", + "rolling_back", + ]; + if (!validStates.includes(state)) { + throw new Error(`Invalid workflow state: ${state}`); + } + return state; +} + +export function serializeWorkflowOutput(output: unknown): Uint8Array { + return cbor.encode(output); +} + +export function deserializeWorkflowOutput(bytes: Uint8Array): T { + try { + return cbor.decode(bytes) as T; + } catch (error) { + throw new Error( + `Failed to deserialize workflow output: ${error instanceof Error ? error.message : String(error)}`, + ); + } +} + +/** + * Structured error type for serialization. + */ +interface SerializedWorkflowError { + name: string; + message: string; + stack?: string; + metadata?: Record; +} + +export function serializeWorkflowError( + error: SerializedWorkflowError, +): Uint8Array { + return cbor.encode(error); +} + +export function deserializeWorkflowError( + bytes: Uint8Array, +): SerializedWorkflowError { + try { + const decoded = cbor.decode(bytes); + // Handle legacy string format + if (typeof decoded === "string") { + return { name: "Error", message: decoded }; + } + assertObject(decoded, "WorkflowError"); + // Validate required fields + const obj = decoded as Record; + assertString(obj.name, "WorkflowError.name"); + assertString(obj.message, "WorkflowError.message"); + return { + name: obj.name, + message: obj.message, + stack: typeof obj.stack === "string" ? obj.stack : undefined, + metadata: + typeof obj.metadata === "object" && obj.metadata !== null + ? (obj.metadata as Record) + : undefined, + }; + } catch { + // If decoding fails, try legacy text format + const decoder = new TextDecoder(); + const message = decoder.decode(bytes); + return { name: "Error", message }; + } +} + +export function serializeWorkflowInput(input: unknown): Uint8Array { + return cbor.encode(input); +} + +export function deserializeWorkflowInput(bytes: Uint8Array): T { + try { + return cbor.decode(bytes) as T; + } catch (error) { + throw new Error( + `Failed to deserialize workflow input: ${error instanceof Error ? error.message : String(error)}`, + ); + } +} + +// === Name Registry Serialization === + +export function serializeName(name: string): Uint8Array { + const encoder = new TextEncoder(); + return encoder.encode(name); +} + +export function deserializeName(bytes: Uint8Array): string { + const decoder = new TextDecoder(); + return decoder.decode(bytes); +} diff --git a/rivetkit-typescript/packages/workflow-engine/schemas/v1.bare b/rivetkit-typescript/packages/workflow-engine/schemas/v1.bare new file mode 100644 index 0000000000..037fea95a9 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/schemas/v1.bare @@ -0,0 +1,203 @@ +# Workflow Engine BARE Schema v1 +# +# This schema defines the binary encoding for workflow engine persistence. +# Types marked with `data` are arbitrary binary blobs (for user-provided data). + +# Opaque user data (CBOR-encoded) +type Cbor data + +# MARK: Location +# Index into the entry name registry +type NameIndex u32 + +# Marker for a loop iteration in a location path +type LoopIterationMarker struct { + loop: NameIndex + iteration: u32 +} + +# A segment in a location path - either a name index or a loop iteration marker +type PathSegment union { + NameIndex | + LoopIterationMarker +} + +# Location identifies where an entry exists in the workflow execution tree +type Location list + +# MARK: Entry Status +type EntryStatus enum { + PENDING + RUNNING + COMPLETED + FAILED + EXHAUSTED +} + +# MARK: Sleep State +type SleepState enum { + PENDING + COMPLETED + INTERRUPTED +} + +# MARK: Branch Status +type BranchStatusType enum { + PENDING + RUNNING + COMPLETED + FAILED + CANCELLED +} + +# MARK: Step Entry +type StepEntry struct { + # Output value (CBOR-encoded arbitrary data) + output: optional + # Error message if step failed + error: optional +} + +# MARK: Loop Entry +type LoopEntry struct { + # Loop state (CBOR-encoded arbitrary data) + state: Cbor + # Current iteration number + iteration: u32 + # Output value if loop completed (CBOR-encoded arbitrary data) + output: optional +} + +# MARK: Sleep Entry +type SleepEntry struct { + # Deadline timestamp in milliseconds + deadline: u64 + # Current sleep state + state: SleepState +} + +# MARK: Message Entry + type MessageEntry struct { + # Message name + name: str + # Message data (CBOR-encoded arbitrary data) + messageData: Cbor + } + + # MARK: Rollback Checkpoint Entry + type RollbackCheckpointEntry struct { + # Checkpoint name + name: str + } + + # MARK: Branch Status + +type BranchStatus struct { + status: BranchStatusType + # Output value if completed (CBOR-encoded arbitrary data) + output: optional + # Error message if failed + error: optional +} + +# MARK: Join Entry +type JoinEntry struct { + # Map of branch name to status + branches: map +} + +# MARK: Race Entry +type RaceEntry struct { + # Name of the winning branch, or null if no winner yet + winner: optional + # Map of branch name to status + branches: map +} + +# MARK: Removed Entry +type RemovedEntry struct { + # Original entry type before removal + originalType: str + # Original entry name + originalName: optional +} + +# MARK: Entry Kind +# Type-specific entry data +type EntryKind union { + StepEntry | + LoopEntry | + SleepEntry | + MessageEntry | + RollbackCheckpointEntry | + JoinEntry | + RaceEntry | + RemovedEntry +} + +# MARK: Entry +# An entry in the workflow history +type Entry struct { + # Unique entry ID + id: str + # Location in the workflow tree + location: Location + # Entry kind and data + kind: EntryKind +} + +# MARK: Entry Metadata +# Metadata for an entry (stored separately, lazily loaded) +type EntryMetadata struct { + status: EntryStatus + # Error message if failed + error: optional + # Number of execution attempts + attempts: u32 + # Last attempt timestamp in milliseconds + lastAttemptAt: u64 + # Creation timestamp in milliseconds + createdAt: u64 + # Completion timestamp in milliseconds + completedAt: optional + # Rollback completion timestamp in milliseconds + rollbackCompletedAt: optional + # Rollback error message if failed + rollbackError: optional +} + +# MARK: Message +# A message in the queue +type Message struct { + # Unique message ID (used as KV key) + id: str + # Message name + name: str + # Message data (CBOR-encoded arbitrary data) + messageData: Cbor + # Timestamp when message was sent in milliseconds + sentAt: u64 +} + +# MARK: Workflow State +type WorkflowState enum { + PENDING + RUNNING + SLEEPING + FAILED + COMPLETED + ROLLING_BACK +} + +# MARK: Workflow Metadata +# Workflow-level metadata stored separately from entries +type WorkflowMetadata struct { + # Current workflow state + state: WorkflowState + # Workflow output if completed (CBOR-encoded arbitrary data) + output: optional + # Error message if failed + error: optional + # Workflow version hash for migration detection + version: optional +} diff --git a/rivetkit-typescript/packages/workflow-engine/schemas/versioned.ts b/rivetkit-typescript/packages/workflow-engine/schemas/versioned.ts new file mode 100644 index 0000000000..cf51f19ab8 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/schemas/versioned.ts @@ -0,0 +1,131 @@ +import { createVersionedDataHandler } from "vbare"; +import * as v1 from "../dist/schemas/v1.js"; + +export const CURRENT_VERSION = 1; + +// Re-export generated types for convenience +export type { + BranchStatus, + Entry, + EntryKind, + EntryMetadata, + JoinEntry, + Location, + LoopEntry, + LoopIterationMarker, + Message, + MessageEntry, + PathSegment, + RaceEntry, + RemovedEntry, + SleepEntry, + StepEntry, + WorkflowMetadata, +} from "../dist/schemas/v1.js"; + +export { + BranchStatusType, + EntryStatus, + SleepState, + WorkflowState, +} from "../dist/schemas/v1.js"; + +// === Entry Handler === + +export const ENTRY_VERSIONED = createVersionedDataHandler({ + deserializeVersion: (bytes, version) => { + switch (version) { + case 1: + return v1.decodeEntry(bytes); + default: + throw new Error(`Unknown Entry version ${version}`); + } + }, + serializeVersion: (data, version) => { + switch (version) { + case 1: + return v1.encodeEntry(data as v1.Entry); + default: + throw new Error(`Unknown Entry version ${version}`); + } + }, + deserializeConverters: () => [], + serializeConverters: () => [], +}); + +// === Entry Metadata Handler === + +export const ENTRY_METADATA_VERSIONED = + createVersionedDataHandler({ + deserializeVersion: (bytes, version) => { + switch (version) { + case 1: + return v1.decodeEntryMetadata(bytes); + default: + throw new Error(`Unknown EntryMetadata version ${version}`); + } + }, + serializeVersion: (data, version) => { + switch (version) { + case 1: + return v1.encodeEntryMetadata(data as v1.EntryMetadata); + default: + throw new Error(`Unknown EntryMetadata version ${version}`); + } + }, + deserializeConverters: () => [], + serializeConverters: () => [], + }); + +// === Message Handler === + +export const MESSAGE_VERSIONED = createVersionedDataHandler({ + deserializeVersion: (bytes, version) => { + switch (version) { + case 1: + return v1.decodeMessage(bytes); + default: + throw new Error(`Unknown Message version ${version}`); + } + }, + serializeVersion: (data, version) => { + switch (version) { + case 1: + return v1.encodeMessage(data as v1.Message); + default: + throw new Error(`Unknown Message version ${version}`); + } + }, + deserializeConverters: () => [], + serializeConverters: () => [], +}); + +// === Workflow Metadata Handler === + +export const WORKFLOW_METADATA_VERSIONED = + createVersionedDataHandler({ + deserializeVersion: (bytes, version) => { + switch (version) { + case 1: + return v1.decodeWorkflowMetadata(bytes); + default: + throw new Error( + `Unknown WorkflowMetadata version ${version}`, + ); + } + }, + serializeVersion: (data, version) => { + switch (version) { + case 1: + return v1.encodeWorkflowMetadata( + data as v1.WorkflowMetadata, + ); + default: + throw new Error( + `Unknown WorkflowMetadata version ${version}`, + ); + } + }, + deserializeConverters: () => [], + serializeConverters: () => [], + }); diff --git a/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts b/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts new file mode 100644 index 0000000000..fecf622c01 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/scripts/compile-bare.ts @@ -0,0 +1,127 @@ +#!/usr/bin/env -S tsx + +/** + * BARE schema compiler for TypeScript + * + * This script compiles .bare schema files to TypeScript using @bare-ts/tools, + * then post-processes the output to: + * 1. Replace @bare-ts/lib import with @rivetkit/bare-ts + * 2. Replace Node.js assert import with a custom assert function + * + * IMPORTANT: Keep the post-processing logic in sync with: + * engine/sdks/rust/runner-protocol/build.rs + */ + +import * as fs from "node:fs/promises"; +import * as path from "node:path"; +import { type Config, transform } from "@bare-ts/tools"; +import { Command } from "commander"; + +const program = new Command(); + +program + .name("bare-compiler") + .description("Compile BARE schemas to TypeScript") + .version("0.0.1"); + +program + .command("compile") + .description("Compile a BARE schema file") + .argument("", "Input BARE schema file") + .option("-o, --output ", "Output file path") + .option("--pedantic", "Enable pedantic mode", false) + .option("--generator ", "Generator type (ts, js, dts, bare)", "ts") + .action(async (input: string, options) => { + try { + const schemaPath = path.resolve(input); + const outputPath = options.output + ? path.resolve(options.output) + : schemaPath.replace(/\.bare$/, ".ts"); + + await compileSchema({ + schemaPath, + outputPath, + config: { + pedantic: options.pedantic, + generator: options.generator, + }, + }); + + console.log(`Successfully compiled ${input} to ${outputPath}`); + } catch (error) { + console.error("Failed to compile schema:", error); + process.exit(1); + } + }); + +program.parse(); + +export interface CompileOptions { + schemaPath: string; + outputPath: string; + config?: Partial; +} + +export async function compileSchema(options: CompileOptions): Promise { + const { schemaPath, outputPath, config = {} } = options; + + const schema = await fs.readFile(schemaPath, "utf-8"); + const outputDir = path.dirname(outputPath); + + await fs.mkdir(outputDir, { recursive: true }); + + const defaultConfig: Partial = { + pedantic: true, + generator: "ts", + ...config, + }; + + let result = transform(schema, defaultConfig); + + result = postProcess(result); + + await fs.writeFile(outputPath, result); +} + +const POST_PROCESS_MARKER = "// @generated - post-processed by compile-bare.ts\n"; + +const ASSERT_FUNCTION = ` +function assert(condition: boolean, message?: string): asserts condition { + if (!condition) throw new Error(message ?? "Assertion failed") +} +`; + +/** + * Post-process the generated TypeScript file to: + * 1. Replace @bare-ts/lib import with @rivetkit/bare-ts + * 2. Replace Node.js assert import with a custom assert function + * + * IMPORTANT: Keep this in sync with engine/sdks/rust/runner-protocol/build.rs + */ +function postProcess(code: string): string { + // Skip if already post-processed + if (code.startsWith(POST_PROCESS_MARKER)) { + return code; + } + + // Replace @bare-ts/lib with @rivetkit/bare-ts + code = code.replace(/@bare-ts\/lib/g, "@rivetkit/bare-ts"); + + // Remove Node.js assert import + code = code.replace(/^import assert from "assert"/m, ""); + + // Add marker and assert function + code = POST_PROCESS_MARKER + code + `\n${ASSERT_FUNCTION}`; + + // Validate post-processing succeeded + if (code.includes("@bare-ts/lib")) { + throw new Error("Failed to replace @bare-ts/lib import"); + } + if (code.includes("import assert from")) { + throw new Error("Failed to remove Node.js assert import"); + } + + return code; +} + +export type { Config } from "@bare-ts/tools"; diff --git a/rivetkit-typescript/packages/workflow-engine/src/context.ts b/rivetkit-typescript/packages/workflow-engine/src/context.ts new file mode 100644 index 0000000000..bcf29915f2 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/context.ts @@ -0,0 +1,1821 @@ +import type { EngineDriver } from "./driver.js"; +import { + CancelledError, + CriticalError, + EntryInProgressError, + EvictedError, + HistoryDivergedError, + JoinError, + MessageWaitError, + RaceError, + RollbackCheckpointError, + RollbackError, + RollbackStopError, + SleepError, + StepExhaustedError, + StepFailedError, +} from "./errors.js"; +import { + appendLoopIteration, + appendName, + emptyLocation, + locationToKey, + registerName, +} from "./location.js"; +import { + consumeMessage, + consumeMessages, + createEntry, + deleteEntriesWithPrefix, + flush, + getEntry, + getOrCreateMetadata, + loadMetadata, + setEntry, +} from "./storage.js"; +import type { + BranchConfig, + BranchOutput, + BranchStatus, + Entry, + EntryKindType, + EntryMetadata, + Location, + LoopConfig, + LoopResult, + RollbackContextInterface, + StepConfig, + Storage, + WorkflowContextInterface, +} from "./types.js"; +import { sleep } from "./utils.js"; + +/** + * Default values for step configuration. + * These are exported so users can reference them when overriding. + */ +export const DEFAULT_MAX_RETRIES = 3; +export const DEFAULT_RETRY_BACKOFF_BASE = 100; +export const DEFAULT_RETRY_BACKOFF_MAX = 30000; +export const DEFAULT_LOOP_COMMIT_INTERVAL = 20; +export const DEFAULT_LOOP_HISTORY_EVERY = 20; +export const DEFAULT_LOOP_HISTORY_KEEP = 20; +export const DEFAULT_STEP_TIMEOUT = 30000; // 30 seconds + +/** + * Calculate backoff delay with exponential backoff. + * Uses deterministic calculation (no jitter) for replay consistency. + */ +function calculateBackoff(attempts: number, base: number, max: number): number { + // Exponential backoff without jitter for determinism + return Math.min(max, base * 2 ** attempts); +} + +/** + * Error thrown when a step times out. + */ +export class StepTimeoutError extends Error { + constructor( + public readonly stepName: string, + public readonly timeoutMs: number, + ) { + super(`Step "${stepName}" timed out after ${timeoutMs}ms`); + this.name = "StepTimeoutError"; + } +} + +/** + * Internal representation of a rollback handler. + */ +export interface RollbackAction { + entryId: string; + name: string; + output: T; + rollback: (ctx: RollbackContextInterface, output: T) => Promise; +} + +/** + * Internal implementation of WorkflowContext. + */ +export class WorkflowContextImpl implements WorkflowContextInterface { + private entryInProgress = false; + private abortController: AbortController; + private currentLocation: Location; + private visitedKeys = new Set(); + private mode: "forward" | "rollback"; + private rollbackActions?: RollbackAction[]; + private rollbackCheckpointSet: boolean; + /** Track names used in current execution to detect duplicates */ + private usedNamesInExecution = new Set(); + + constructor( + public readonly workflowId: string, + private storage: Storage, + private driver: EngineDriver, + location: Location = emptyLocation(), + abortController?: AbortController, + mode: "forward" | "rollback" = "forward", + rollbackActions?: RollbackAction[], + rollbackCheckpointSet = false, + ) { + this.currentLocation = location; + this.abortController = abortController ?? new AbortController(); + this.mode = mode; + this.rollbackActions = rollbackActions; + this.rollbackCheckpointSet = rollbackCheckpointSet; + } + + get abortSignal(): AbortSignal { + return this.abortController.signal; + } + + isEvicted(): boolean { + return this.abortSignal.aborted; + } + + private assertNotInProgress(): void { + if (this.entryInProgress) { + throw new EntryInProgressError(); + } + } + + private checkEvicted(): void { + if (this.abortSignal.aborted) { + throw new EvictedError(); + } + } + + /** + * Create a new branch context for parallel/nested execution. + */ + createBranch( + location: Location, + abortController?: AbortController, + ): WorkflowContextImpl { + return new WorkflowContextImpl( + this.workflowId, + this.storage, + this.driver, + location, + abortController ?? this.abortController, + this.mode, + this.rollbackActions, + this.rollbackCheckpointSet, + ); + } + + /** + * Mark a key as visited. + */ + private markVisited(key: string): void { + this.visitedKeys.add(key); + } + + /** + * Check if a name has already been used at the current location in this execution. + * Throws HistoryDivergedError if duplicate detected. + */ + private checkDuplicateName(name: string): void { + const fullKey = + locationToKey(this.storage, this.currentLocation) + "/" + name; + if (this.usedNamesInExecution.has(fullKey)) { + throw new HistoryDivergedError( + `Duplicate entry name "${name}" at location "${locationToKey(this.storage, this.currentLocation)}". ` + + `Each step/loop/sleep/listen/join/race must have a unique name within its scope.`, + ); + } + this.usedNamesInExecution.add(fullKey); + } + + private stopRollback(): never { + throw new RollbackStopError(); + } + + private stopRollbackIfMissing(entry: Entry | undefined): void { + if (this.mode === "rollback" && !entry) { + this.stopRollback(); + } + } + + private stopRollbackIfIncomplete(condition: boolean): void { + if (this.mode === "rollback" && condition) { + this.stopRollback(); + } + } + + private registerRollbackAction( + config: StepConfig, + entryId: string, + output: T, + metadata: EntryMetadata, + ): void { + if (!config.rollback) { + return; + } + if (metadata.rollbackCompletedAt !== undefined) { + return; + } + this.rollbackActions?.push({ + entryId, + name: config.name, + output: output as unknown, + rollback: config.rollback as ( + ctx: RollbackContextInterface, + output: unknown, + ) => Promise, + }); + } + + /** + * Ensure a rollback checkpoint exists before registering rollback handlers. + */ + private ensureRollbackCheckpoint(config: StepConfig): void { + if (!config.rollback) { + return; + } + if (!this.rollbackCheckpointSet) { + throw new RollbackCheckpointError(); + } + } + + /** + * Validate that all expected entries in the branch were visited. + * Throws HistoryDivergedError if there are unvisited entries. + */ + validateComplete(): void { + const prefix = locationToKey(this.storage, this.currentLocation); + + for (const key of this.storage.history.entries.keys()) { + // Check if this key is under our current location prefix + // Handle root prefix (empty string) specially - all keys are under root + const isUnderPrefix = + prefix === "" + ? true // Root: all keys are children + : key.startsWith(prefix + "/") || key === prefix; + + if (isUnderPrefix) { + if (!this.visitedKeys.has(key)) { + // Entry exists in history but wasn't visited + // This means workflow code may have changed + throw new HistoryDivergedError( + `Entry "${key}" exists in history but was not visited. ` + + `Workflow code may have changed. Use ctx.removed() to handle migrations.`, + ); + } + } + } + } + + /** + * Evict the workflow. + */ + evict(): void { + this.abortController.abort(new EvictedError()); + } + + /** + * Wait for eviction message. + * + * The event listener uses { once: true } to auto-remove after firing, + * preventing memory leaks if this method is called multiple times. + */ + waitForEviction(): Promise { + return new Promise((_, reject) => { + if (this.abortSignal.aborted) { + reject(new EvictedError()); + return; + } + this.abortSignal.addEventListener( + "abort", + () => { + reject(new EvictedError()); + }, + { once: true }, + ); + }); + } + + // === Step === + + async step( + nameOrConfig: string | StepConfig, + run?: () => Promise, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + const config: StepConfig = + typeof nameOrConfig === "string" + ? { name: nameOrConfig, run: run! } + : nameOrConfig; + + this.entryInProgress = true; + try { + return await this.executeStep(config); + } finally { + this.entryInProgress = false; + } + } + + private async executeStep(config: StepConfig): Promise { + this.ensureRollbackCheckpoint(config); + if (this.mode === "rollback") { + return await this.executeStepRollback(config); + } + + // Check for duplicate name in current execution + this.checkDuplicateName(config.name); + + const location = appendName( + this.storage, + this.currentLocation, + config.name, + ); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + if (existing) { + if (existing.kind.type !== "step") { + throw new HistoryDivergedError( + `Expected step "${config.name}" at ${key}, found ${existing.kind.type}`, + ); + } + + const stepData = existing.kind.data; + + // Replay successful result + if (stepData.output !== undefined) { + return stepData.output as T; + } + + // Check if we should retry + const metadata = await loadMetadata( + this.storage, + this.driver, + existing.id, + ); + const maxRetries = config.maxRetries ?? DEFAULT_MAX_RETRIES; + + if (metadata.attempts >= maxRetries) { + throw new StepExhaustedError(config.name, stepData.error); + } + + // Calculate backoff and yield to scheduler + // This allows the workflow to be evicted during backoff + const backoffDelay = calculateBackoff( + metadata.attempts, + config.retryBackoffBase ?? DEFAULT_RETRY_BACKOFF_BASE, + config.retryBackoffMax ?? DEFAULT_RETRY_BACKOFF_MAX, + ); + const retryAt = metadata.lastAttemptAt + backoffDelay; + const now = Date.now(); + + if (now < retryAt) { + // Yield to scheduler - will be woken up at retryAt + throw new SleepError(retryAt); + } + } + + // Execute the step + const entry = + existing ?? createEntry(location, { type: "step", data: {} }); + if (!existing) { + // New entry - register name + const nameIndex = registerName(this.storage, config.name); + entry.location = [...location]; + entry.location[entry.location.length - 1] = nameIndex; + setEntry(this.storage, location, entry); + } + + const metadata = getOrCreateMetadata(this.storage, entry.id); + metadata.status = "running"; + metadata.attempts++; + metadata.lastAttemptAt = Date.now(); + metadata.dirty = true; + + // Get timeout configuration + const timeout = config.timeout ?? DEFAULT_STEP_TIMEOUT; + + try { + // Execute with timeout + const output = await this.executeWithTimeout( + config.run(), + timeout, + config.name, + ); + + if (entry.kind.type === "step") { + entry.kind.data.output = output; + } + entry.dirty = true; + metadata.status = "completed"; + metadata.completedAt = Date.now(); + + // Ephemeral steps don't trigger an immediate flush. This avoids the + // synchronous write overhead for transient operations. Note that the + // step's entry is still marked dirty and WILL be persisted on the + // next flush from a non-ephemeral operation. The purpose of ephemeral + // is to batch writes, not to avoid persistence entirely. + if (!config.ephemeral) { + await flush(this.storage, this.driver); + } + + return output; + } catch (error) { + // Timeout errors are treated as critical (no retry) + if (error instanceof StepTimeoutError) { + if (entry.kind.type === "step") { + entry.kind.data.error = String(error); + } + entry.dirty = true; + metadata.status = "exhausted"; + await flush(this.storage, this.driver); + throw new CriticalError(error.message); + } + + if ( + error instanceof CriticalError || + error instanceof RollbackError + ) { + if (entry.kind.type === "step") { + entry.kind.data.error = String(error); + } + entry.dirty = true; + metadata.status = "exhausted"; + await flush(this.storage, this.driver); + throw error; + } + + if (entry.kind.type === "step") { + entry.kind.data.error = String(error); + } + entry.dirty = true; + metadata.status = "failed"; + + await flush(this.storage, this.driver); + + throw new StepFailedError(config.name, error, metadata.attempts); + } + } + + /** + * Execute a promise with timeout. + * + * Note: This does NOT cancel the underlying operation. JavaScript Promises + * cannot be cancelled once started. When a timeout occurs: + * - The step is marked as failed with StepTimeoutError + * - The underlying async operation continues running in the background + * - Any side effects from the operation may still occur + * + * For cancellable operations, pass ctx.abortSignal to APIs that support AbortSignal: + * + * return fetch(url, { signal: ctx.abortSignal }); + + * }); + * + * Or check ctx.isEvicted() periodically in long-running loops. + */ + private async executeStepRollback(config: StepConfig): Promise { + this.checkDuplicateName(config.name); + this.ensureRollbackCheckpoint(config); + + const location = appendName( + this.storage, + this.currentLocation, + config.name, + ); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + this.markVisited(key); + + if (!existing || existing.kind.type !== "step") { + this.stopRollback(); + } + + const metadata = await loadMetadata( + this.storage, + this.driver, + existing.id, + ); + if (metadata.status !== "completed") { + this.stopRollback(); + } + + const output = existing.kind.data.output as T; + this.registerRollbackAction(config, existing.id, output, metadata); + + return output; + } + + private async executeWithTimeout( + promise: Promise, + timeoutMs: number, + stepName: string, + ): Promise { + if (timeoutMs <= 0) { + return promise; + } + + let timeoutId: ReturnType | undefined; + const timeoutPromise = new Promise((_, reject) => { + timeoutId = setTimeout(() => { + reject(new StepTimeoutError(stepName, timeoutMs)); + }, timeoutMs); + }); + + try { + return await Promise.race([promise, timeoutPromise]); + } finally { + if (timeoutId !== undefined) { + clearTimeout(timeoutId); + } + } + } + + // === Loop === + + async loop( + nameOrConfig: string | LoopConfig, + run?: ( + ctx: WorkflowContextInterface, + ) => Promise>, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + const config: LoopConfig = + typeof nameOrConfig === "string" + ? { name: nameOrConfig, run: run as LoopConfig["run"] } + : nameOrConfig; + + this.entryInProgress = true; + try { + return await this.executeLoop(config); + } finally { + this.entryInProgress = false; + } + } + + private async executeLoop(config: LoopConfig): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(config.name); + + const location = appendName( + this.storage, + this.currentLocation, + config.name, + ); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + let entry: Entry; + let state: S; + let iteration: number; + let rollbackSingleIteration = false; + let rollbackIterationRan = false; + let rollbackOutput: T | undefined; + const rollbackMode = this.mode === "rollback"; + + if (existing) { + if (existing.kind.type !== "loop") { + throw new HistoryDivergedError( + `Expected loop "${config.name}" at ${key}, found ${existing.kind.type}`, + ); + } + + const loopData = existing.kind.data; + + if (rollbackMode) { + if (loopData.output !== undefined) { + return loopData.output as T; + } + rollbackSingleIteration = true; + rollbackIterationRan = false; + rollbackOutput = undefined; + } + + // Loop already completed + if (loopData.output !== undefined) { + return loopData.output as T; + } + + // Resume from saved state + entry = existing; + state = loopData.state as S; + iteration = loopData.iteration; + if (rollbackMode) { + rollbackOutput = loopData.output as T | undefined; + rollbackIterationRan = rollbackOutput !== undefined; + } + } else { + this.stopRollbackIfIncomplete(true); + + // New loop + state = config.state as S; + iteration = 0; + entry = createEntry(location, { + type: "loop", + data: { state, iteration }, + }); + setEntry(this.storage, location, entry); + } + + // TODO: Add validation for commitInterval (must be > 0) + const commitInterval = + config.commitInterval ?? DEFAULT_LOOP_COMMIT_INTERVAL; + const historyEvery = + config.historyEvery ?? + config.commitInterval ?? + DEFAULT_LOOP_HISTORY_EVERY; + const historyKeep = + config.historyKeep ?? + config.commitInterval ?? + DEFAULT_LOOP_HISTORY_KEEP; + + // Execute loop iterations + while (true) { + if (rollbackMode && rollbackSingleIteration) { + if (rollbackIterationRan) { + return rollbackOutput as T; + } + this.stopRollbackIfIncomplete(true); + } + this.checkEvicted(); + + // Create branch for this iteration + const iterationLocation = appendLoopIteration( + this.storage, + location, + config.name, + iteration, + ); + const branchCtx = this.createBranch(iterationLocation); + + // Execute iteration + const result = await config.run(branchCtx, state); + + // Validate branch completed cleanly + branchCtx.validateComplete(); + + if ("break" in result && result.break) { + // Loop complete + if (entry.kind.type === "loop") { + entry.kind.data.output = result.value; + entry.kind.data.state = state; + entry.kind.data.iteration = iteration; + } + entry.dirty = true; + + await flush(this.storage, this.driver); + await this.forgetOldIterations( + location, + iteration + 1, + historyEvery, + historyKeep, + ); + + if (rollbackMode && rollbackSingleIteration) { + rollbackOutput = result.value; + rollbackIterationRan = true; + continue; + } + + return result.value; + } + + // Continue with new state + if ("continue" in result && result.continue) { + state = result.state; + } + iteration++; + + // Periodic commit + if (iteration % commitInterval === 0) { + if (entry.kind.type === "loop") { + entry.kind.data.state = state; + entry.kind.data.iteration = iteration; + } + entry.dirty = true; + + await flush(this.storage, this.driver); + await this.forgetOldIterations( + location, + iteration, + historyEvery, + historyKeep, + ); + } + } + } + + /** + * Delete old loop iteration entries to save storage space. + * + * Loop locations always end with a NameIndex (number) because loops are + * created via appendName(). Even for nested loops, the innermost loop's + * location ends with its name index: + * + * ctx.loop("outer") → location: [outerIndex] + * iteration 0 → location: [{ loop: outerIndex, iteration: 0 }] + * ctx.loop("inner") → location: [{ loop: outerIndex, iteration: 0 }, innerIndex] + * + * This function removes iterations older than (currentIteration - historyKeep) + * every historyEvery iterations. + */ + private async forgetOldIterations( + loopLocation: Location, + currentIteration: number, + historyEvery: number, + historyKeep: number, + ): Promise { + if (historyEvery <= 0 || historyKeep <= 0) { + return; + } + if (currentIteration === 0 || currentIteration % historyEvery !== 0) { + return; + } + const keepFrom = Math.max(0, currentIteration - historyKeep); + // Get the loop name index from the last segment of loopLocation. + // This is always a NameIndex (number) because loop entries are created + // via appendName(), not appendLoopIteration(). + const loopSegment = loopLocation[loopLocation.length - 1]; + if (typeof loopSegment !== "number") { + throw new Error("Expected loop location to end with a name index"); + } + + for (let i = 0; i < keepFrom; i++) { + const iterationLocation: Location = [ + ...loopLocation, + { loop: loopSegment, iteration: i }, + ]; + await deleteEntriesWithPrefix( + this.storage, + this.driver, + iterationLocation, + ); + } + } + + // === Sleep === + + async sleep(name: string, durationMs: number): Promise { + const deadline = Date.now() + durationMs; + return this.sleepUntil(name, deadline); + } + + async sleepUntil(name: string, timestampMs: number): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + await this.executeSleep(name, timestampMs); + } finally { + this.entryInProgress = false; + } + } + + private async executeSleep(name: string, deadline: number): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + const location = appendName(this.storage, this.currentLocation, name); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + let entry: Entry; + + if (existing) { + if (existing.kind.type !== "sleep") { + throw new HistoryDivergedError( + `Expected sleep "${name}" at ${key}, found ${existing.kind.type}`, + ); + } + + const sleepData = existing.kind.data; + + if (this.mode === "rollback") { + this.stopRollbackIfIncomplete(sleepData.state === "pending"); + return; + } + + // Already completed or interrupted + if (sleepData.state !== "pending") { + return; + } + + // Use stored deadline + deadline = sleepData.deadline; + entry = existing; + } else { + this.stopRollbackIfIncomplete(true); + + entry = createEntry(location, { + type: "sleep", + data: { deadline, state: "pending" }, + }); + setEntry(this.storage, location, entry); + entry.dirty = true; + await flush(this.storage, this.driver); + } + + const now = Date.now(); + const remaining = deadline - now; + + if (remaining <= 0) { + // Deadline passed + if (entry.kind.type === "sleep") { + entry.kind.data.state = "completed"; + } + entry.dirty = true; + await flush(this.storage, this.driver); + return; + } + + // Short sleep: wait in memory + if (remaining < this.driver.workerPollInterval) { + await Promise.race([sleep(remaining), this.waitForEviction()]); + + this.checkEvicted(); + + if (entry.kind.type === "sleep") { + entry.kind.data.state = "completed"; + } + entry.dirty = true; + return; + } + + // Long sleep: yield to scheduler + throw new SleepError(deadline); + } + + // === Rollback Checkpoint === + + async rollbackCheckpoint(name: string): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + await this.executeRollbackCheckpoint(name); + } finally { + this.entryInProgress = false; + } + } + + private async executeRollbackCheckpoint(name: string): Promise { + this.checkDuplicateName(name); + + const location = appendName(this.storage, this.currentLocation, name); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + this.markVisited(key); + + if (existing) { + if (existing.kind.type !== "rollback_checkpoint") { + throw new HistoryDivergedError( + `Expected rollback checkpoint "${name}" at ${key}, found ${existing.kind.type}`, + ); + } + this.rollbackCheckpointSet = true; + return; + } + + if (this.mode === "rollback") { + throw new HistoryDivergedError( + `Missing rollback checkpoint "${name}" at ${key}`, + ); + } + + const entry = createEntry(location, { + type: "rollback_checkpoint", + data: { name }, + }); + setEntry(this.storage, location, entry); + entry.dirty = true; + await flush(this.storage, this.driver); + + this.rollbackCheckpointSet = true; + } + + // === Listen === + // + // IMPORTANT: Messages are loaded once at workflow start (in loadStorage). + // If a message is sent via handle.message() DURING workflow execution, + // it won't be visible until the next execution. The workflow will yield + // (SleepError/MessageWaitError), then on the next run, loadStorage() will + // pick up the new message. This is intentional - no polling during execution. + + async listen(name: string, messageName: string): Promise { + const messages = await this.listenN(name, messageName, 1); + return messages[0]; + } + + async listenN( + name: string, + messageName: string, + limit: number, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + return await this.executeListenN(name, messageName, limit); + } finally { + this.entryInProgress = false; + } + } + + private async executeListenN( + name: string, + messageName: string, + limit: number, + ): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + // Check for replay: first check if we have a count entry + const countLocation = appendName( + this.storage, + this.currentLocation, + `${name}:count`, + ); + const countKey = locationToKey(this.storage, countLocation); + const existingCount = this.storage.history.entries.get(countKey); + + // Mark count entry as visited + this.markVisited(countKey); + + this.stopRollbackIfMissing(existingCount); + + if (existingCount && existingCount.kind.type === "message") { + // Replay: read all recorded messages + const count = existingCount.kind.data.data as number; + const results: T[] = []; + + for (let i = 0; i < count; i++) { + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); + const messageKey = locationToKey(this.storage, messageLocation); + + // Mark each message entry as visited + this.markVisited(messageKey); + + const existingMessage = + this.storage.history.entries.get(messageKey); + if ( + existingMessage && + existingMessage.kind.type === "message" + ) { + results.push(existingMessage.kind.data.data as T); + } + } + + return results; + } + + // Try to consume messages immediately + const messages = await consumeMessages( + this.storage, + this.driver, + messageName, + limit, + ); + + if (messages.length > 0) { + // Record each message in history with indexed names + for (let i = 0; i < messages.length; i++) { + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); + const messageEntry = createEntry(messageLocation, { + type: "message", + data: { name: messageName, data: messages[i].data }, + }); + setEntry(this.storage, messageLocation, messageEntry); + + // Mark as visited + this.markVisited(locationToKey(this.storage, messageLocation)); + } + + // Record the count for replay + const countEntry = createEntry(countLocation, { + type: "message", + data: { name: `${messageName}:count`, data: messages.length }, + }); + setEntry(this.storage, countLocation, countEntry); + + await flush(this.storage, this.driver); + + return messages.map((message) => message.data as T); + } + + // No messages found, throw to yield to scheduler + throw new MessageWaitError([messageName]); + } + + async listenWithTimeout( + name: string, + messageName: string, + timeoutMs: number, + ): Promise { + const deadline = Date.now() + timeoutMs; + return this.listenUntil(name, messageName, deadline); + } + + async listenUntil( + name: string, + messageName: string, + timestampMs: number, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + return await this.executeListenUntil( + name, + messageName, + timestampMs, + ); + } finally { + this.entryInProgress = false; + } + } + + private async executeListenUntil( + name: string, + messageName: string, + deadline: number, + ): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + const sleepLocation = appendName( + this.storage, + this.currentLocation, + name, + ); + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:message`, + ); + const sleepKey = locationToKey(this.storage, sleepLocation); + const messageKey = locationToKey(this.storage, messageLocation); + + // Mark entries as visited for validateComplete + this.markVisited(sleepKey); + this.markVisited(messageKey); + + const existingSleep = this.storage.history.entries.get(sleepKey); + + this.stopRollbackIfMissing(existingSleep); + + // Check for replay + if (existingSleep && existingSleep.kind.type === "sleep") { + const sleepData = existingSleep.kind.data; + + if (sleepData.state === "completed") { + return null; + } + + if (sleepData.state === "interrupted") { + const existingMessage = + this.storage.history.entries.get(messageKey); + if ( + existingMessage && + existingMessage.kind.type === "message" + ) { + return existingMessage.kind.data.data as T; + } + throw new HistoryDivergedError( + "Expected message entry after interrupted sleep", + ); + } + + this.stopRollbackIfIncomplete(true); + + deadline = sleepData.deadline; + } else { + this.stopRollbackIfIncomplete(true); + + // Create sleep entry + const sleepEntry = createEntry(sleepLocation, { + type: "sleep", + data: { deadline, state: "pending" }, + }); + setEntry(this.storage, sleepLocation, sleepEntry); + sleepEntry.dirty = true; + await flush(this.storage, this.driver); + } + + const now = Date.now(); + const remaining = deadline - now; + + // Deadline passed, check for message one more time + if (remaining <= 0) { + const message = await consumeMessage( + this.storage, + this.driver, + messageName, + ); + const sleepEntry = getEntry(this.storage, sleepLocation)!; + + if (message) { + if (sleepEntry.kind.type === "sleep") { + sleepEntry.kind.data.state = "interrupted"; + } + sleepEntry.dirty = true; + + const messageEntry = createEntry(messageLocation, { + type: "message", + data: { name: messageName, data: message.data }, + }); + setEntry(this.storage, messageLocation, messageEntry); + await flush(this.storage, this.driver); + + return message.data as T; + } + + if (sleepEntry.kind.type === "sleep") { + sleepEntry.kind.data.state = "completed"; + } + sleepEntry.dirty = true; + await flush(this.storage, this.driver); + return null; + } + + // Check for message (messages are loaded at workflow start, no polling needed) + const message = await consumeMessage( + this.storage, + this.driver, + messageName, + ); + if (message) { + const sleepEntry = getEntry(this.storage, sleepLocation)!; + if (sleepEntry.kind.type === "sleep") { + sleepEntry.kind.data.state = "interrupted"; + } + sleepEntry.dirty = true; + + const messageEntry = createEntry(messageLocation, { + type: "message", + data: { name: messageName, data: message.data }, + }); + setEntry(this.storage, messageLocation, messageEntry); + await flush(this.storage, this.driver); + + return message.data as T; + } + + // Message not available, yield to scheduler until deadline + throw new SleepError(deadline); + } + + async listenNWithTimeout( + name: string, + messageName: string, + limit: number, + timeoutMs: number, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + return await this.executeListenNWithTimeout( + name, + messageName, + limit, + timeoutMs, + ); + } finally { + this.entryInProgress = false; + } + } + + private async executeListenNWithTimeout( + name: string, + messageName: string, + limit: number, + timeoutMs: number, + ): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + // Use a sleep entry to store the deadline for replay + const sleepLocation = appendName( + this.storage, + this.currentLocation, + `${name}:deadline`, + ); + const sleepKey = locationToKey(this.storage, sleepLocation); + const existingSleep = this.storage.history.entries.get(sleepKey); + + this.markVisited(sleepKey); + + this.stopRollbackIfMissing(existingSleep); + + let deadline: number; + + if (existingSleep && existingSleep.kind.type === "sleep") { + // Replay: use stored deadline + deadline = existingSleep.kind.data.deadline; + } else { + // New execution: calculate and store deadline + deadline = Date.now() + timeoutMs; + const sleepEntry = createEntry(sleepLocation, { + type: "sleep", + data: { deadline, state: "pending" }, + }); + setEntry(this.storage, sleepLocation, sleepEntry); + sleepEntry.dirty = true; + } + + return this.executeListenNUntilImpl( + name, + messageName, + limit, + deadline, + ); + } + + async listenNUntil( + name: string, + messageName: string, + limit: number, + timestampMs: number, + ): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + this.entryInProgress = true; + try { + return await this.executeListenNUntilImpl( + name, + messageName, + limit, + timestampMs, + ); + } finally { + this.entryInProgress = false; + } + } + + /** + * Internal implementation for listenNUntil with proper replay support. + * Stores the count and individual messages for deterministic replay. + */ + private async executeListenNUntilImpl( + name: string, + messageName: string, + limit: number, + deadline: number, + ): Promise { + // Check for replay: look for count entry + const countLocation = appendName( + this.storage, + this.currentLocation, + `${name}:count`, + ); + const countKey = locationToKey(this.storage, countLocation); + const existingCount = this.storage.history.entries.get(countKey); + + this.markVisited(countKey); + + this.stopRollbackIfMissing(existingCount); + + if (existingCount && existingCount.kind.type === "message") { + // Replay: read all recorded messages + const count = existingCount.kind.data.data as number; + const results: T[] = []; + + for (let i = 0; i < count; i++) { + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); + const messageKey = locationToKey(this.storage, messageLocation); + + this.markVisited(messageKey); + + const existingMessage = + this.storage.history.entries.get(messageKey); + if ( + existingMessage && + existingMessage.kind.type === "message" + ) { + results.push(existingMessage.kind.data.data as T); + } + } + + return results; + } + + // New execution: collect messages until timeout or limit reached + const results: T[] = []; + + for (let i = 0; i < limit; i++) { + const now = Date.now(); + if (now >= deadline) { + break; + } + + // Try to consume a message + const message = await consumeMessage( + this.storage, + this.driver, + messageName, + ); + if (!message) { + // No message available - check if we should wait + if (results.length === 0) { + // No messages yet - yield to scheduler until deadline + throw new SleepError(deadline); + } + // We have some messages - return what we have + break; + } + + // Record the message + const messageLocation = appendName( + this.storage, + this.currentLocation, + `${name}:${i}`, + ); + const messageEntry = createEntry(messageLocation, { + type: "message", + data: { name: messageName, data: message.data }, + }); + setEntry(this.storage, messageLocation, messageEntry); + this.markVisited(locationToKey(this.storage, messageLocation)); + + results.push(message.data as T); + } + + // Record the count for replay + const countEntry = createEntry(countLocation, { + type: "message", + data: { name: `${messageName}:count`, data: results.length }, + }); + setEntry(this.storage, countLocation, countEntry); + + await flush(this.storage, this.driver); + + return results; + } + + // === Join === + + async join>>( + name: string, + branches: T, + ): Promise<{ [K in keyof T]: BranchOutput }> { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + return await this.executeJoin(name, branches); + } finally { + this.entryInProgress = false; + } + } + + private async executeJoin>>( + name: string, + branches: T, + ): Promise<{ [K in keyof T]: BranchOutput }> { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + const location = appendName(this.storage, this.currentLocation, name); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + this.stopRollbackIfMissing(existing); + + let entry: Entry; + + if (existing) { + if (existing.kind.type !== "join") { + throw new HistoryDivergedError( + `Expected join "${name}" at ${key}, found ${existing.kind.type}`, + ); + } + entry = existing; + } else { + entry = createEntry(location, { + type: "join", + data: { + branches: Object.fromEntries( + Object.keys(branches).map((k) => [ + k, + { status: "pending" as const }, + ]), + ), + }, + }); + setEntry(this.storage, location, entry); + entry.dirty = true; + } + + if (entry.kind.type !== "join") { + throw new HistoryDivergedError("Entry type mismatch"); + } + + this.stopRollbackIfIncomplete( + Object.values(entry.kind.data.branches).some( + (branch) => branch.status !== "completed", + ), + ); + + const joinData = entry.kind.data; + const results: Record = {}; + const errors: Record = {}; + + // Execute all branches in parallel + const branchPromises = Object.entries(branches).map( + async ([branchName, config]) => { + const branchStatus = joinData.branches[branchName]; + + // Already completed + if (branchStatus.status === "completed") { + results[branchName] = branchStatus.output; + return; + } + + // Already failed + if (branchStatus.status === "failed") { + errors[branchName] = new Error(branchStatus.error); + return; + } + + // Execute branch + const branchLocation = appendName( + this.storage, + location, + branchName, + ); + const branchCtx = this.createBranch(branchLocation); + + branchStatus.status = "running"; + entry.dirty = true; + + try { + const output = await config.run(branchCtx); + branchCtx.validateComplete(); + + branchStatus.status = "completed"; + branchStatus.output = output; + results[branchName] = output; + } catch (error) { + branchStatus.status = "failed"; + branchStatus.error = String(error); + errors[branchName] = error as Error; + } + + entry.dirty = true; + }, + ); + + // Wait for ALL branches (no short-circuit on error) + await Promise.allSettled(branchPromises); + await flush(this.storage, this.driver); + + // Throw if any branches failed + if (Object.keys(errors).length > 0) { + throw new JoinError(errors); + } + + return results as { [K in keyof T]: BranchOutput }; + } + + // === Race === + + async race( + name: string, + branches: Array<{ + name: string; + run: (ctx: WorkflowContextInterface) => Promise; + }>, + ): Promise<{ winner: string; value: T }> { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + return await this.executeRace(name, branches); + } finally { + this.entryInProgress = false; + } + } + + private async executeRace( + name: string, + branches: Array<{ + name: string; + run: (ctx: WorkflowContextInterface) => Promise; + }>, + ): Promise<{ winner: string; value: T }> { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + const location = appendName(this.storage, this.currentLocation, name); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + this.stopRollbackIfMissing(existing); + + let entry: Entry; + + if (existing) { + if (existing.kind.type !== "race") { + throw new HistoryDivergedError( + `Expected race "${name}" at ${key}, found ${existing.kind.type}`, + ); + } + entry = existing; + + // Check if we already have a winner + const raceKind = existing.kind; + if (raceKind.data.winner !== null) { + const winnerStatus = + raceKind.data.branches[raceKind.data.winner]; + return { + winner: raceKind.data.winner, + value: winnerStatus.output as T, + }; + } + + this.stopRollbackIfIncomplete(true); + } else { + entry = createEntry(location, { + type: "race", + data: { + winner: null, + branches: Object.fromEntries( + branches.map((b) => [ + b.name, + { status: "pending" as const }, + ]), + ), + }, + }); + setEntry(this.storage, location, entry); + entry.dirty = true; + } + + if (entry.kind.type !== "race") { + throw new HistoryDivergedError("Entry type mismatch"); + } + + const raceData = entry.kind.data; + + // Create abort controller for cancellation + const raceAbortController = new AbortController(); + + // Track all branch promises to wait for cleanup + const branchPromises: Promise[] = []; + + // Track winner info + let winnerName: string | null = null; + let winnerValue: T | null = null; + let settled = false; + let pendingCount = branches.length; + const errors: Record = {}; + const lateErrors: Array<{ name: string; error: string }> = []; + + // Check for replay winners first + for (const branch of branches) { + const branchStatus = raceData.branches[branch.name]; + if ( + branchStatus.status !== "pending" && + branchStatus.status !== "running" + ) { + pendingCount--; + if (branchStatus.status === "completed" && !settled) { + settled = true; + winnerName = branch.name; + winnerValue = branchStatus.output as T; + } + } + } + + // If we found a replay winner, return immediately + if (settled && winnerName !== null && winnerValue !== null) { + return { winner: winnerName, value: winnerValue }; + } + + // Execute branches that need to run + for (const branch of branches) { + const branchStatus = raceData.branches[branch.name]; + + // Skip already completed/cancelled + if ( + branchStatus.status !== "pending" && + branchStatus.status !== "running" + ) { + continue; + } + + const branchLocation = appendName( + this.storage, + location, + branch.name, + ); + const branchCtx = this.createBranch( + branchLocation, + raceAbortController, + ); + + branchStatus.status = "running"; + entry.dirty = true; + + const branchPromise = branch.run(branchCtx).then( + async (output) => { + if (settled) { + // This branch completed after a winner was determined + // Still record the completion for observability + branchStatus.status = "completed"; + branchStatus.output = output; + entry.dirty = true; + return; + } + settled = true; + winnerName = branch.name; + winnerValue = output; + + branchCtx.validateComplete(); + + branchStatus.status = "completed"; + branchStatus.output = output; + raceData.winner = branch.name; + entry.dirty = true; + + // Cancel other branches + raceAbortController.abort(); + }, + (error) => { + pendingCount--; + + if ( + error instanceof CancelledError || + error instanceof EvictedError + ) { + branchStatus.status = "cancelled"; + } else { + branchStatus.status = "failed"; + branchStatus.error = String(error); + + if (settled) { + // Track late errors for observability + lateErrors.push({ + name: branch.name, + error: String(error), + }); + } else { + errors[branch.name] = error; + } + } + entry.dirty = true; + + // All branches failed (only if no winner yet) + if (pendingCount === 0 && !settled) { + settled = true; + } + }, + ); + + branchPromises.push(branchPromise); + } + + // Wait for all branches to complete or be cancelled + await Promise.allSettled(branchPromises); + + // Clean up entries from non-winning branches + if (winnerName !== null) { + for (const branch of branches) { + if (branch.name !== winnerName) { + const branchLocation = appendName( + this.storage, + location, + branch.name, + ); + await deleteEntriesWithPrefix( + this.storage, + this.driver, + branchLocation, + ); + } + } + } + + // Flush final state + await flush(this.storage, this.driver); + + // Log late errors if any (these occurred after a winner was determined) + if (lateErrors.length > 0) { + console.warn( + `Race "${name}" had ${lateErrors.length} branch(es) fail after winner was determined:`, + lateErrors, + ); + } + + // Return result or throw error + if (winnerName !== null && winnerValue !== null) { + return { winner: winnerName, value: winnerValue }; + } + + // All branches failed + throw new RaceError( + "All branches failed", + Object.entries(errors).map(([name, error]) => ({ + name, + error: String(error), + })), + ); + } + + // === Removed === + + async removed(name: string, originalType: EntryKindType): Promise { + this.assertNotInProgress(); + this.checkEvicted(); + + this.entryInProgress = true; + try { + await this.executeRemoved(name, originalType); + } finally { + this.entryInProgress = false; + } + } + + private async executeRemoved( + name: string, + originalType: EntryKindType, + ): Promise { + // Check for duplicate name in current execution + this.checkDuplicateName(name); + + const location = appendName(this.storage, this.currentLocation, name); + const key = locationToKey(this.storage, location); + const existing = this.storage.history.entries.get(key); + + // Mark this entry as visited for validateComplete + this.markVisited(key); + + this.stopRollbackIfMissing(existing); + + if (existing) { + // Validate the existing entry matches what we expect + if ( + existing.kind.type !== "removed" && + existing.kind.type !== originalType + ) { + throw new HistoryDivergedError( + `Expected ${originalType} or removed at ${key}, found ${existing.kind.type}`, + ); + } + + // If it's not already marked as removed, we just skip it + return; + } + + // Create a removed entry placeholder + const entry = createEntry(location, { + type: "removed", + data: { originalType, originalName: name }, + }); + setEntry(this.storage, location, entry); + await flush(this.storage, this.driver); + } +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/driver.ts b/rivetkit-typescript/packages/workflow-engine/src/driver.ts new file mode 100644 index 0000000000..2ed644b40b --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/driver.ts @@ -0,0 +1,88 @@ +/** + * A key-value entry returned from list operations. + */ +export interface KVEntry { + key: Uint8Array; + value: Uint8Array; +} + +/** + * A write operation for batch writes. + */ +export interface KVWrite { + key: Uint8Array; + value: Uint8Array; +} + +/** + * The engine driver provides the KV and scheduling interface. + * Implementations must provide these methods to integrate with different backends. + * + * IMPORTANT: Each workflow instance must have its own isolated driver/KV namespace. + * The workflow engine is the sole reader/writer of its KV during execution. + * KV operations do not include workflow IDs because isolation is provided externally + * by the host system (e.g., Cloudflare Durable Objects, dedicated actor processes). + * + * External systems may only write messages to the KV (via WorkflowHandle.message()). + * See architecture.md "Isolation Model" for details. + */ +export interface EngineDriver { + // === KV Operations === + + /** + * Get a value by key. + * Returns null if the key doesn't exist. + */ + get(key: Uint8Array): Promise; + + /** + * Set a value by key. + */ + set(key: Uint8Array, value: Uint8Array): Promise; + + /** + * Delete a key. + */ + delete(key: Uint8Array): Promise; + + /** + * Delete all keys with a given prefix. + */ + deletePrefix(prefix: Uint8Array): Promise; + + /** + * List all key-value pairs with a given prefix. + * + * IMPORTANT: Results MUST be sorted by key in lexicographic byte order. + * The workflow engine relies on this ordering for correct message FIFO + * processing and name registry reconstruction. Failing to sort will + * cause non-deterministic replay behavior. + */ + list(prefix: Uint8Array): Promise; + + /** + * Batch write multiple key-value pairs. + * Should be atomic if possible. + */ + batch(writes: KVWrite[]): Promise; + + // === Scheduling === + + /** + * Set an alarm to wake the workflow at a specific time. + * @param workflowId The workflow to wake + * @param wakeAt Timestamp in milliseconds when to wake + */ + setAlarm(workflowId: string, wakeAt: number): Promise; + + /** + * Clear any pending alarm for a workflow. + */ + clearAlarm(workflowId: string): Promise; + + /** + * How often the worker polls for work (in milliseconds). + * Affects the threshold for in-memory vs scheduled sleeps. + */ + readonly workerPollInterval: number; +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/errors.ts b/rivetkit-typescript/packages/workflow-engine/src/errors.ts new file mode 100644 index 0000000000..9865355a13 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/errors.ts @@ -0,0 +1,162 @@ +/** + * Thrown from steps to prevent retry. + * Use this when an error is unrecoverable and retrying would be pointless. + */ +export class CriticalError extends Error { + constructor(message: string) { + super(message); + this.name = "CriticalError"; + } +} + +/** + * Thrown from steps to force rollback without retry. + */ +export class RollbackError extends Error { + constructor(message: string) { + super(message); + this.name = "RollbackError"; + } +} + +/** + * Thrown when rollback is used without a checkpoint. + */ +export class RollbackCheckpointError extends Error { + constructor() { + super("Rollback requires a checkpoint before any rollback step"); + this.name = "RollbackCheckpointError"; + } +} + +/** + * Internal: Workflow should sleep until deadline. + * This is thrown to yield control back to the scheduler. + */ +export class SleepError extends Error { + constructor(public readonly deadline: number) { + super(`Sleeping until ${deadline}`); + this.name = "SleepError"; + } +} + +/** + * Internal: Workflow is waiting for messages. + * This is thrown to yield control back to the scheduler. + */ +export class MessageWaitError extends Error { + constructor(public readonly messageNames: string[]) { + super(`Waiting for messages: ${messageNames.join(", ")}`); + this.name = "MessageWaitError"; + } +} + +/** + * Internal: Workflow was evicted. + * This is thrown when the workflow is being gracefully stopped. + */ +export class EvictedError extends Error { + constructor() { + super("Workflow evicted"); + this.name = "EvictedError"; + } +} + +/** + * Internal: Stop rollback traversal. + */ +export class RollbackStopError extends Error { + constructor() { + super("Rollback traversal halted"); + this.name = "RollbackStopError"; + } +} + +/** + * Workflow code changed incompatibly. + * Thrown when history doesn't match the current workflow code. + */ +export class HistoryDivergedError extends Error { + constructor(message: string) { + super(message); + this.name = "HistoryDivergedError"; + } +} + +/** + * Step exhausted all retries. + */ +export class StepExhaustedError extends Error { + constructor( + public readonly stepName: string, + public readonly lastError?: string, + ) { + super( + `Step "${stepName}" exhausted retries: ${lastError ?? "unknown error"}`, + ); + this.name = "StepExhaustedError"; + } +} + +/** + * Step failed (will be retried). + * Internal error used to trigger retry logic. + */ +export class StepFailedError extends Error { + constructor( + public readonly stepName: string, + public readonly originalError: unknown, + public readonly attempts: number, + ) { + super(`Step "${stepName}" failed (attempt ${attempts})`); + this.name = "StepFailedError"; + this.cause = originalError; + } +} + +/** + * Join had branch failures. + */ +export class JoinError extends Error { + constructor(public readonly errors: Record) { + super(`Join failed: ${Object.keys(errors).join(", ")}`); + this.name = "JoinError"; + } +} + +/** + * Race had all branches fail. + */ +export class RaceError extends Error { + constructor( + message: string, + public readonly errors: Array<{ name: string; error: string }>, + ) { + super(message); + this.name = "RaceError"; + } +} + +/** + * Branch was cancelled (used by race). + */ +export class CancelledError extends Error { + constructor() { + super("Branch cancelled"); + this.name = "CancelledError"; + } +} + +/** + * Entry is currently being processed. + * Thrown when user forgets to await a step. + */ +export class EntryInProgressError extends Error { + constructor() { + super( + "Cannot start a new workflow entry while another is in progress. " + + "Did you forget to await the previous step/loop/sleep?", + ); + this.name = "EntryInProgressError"; + } +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/index.ts b/rivetkit-typescript/packages/workflow-engine/src/index.ts new file mode 100644 index 0000000000..a79f342b95 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/index.ts @@ -0,0 +1,862 @@ +// Types + +// Context +export { + DEFAULT_LOOP_COMMIT_INTERVAL, + DEFAULT_LOOP_HISTORY_EVERY, + DEFAULT_LOOP_HISTORY_KEEP, + DEFAULT_MAX_RETRIES, + DEFAULT_RETRY_BACKOFF_BASE, + DEFAULT_RETRY_BACKOFF_MAX, + DEFAULT_STEP_TIMEOUT, + WorkflowContextImpl, +} from "./context.js"; +// Driver +export type { EngineDriver, KVEntry, KVWrite } from "./driver.js"; +// Errors +export { + CancelledError, + CriticalError, + EntryInProgressError, + EvictedError, + HistoryDivergedError, + JoinError, + MessageWaitError, + RaceError, + RollbackCheckpointError, + RollbackError, + SleepError, + StepExhaustedError, + StepFailedError, +} from "./errors.js"; + +// Location utilities +export { + appendLoopIteration, + appendName, + emptyLocation, + isLocationPrefix, + isLoopIterationMarker, + locationsEqual, + locationToKey, + parentLocation, + registerName, + resolveName, +} from "./location.js"; + +// Storage utilities +export { + addMessage, + consumeMessage, + consumeMessages, + createEntry, + createStorage, + deleteEntriesWithPrefix, + flush, + generateId, + getEntry, + getOrCreateMetadata, + loadMetadata, + loadStorage, + setEntry, +} from "./storage.js"; +export type { + BranchConfig, + BranchOutput, + BranchStatus, + Entry, + EntryKind, + EntryKindType, + EntryMetadata, + EntryStatus, + History, + JoinEntry, + Location, + LoopConfig, + LoopEntry, + LoopIterationMarker, + LoopResult, + Message, + MessageEntry, + NameIndex, + PathSegment, + RaceEntry, + RemovedEntry, + RollbackCheckpointEntry, + RollbackContextInterface, + RunWorkflowOptions, + SleepEntry, + SleepState, + StepConfig, + StepEntry, + Storage, + WorkflowContextInterface, + WorkflowFunction, + WorkflowHandle, + WorkflowResult, + WorkflowRunMode, + WorkflowState, +} from "./types.js"; + +// Loop result helpers +export const Loop = { + continue: (state: S): { continue: true; state: S } => ({ + continue: true, + state, + }), + break: (value: T): { break: true; value: T } => ({ + break: true, + value, + }), +}; + +import { + deserializeEntryMetadata, + deserializeWorkflowInput, + deserializeWorkflowOutput, + deserializeWorkflowState, + serializeEntryMetadata, + serializeMessage, + serializeWorkflowInput, + serializeWorkflowState, +} from "../schemas/serde.js"; +import { type RollbackAction, WorkflowContextImpl } from "./context.js"; +// Main workflow runner +import type { EngineDriver } from "./driver.js"; +import { + EvictedError, + MessageWaitError, + RollbackCheckpointError, + RollbackStopError, + SleepError, + StepFailedError, +} from "./errors.js"; +import { + buildEntryMetadataPrefix, + buildMessageKey, + buildWorkflowErrorKey, + buildWorkflowInputKey, + buildWorkflowOutputKey, + buildWorkflowStateKey, +} from "./keys.js"; +import { flush, generateId, loadMetadata, loadStorage } from "./storage.js"; +import type { + RollbackContextInterface, + RunWorkflowOptions, + Storage, + WorkflowFunction, + WorkflowHandle, + WorkflowResult, + WorkflowRunMode, + WorkflowState, +} from "./types.js"; +import { setLongTimeout } from "./utils.js"; + +/** + * Run a workflow and return a handle for managing it. + * + * The workflow starts executing immediately. Use the returned handle to: + * - `handle.result` - Await workflow completion (or yield in `yield` mode) + * - `handle.message()` - Send messages to the workflow + * - `handle.wake()` - Wake the workflow early + * - `handle.evict()` - Request graceful shutdown + * - `handle.getOutput()` / `handle.getState()` - Query status + */ +interface LiveRuntime { + pendingMessageNames: string[]; + messageWaiters: Array<{ names: string[]; resolve: () => void }>; + sleepWaiter?: () => void; + isSleeping: boolean; +} + +function createLiveRuntime(): LiveRuntime { + return { + pendingMessageNames: [], + messageWaiters: [], + isSleeping: false, + }; +} + +function notifyMessage(runtime: LiveRuntime, name: string): void { + runtime.pendingMessageNames.push(name); + + for (let i = 0; i < runtime.messageWaiters.length; i++) { + const waiter = runtime.messageWaiters[i]; + const matchIndex = runtime.pendingMessageNames.findIndex((pending) => + waiter.names.includes(pending), + ); + if (matchIndex !== -1) { + runtime.pendingMessageNames.splice(matchIndex, 1); + runtime.messageWaiters.splice(i, 1); + waiter.resolve(); + break; + } + } +} + +async function waitForMessage( + runtime: LiveRuntime, + names: string[], + abortSignal: AbortSignal, +): Promise { + if (abortSignal.aborted) { + throw new EvictedError(); + } + + const matchIndex = runtime.pendingMessageNames.findIndex((pending) => + names.includes(pending), + ); + if (matchIndex !== -1) { + runtime.pendingMessageNames.splice(matchIndex, 1); + return; + } + + let resolveWaiter: (() => void) | undefined; + const waiter = { + names, + resolve: () => { + resolveWaiter?.(); + }, + }; + + const waiterPromise = new Promise((resolve) => { + resolveWaiter = resolve; + }); + + runtime.messageWaiters.push(waiter); + + try { + await awaitWithEviction(waiterPromise, abortSignal); + } finally { + const index = runtime.messageWaiters.indexOf(waiter); + if (index !== -1) { + runtime.messageWaiters.splice(index, 1); + } + } + + if (abortSignal.aborted) { + throw new EvictedError(); + } +} + +function createEvictionWait(signal: AbortSignal): { + promise: Promise; + cleanup: () => void; +} { + if (signal.aborted) { + return { + promise: Promise.reject(new EvictedError()), + cleanup: () => {}, + }; + } + + let onAbort: (() => void) | undefined; + const promise = new Promise((_, reject) => { + onAbort = () => { + reject(new EvictedError()); + }; + signal.addEventListener("abort", onAbort, { once: true }); + }); + + return { + promise, + cleanup: () => { + if (onAbort) { + signal.removeEventListener("abort", onAbort); + } + }, + }; +} + +function createRollbackContext( + workflowId: string, + abortController: AbortController, +): RollbackContextInterface { + return { + workflowId, + abortSignal: abortController.signal, + isEvicted: () => abortController.signal.aborted, + }; +} + +async function awaitWithEviction( + promise: Promise, + abortSignal: AbortSignal, +): Promise { + const { promise: evictionPromise, cleanup } = + createEvictionWait(abortSignal); + try { + return await Promise.race([promise, evictionPromise]); + } finally { + cleanup(); + } +} + +async function executeRollback( + workflowId: string, + workflowFn: WorkflowFunction, + input: TInput, + driver: EngineDriver, + abortController: AbortController, + storage: Storage, +): Promise { + const rollbackActions: RollbackAction[] = []; + const ctx = new WorkflowContextImpl( + workflowId, + storage, + driver, + undefined, + abortController, + "rollback", + rollbackActions, + ); + + try { + await workflowFn(ctx, input); + } catch (error) { + if (error instanceof EvictedError) { + throw error; + } + if (error instanceof RollbackStopError) { + // Stop replay once we hit incomplete history during rollback. + } else { + // Ignore workflow errors during rollback replay. + } + } + + if (rollbackActions.length === 0) { + return; + } + + const rollbackContext = createRollbackContext(workflowId, abortController); + + for (let i = rollbackActions.length - 1; i >= 0; i--) { + if (abortController.signal.aborted) { + throw new EvictedError(); + } + + const action = rollbackActions[i]; + const metadata = await loadMetadata(storage, driver, action.entryId); + if (metadata.rollbackCompletedAt !== undefined) { + continue; + } + + try { + await awaitWithEviction( + action.rollback(rollbackContext, action.output), + abortController.signal, + ); + metadata.rollbackCompletedAt = Date.now(); + metadata.rollbackError = undefined; + } catch (error) { + if (error instanceof EvictedError) { + throw error; + } + metadata.rollbackError = + error instanceof Error ? error.message : String(error); + throw error; + } finally { + metadata.dirty = true; + await flush(storage, driver); + } + } +} + +async function setSleepState( + storage: Storage, + driver: EngineDriver, + workflowId: string, + deadline: number, +): Promise> { + storage.state = "sleeping"; + await flush(storage, driver); + await driver.setAlarm(workflowId, deadline); + + return { state: "sleeping", sleepUntil: deadline }; +} + +async function setMessageWaitState( + storage: Storage, + driver: EngineDriver, + messageNames: string[], +): Promise> { + storage.state = "sleeping"; + await flush(storage, driver); + + return { state: "sleeping", waitingForMessages: messageNames }; +} + +async function setEvictedState( + storage: Storage, + driver: EngineDriver, +): Promise> { + await flush(storage, driver); + return { state: storage.state }; +} + +async function setRetryState( + storage: Storage, + driver: EngineDriver, + workflowId: string, +): Promise> { + storage.state = "sleeping"; + await flush(storage, driver); + + const retryAt = Date.now() + 100; + await driver.setAlarm(workflowId, retryAt); + + return { state: "sleeping", sleepUntil: retryAt }; +} + +async function setFailedState( + storage: Storage, + driver: EngineDriver, + error: unknown, +): Promise { + storage.state = "failed"; + storage.error = extractErrorInfo(error); + await flush(storage, driver); +} + +async function waitForSleep( + runtime: LiveRuntime, + deadline: number, + abortSignal: AbortSignal, +): Promise { + while (true) { + const remaining = deadline - Date.now(); + if (remaining <= 0) { + return; + } + + let timeoutHandle: ReturnType | undefined; + const timeoutPromise = new Promise((resolve) => { + timeoutHandle = setLongTimeout(resolve, remaining); + }); + + const wakePromise = new Promise((resolve) => { + runtime.sleepWaiter = resolve; + }); + runtime.isSleeping = true; + + try { + await awaitWithEviction( + Promise.race([timeoutPromise, wakePromise]), + abortSignal, + ); + } finally { + runtime.isSleeping = false; + runtime.sleepWaiter = undefined; + timeoutHandle?.abort(); + } + + if (abortSignal.aborted) { + throw new EvictedError(); + } + + if (Date.now() >= deadline) { + return; + } + } +} + +async function executeLiveWorkflow( + workflowId: string, + workflowFn: WorkflowFunction, + input: TInput, + driver: EngineDriver, + abortController: AbortController, + runtime: LiveRuntime, +): Promise> { + let lastResult: WorkflowResult | undefined; + + while (true) { + const result = await executeWorkflow( + workflowId, + workflowFn, + input, + driver, + abortController, + ); + lastResult = result; + + if (result.state !== "sleeping") { + return result; + } + + if (result.waitingForMessages && result.waitingForMessages.length > 0) { + try { + await waitForMessage( + runtime, + result.waitingForMessages, + abortController.signal, + ); + } catch (error) { + if (error instanceof EvictedError) { + return lastResult; + } + throw error; + } + continue; + } + + if (result.sleepUntil !== undefined) { + try { + await waitForSleep( + runtime, + result.sleepUntil, + abortController.signal, + ); + } catch (error) { + if (error instanceof EvictedError) { + return lastResult; + } + throw error; + } + continue; + } + + return result; + } +} + +export function runWorkflow( + workflowId: string, + workflowFn: WorkflowFunction, + input: TInput, + driver: EngineDriver, + options: RunWorkflowOptions = {}, +): WorkflowHandle { + const abortController = new AbortController(); + const mode: WorkflowRunMode = options.mode ?? "yield"; + const liveRuntime = mode === "live" ? createLiveRuntime() : undefined; + + const resultPromise = + mode === "live" && liveRuntime + ? executeLiveWorkflow( + workflowId, + workflowFn, + input, + driver, + abortController, + liveRuntime, + ) + : executeWorkflow( + workflowId, + workflowFn, + input, + driver, + abortController, + ); + + return { + workflowId, + result: resultPromise, + + async message(name: string, data: unknown): Promise { + const messageId = generateId(); + await driver.set( + buildMessageKey(messageId), + serializeMessage({ + id: messageId, + name, + data, + sentAt: Date.now(), + }), + ); + + if (liveRuntime) { + notifyMessage(liveRuntime, name); + } + }, + + async wake(): Promise { + if (liveRuntime) { + if (liveRuntime.isSleeping && liveRuntime.sleepWaiter) { + liveRuntime.sleepWaiter(); + } + return; + } + await driver.setAlarm(workflowId, Date.now()); + }, + + async recover(): Promise { + const stateValue = await driver.get(buildWorkflowStateKey()); + const state = stateValue + ? deserializeWorkflowState(stateValue) + : "pending"; + + if (state !== "failed") { + return; + } + + const metadataEntries = await driver.list( + buildEntryMetadataPrefix(), + ); + const writes: { key: Uint8Array; value: Uint8Array }[] = []; + + for (const entry of metadataEntries) { + const metadata = deserializeEntryMetadata(entry.value); + if ( + metadata.status !== "failed" && + metadata.status !== "exhausted" + ) { + continue; + } + + metadata.status = "pending"; + metadata.attempts = 0; + metadata.lastAttemptAt = 0; + metadata.error = undefined; + metadata.dirty = false; + + writes.push({ + key: entry.key, + value: serializeEntryMetadata(metadata), + }); + } + + if (writes.length > 0) { + await driver.batch(writes); + } + + await driver.delete(buildWorkflowErrorKey()); + await driver.set( + buildWorkflowStateKey(), + serializeWorkflowState("sleeping"), + ); + + if (liveRuntime) { + if (liveRuntime.isSleeping && liveRuntime.sleepWaiter) { + liveRuntime.sleepWaiter(); + } + return; + } + + await driver.setAlarm(workflowId, Date.now()); + }, + + evict(): void { + abortController.abort(new EvictedError()); + }, + + async cancel(): Promise { + abortController.abort(new EvictedError()); + + await driver.set( + buildWorkflowStateKey(), + serializeWorkflowState("cancelled"), + ); + + await driver.clearAlarm(workflowId); + }, + + async getOutput(): Promise { + const value = await driver.get(buildWorkflowOutputKey()); + if (!value) { + return undefined; + } + return deserializeWorkflowOutput(value); + }, + + async getState(): Promise { + const value = await driver.get(buildWorkflowStateKey()); + if (!value) { + return "pending"; + } + return deserializeWorkflowState(value); + }, + }; +} + +/** + * Internal: Execute the workflow and return the result. + */ +async function executeWorkflow( + workflowId: string, + workflowFn: WorkflowFunction, + input: TInput, + driver: EngineDriver, + abortController: AbortController, +): Promise> { + const storage = await loadStorage(driver); + + // Check if workflow was cancelled + if (storage.state === "cancelled") { + throw new EvictedError(); + } + + // Input persistence: store on first run, use stored input on resume + const storedInputBytes = await driver.get(buildWorkflowInputKey()); + let effectiveInput: TInput; + + if (storedInputBytes) { + // Resume: use stored input for deterministic replay + effectiveInput = deserializeWorkflowInput(storedInputBytes); + } else { + // First run: store the input + effectiveInput = input; + await driver.set( + buildWorkflowInputKey(), + serializeWorkflowInput(input), + ); + } + + if (storage.state === "rolling_back") { + try { + await executeRollback( + workflowId, + workflowFn, + effectiveInput, + driver, + abortController, + storage, + ); + } catch (error) { + if (error instanceof EvictedError) { + return { state: storage.state }; + } + throw error; + } + + storage.state = "failed"; + await flush(storage, driver); + + const storedError = storage.error + ? new Error(storage.error.message) + : new Error("Workflow failed"); + if (storage.error?.name) { + storedError.name = storage.error.name; + } + throw storedError; + } + + const ctx = new WorkflowContextImpl( + workflowId, + storage, + driver, + undefined, + abortController, + ); + + storage.state = "running"; + + try { + const output = await workflowFn(ctx, effectiveInput); + + storage.state = "completed"; + storage.output = output; + await flush(storage, driver); + await driver.clearAlarm(workflowId); + + return { state: "completed", output }; + } catch (error) { + if (error instanceof SleepError) { + return await setSleepState( + storage, + driver, + workflowId, + error.deadline, + ); + } + + if (error instanceof MessageWaitError) { + return await setMessageWaitState( + storage, + driver, + error.messageNames, + ); + } + + if (error instanceof EvictedError) { + return await setEvictedState(storage, driver); + } + + if (error instanceof StepFailedError) { + return await setRetryState(storage, driver, workflowId); + } + + if (error instanceof RollbackCheckpointError) { + await setFailedState(storage, driver, error); + throw error; + } + + // Unrecoverable error + storage.error = extractErrorInfo(error); + storage.state = "rolling_back"; + await flush(storage, driver); + + try { + await executeRollback( + workflowId, + workflowFn, + effectiveInput, + driver, + abortController, + storage, + ); + } catch (rollbackError) { + if (rollbackError instanceof EvictedError) { + return { state: storage.state }; + } + throw rollbackError; + } + + storage.state = "failed"; + await flush(storage, driver); + + throw error; + } +} + +/** + * Extract structured error information from an error. + */ +function extractErrorInfo(error: unknown): { + name: string; + message: string; + stack?: string; + metadata?: Record; +} { + if (error instanceof Error) { + const result: { + name: string; + message: string; + stack?: string; + metadata?: Record; + } = { + name: error.name, + message: error.message, + stack: error.stack, + }; + + // Extract custom properties from error + const metadata: Record = {}; + for (const key of Object.keys(error)) { + if (key !== "name" && key !== "message" && key !== "stack") { + const value = (error as unknown as Record)[ + key + ]; + // Only include serializable values + if ( + typeof value === "string" || + typeof value === "number" || + typeof value === "boolean" || + value === null + ) { + metadata[key] = value; + } + } + } + if (Object.keys(metadata).length > 0) { + result.metadata = metadata; + } + + return result; + } + + return { + name: "Error", + message: String(error), + }; +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/keys.ts b/rivetkit-typescript/packages/workflow-engine/src/keys.ts new file mode 100644 index 0000000000..cf7a36f81d --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/keys.ts @@ -0,0 +1,311 @@ +/** + * Binary key encoding/decoding using fdb-tuple. + * All keys are encoded as tuples with integer prefixes for proper sorting. + */ + +import * as tuple from "fdb-tuple"; +import type { Location, LoopIterationMarker, PathSegment } from "./types.js"; + +// === Key Prefixes === +// Using integers for compact encoding and proper sorting + +export const KEY_PREFIX = { + NAMES: 1, // Name registry: [1, index] + HISTORY: 2, // History entries: [2, ...locationSegments] + MESSAGES: 3, // Message queue: [3, index] + WORKFLOW: 4, // Workflow metadata: [4, field] + ENTRY_METADATA: 5, // Entry metadata: [5, entryId] +} as const; + +// Workflow metadata field identifiers +export const WORKFLOW_FIELD = { + STATE: 1, + OUTPUT: 2, + ERROR: 3, + VERSION: 4, + INPUT: 5, +} as const; + +// === Type Definitions === + +// fdb-tuple's TupleItem type - we use a subset +type TupleItem = string | number | boolean | null | TupleItem[]; + +// === Location Segment Encoding === + +/** + * Convert a path segment to tuple elements. + * - NameIndex (number) → just the number + * - LoopIterationMarker → nested tuple [loopIdx, iteration] + */ +function segmentToTuple(segment: PathSegment): TupleItem { + if (typeof segment === "number") { + return segment; + } + // LoopIterationMarker + return [segment.loop, segment.iteration]; +} + +/** + * Convert tuple elements back to a path segment. + */ +function tupleToSegment(element: TupleItem): PathSegment { + if (typeof element === "number") { + return element; + } + if (Array.isArray(element) && element.length === 2) { + const [loop, iteration] = element; + if (typeof loop === "number" && typeof iteration === "number") { + return { loop, iteration } as LoopIterationMarker; + } + } + throw new Error( + `Invalid path segment tuple element: ${JSON.stringify(element)}`, + ); +} + +/** + * Convert a location to tuple elements. + */ +function locationToTupleElements(location: Location): TupleItem[] { + return location.map(segmentToTuple); +} + +/** + * Convert tuple elements back to a location. + */ +function tupleElementsToLocation(elements: TupleItem[]): Location { + return elements.map(tupleToSegment); +} + +// === Helper Functions === + +/** + * Convert Buffer to Uint8Array. + */ +function bufferToUint8Array(buf: Buffer): Uint8Array { + return new Uint8Array(buf.buffer, buf.byteOffset, buf.byteLength); +} + +/** + * Convert Uint8Array to Buffer. + */ +function uint8ArrayToBuffer(arr: Uint8Array): Buffer { + return Buffer.from(arr.buffer, arr.byteOffset, arr.byteLength); +} + +/** + * Pack tuple items and return as Uint8Array. + */ +function pack(items: TupleItem | TupleItem[]): Uint8Array { + const buf = tuple.pack(items); + return bufferToUint8Array(buf); +} + +/** + * Unpack a Uint8Array and return tuple items. + */ +function unpack(data: Uint8Array): TupleItem[] { + const buf = uint8ArrayToBuffer(data); + return tuple.unpack(buf) as TupleItem[]; +} + +// === Key Builders === + +/** + * Build a key for the name registry. + * Key: [1, index] + */ +export function buildNameKey(index: number): Uint8Array { + return pack([KEY_PREFIX.NAMES, index]); +} + +/** + * Build a prefix for listing all names. + * Prefix: [1] + */ +export function buildNamePrefix(): Uint8Array { + return pack([KEY_PREFIX.NAMES]); +} + +/** + * Build a key for a history entry. + * Key: [2, ...locationSegments] + */ +export function buildHistoryKey(location: Location): Uint8Array { + return pack([KEY_PREFIX.HISTORY, ...locationToTupleElements(location)]); +} + +/** + * Build a prefix for listing history entries under a location. + * Prefix: [2, ...locationSegments] + */ +export function buildHistoryPrefix(location: Location): Uint8Array { + return pack([KEY_PREFIX.HISTORY, ...locationToTupleElements(location)]); +} + +/** + * Build a prefix for listing all history entries. + * Prefix: [2] + */ +export function buildHistoryPrefixAll(): Uint8Array { + return pack([KEY_PREFIX.HISTORY]); +} + +/** + * Build a key for a message. + * Key: [3, messageId] + * + * Message IDs can be either: + * - A string UUID (for new messages added via handle.message()) + * - Used with the prefix to list all messages + */ +export function buildMessageKey(messageId: string): Uint8Array { + return pack([KEY_PREFIX.MESSAGES, messageId]); +} + +/** + * Build a prefix for listing all messages. + * Prefix: [3] + */ +export function buildMessagePrefix(): Uint8Array { + return pack([KEY_PREFIX.MESSAGES]); +} + +/** + * Build a key for workflow state. + * Key: [4, 1] + */ +export function buildWorkflowStateKey(): Uint8Array { + return pack([KEY_PREFIX.WORKFLOW, WORKFLOW_FIELD.STATE]); +} + +/** + * Build a key for workflow output. + * Key: [4, 2] + */ +export function buildWorkflowOutputKey(): Uint8Array { + return pack([KEY_PREFIX.WORKFLOW, WORKFLOW_FIELD.OUTPUT]); +} + +/** + * Build a key for workflow error. + * Key: [4, 3] + */ +export function buildWorkflowErrorKey(): Uint8Array { + return pack([KEY_PREFIX.WORKFLOW, WORKFLOW_FIELD.ERROR]); +} + +/** + * Build a key for workflow input. + * Key: [4, 5] + */ +export function buildWorkflowInputKey(): Uint8Array { + return pack([KEY_PREFIX.WORKFLOW, WORKFLOW_FIELD.INPUT]); +} + +/** + * Build a key for entry metadata. + * Key: [5, entryId] + */ +export function buildEntryMetadataKey(entryId: string): Uint8Array { + return pack([KEY_PREFIX.ENTRY_METADATA, entryId]); +} + +/** + * Build a prefix for listing all entry metadata. + * Prefix: [5] + */ +export function buildEntryMetadataPrefix(): Uint8Array { + return pack([KEY_PREFIX.ENTRY_METADATA]); +} + +// === Key Parsers === + +/** + * Parse a name key and return the index. + * Key: [1, index] → index + */ +export function parseNameKey(key: Uint8Array): number { + const elements = unpack(key); + if (elements.length !== 2 || elements[0] !== KEY_PREFIX.NAMES) { + throw new Error("Invalid name key"); + } + return elements[1] as number; +} + +/** + * Parse a history key and return the location. + * Key: [2, ...segments] → Location + */ +export function parseHistoryKey(key: Uint8Array): Location { + const elements = unpack(key); + if (elements.length < 1 || elements[0] !== KEY_PREFIX.HISTORY) { + throw new Error("Invalid history key"); + } + return tupleElementsToLocation(elements.slice(1)); +} + +/** + * Parse a message key and return the message ID. + * Key: [3, messageId] → messageId + */ +export function parseMessageKey(key: Uint8Array): string { + const elements = unpack(key); + if (elements.length !== 2 || elements[0] !== KEY_PREFIX.MESSAGES) { + throw new Error("Invalid message key"); + } + return elements[1] as string; +} + +/** + * Parse an entry metadata key and return the entry ID. + * Key: [5, entryId] → entryId + */ +export function parseEntryMetadataKey(key: Uint8Array): string { + const elements = unpack(key); + if (elements.length !== 2 || elements[0] !== KEY_PREFIX.ENTRY_METADATA) { + throw new Error("Invalid entry metadata key"); + } + return elements[1] as string; +} + +// === Key Comparison Utilities === + +/** + * Check if a key starts with a prefix. + */ +export function keyStartsWith(key: Uint8Array, prefix: Uint8Array): boolean { + if (key.length < prefix.length) { + return false; + } + for (let i = 0; i < prefix.length; i++) { + if (key[i] !== prefix[i]) { + return false; + } + } + return true; +} + +/** + * Compare two keys lexicographically. + * Returns negative if a < b, 0 if a === b, positive if a > b. + */ +export function compareKeys(a: Uint8Array, b: Uint8Array): number { + const minLen = Math.min(a.length, b.length); + for (let i = 0; i < minLen; i++) { + if (a[i] !== b[i]) { + return a[i] - b[i]; + } + } + return a.length - b.length; +} + +/** + * Convert a key to a hex string for debugging. + */ +export function keyToHex(key: Uint8Array): string { + return Array.from(key) + .map((b) => b.toString(16).padStart(2, "0")) + .join(""); +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/location.ts b/rivetkit-typescript/packages/workflow-engine/src/location.ts new file mode 100644 index 0000000000..9898621e71 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/location.ts @@ -0,0 +1,168 @@ +import type { + Location, + LoopIterationMarker, + NameIndex, + PathSegment, + Storage, +} from "./types.js"; + +/** + * Check if a path segment is a loop iteration marker. + */ +export function isLoopIterationMarker( + segment: PathSegment, +): segment is LoopIterationMarker { + return typeof segment === "object" && "loop" in segment; +} + +/** + * Register a name in the registry and return its index. + * If the name already exists, returns the existing index. + */ +export function registerName(storage: Storage, name: string): NameIndex { + const existing = storage.nameRegistry.indexOf(name); + if (existing !== -1) { + return existing; + } + storage.nameRegistry.push(name); + return storage.nameRegistry.length - 1; +} + +/** + * Resolve a name index to its string value. + */ +export function resolveName(storage: Storage, index: NameIndex): string { + const name = storage.nameRegistry[index]; + if (name === undefined) { + throw new Error(`Name index ${index} not found in registry`); + } + return name; +} + +/** + * Convert a location to a KV key string. + * Named entries use their string name, loop iterations use ~N format. + */ +export function locationToKey(storage: Storage, location: Location): string { + return location + .map((segment) => { + if (typeof segment === "number") { + return resolveName(storage, segment); + } + return `~${segment.iteration}`; + }) + .join("/"); +} + +/** + * Append a named segment to a location. + */ +export function appendName( + storage: Storage, + location: Location, + name: string, +): Location { + const nameIndex = registerName(storage, name); + return [...location, nameIndex]; +} + +/** + * Append a loop iteration segment to a location. + */ +export function appendLoopIteration( + storage: Storage, + location: Location, + loopName: string, + iteration: number, +): Location { + const loopIndex = registerName(storage, loopName); + return [...location, { loop: loopIndex, iteration }]; +} + +/** + * Create an empty location (root). + */ +export function emptyLocation(): Location { + return []; +} + +/** + * Get the parent location (all segments except the last). + */ +export function parentLocation(location: Location): Location { + return location.slice(0, -1); +} + +/** + * Check if one location is a prefix of another. + */ +export function isLocationPrefix( + prefix: Location, + location: Location, +): boolean { + if (prefix.length > location.length) { + return false; + } + for (let i = 0; i < prefix.length; i++) { + const prefixSegment = prefix[i]; + const locationSegment = location[i]; + + if (typeof prefixSegment === "number" && typeof locationSegment === "number") { + if (prefixSegment !== locationSegment) { + return false; + } + } else if ( + isLoopIterationMarker(prefixSegment) && + isLoopIterationMarker(locationSegment) + ) { + if ( + prefixSegment.loop !== locationSegment.loop || + prefixSegment.iteration !== locationSegment.iteration + ) { + return false; + } + } else { + return false; + } + } + return true; +} + +/** + * Compare two locations for equality. + */ +export function locationsEqual(a: Location, b: Location): boolean { + if (a.length !== b.length) { + return false; + } + return isLocationPrefix(a, b); +} + +/** + * Get all entry keys that are children of a given location. + * + * Note: Returns a map of key → entry for convenience, not key → location. + * The location can be retrieved from the entry itself via entry.location. + */ +export function getChildEntries( + storage: Storage, + parentLoc: Location, +): Map { + const parentKey = locationToKey(storage, parentLoc); + const children = new Map(); + + for (const [key, entry] of storage.history.entries) { + // Handle empty parent (root) - all entries are children + const isChild = + parentKey === "" + ? true + : key.startsWith(parentKey + "/") || key === parentKey; + + if (isChild) { + // Return the actual entry's location, not the parent location + children.set(key, entry.location); + } + } + + return children; +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/storage.ts b/rivetkit-typescript/packages/workflow-engine/src/storage.ts new file mode 100644 index 0000000000..55b12cecf6 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/storage.ts @@ -0,0 +1,436 @@ +import { + deserializeEntry, + deserializeEntryMetadata, + deserializeMessage, + deserializeName, + deserializeWorkflowError, + deserializeWorkflowOutput, + deserializeWorkflowState, + serializeEntry, + serializeEntryMetadata, + serializeMessage, + serializeName, + serializeWorkflowError, + serializeWorkflowOutput, + serializeWorkflowState, +} from "../schemas/serde.js"; +import type { EngineDriver, KVWrite } from "./driver.js"; +import { + buildEntryMetadataKey, + buildHistoryKey, + buildHistoryPrefix, + buildHistoryPrefixAll, + buildMessageKey, + buildMessagePrefix, + buildNameKey, + buildNamePrefix, + buildWorkflowErrorKey, + buildWorkflowOutputKey, + buildWorkflowStateKey, + compareKeys, + parseNameKey, +} from "./keys.js"; +import { isLocationPrefix, locationToKey } from "./location.js"; +import type { + Entry, + EntryKind, + EntryMetadata, + Location, + Message, + Storage, +} from "./types.js"; + +/** + * Create an empty storage instance. + */ +export function createStorage(): Storage { + return { + nameRegistry: [], + flushedNameCount: 0, + history: { entries: new Map() }, + entryMetadata: new Map(), + messages: [], + output: undefined, + state: "pending", + flushedState: undefined, + error: undefined, + flushedError: undefined, + flushedOutput: undefined, + }; +} + +/** + * Generate a UUID v4. + */ +export function generateId(): string { + return crypto.randomUUID(); +} + +/** + * Create a new entry. + */ +export function createEntry(location: Location, kind: EntryKind): Entry { + return { + id: generateId(), + location, + kind, + dirty: true, + }; +} + +/** + * Create or get metadata for an entry. + */ +export function getOrCreateMetadata( + storage: Storage, + entryId: string, +): EntryMetadata { + let metadata = storage.entryMetadata.get(entryId); + if (!metadata) { + metadata = { + status: "pending", + attempts: 0, + lastAttemptAt: 0, + createdAt: Date.now(), + rollbackCompletedAt: undefined, + rollbackError: undefined, + dirty: true, + }; + storage.entryMetadata.set(entryId, metadata); + } + return metadata; +} + +/** + * Load storage from the driver. + */ +export async function loadStorage(driver: EngineDriver): Promise { + const storage = createStorage(); + + // Load name registry + const nameEntries = await driver.list(buildNamePrefix()); + // Sort by index to ensure correct order + nameEntries.sort((a, b) => compareKeys(a.key, b.key)); + for (const entry of nameEntries) { + const index = parseNameKey(entry.key); + storage.nameRegistry[index] = deserializeName(entry.value); + } + // Track how many names are already persisted + storage.flushedNameCount = storage.nameRegistry.length; + + // Load history entries + const historyEntries = await driver.list(buildHistoryPrefixAll()); + for (const entry of historyEntries) { + const parsed = deserializeEntry(entry.value); + parsed.dirty = false; + // Use locationToKey to match how context.ts looks up entries + const key = locationToKey(storage, parsed.location); + storage.history.entries.set(key, parsed); + } + + // Load messages + const messageEntries = await driver.list(buildMessagePrefix()); + // Sort by index to ensure correct FIFO order + messageEntries.sort((a, b) => compareKeys(a.key, b.key)); + for (const entry of messageEntries) { + const message = deserializeMessage(entry.value); + storage.messages.push(message); + } + + // Load workflow state + const stateValue = await driver.get(buildWorkflowStateKey()); + if (stateValue) { + storage.state = deserializeWorkflowState(stateValue); + storage.flushedState = storage.state; + } + + // Load output if present + const outputValue = await driver.get(buildWorkflowOutputKey()); + if (outputValue) { + storage.output = deserializeWorkflowOutput(outputValue); + storage.flushedOutput = storage.output; + } + + // Load error if present + const errorValue = await driver.get(buildWorkflowErrorKey()); + if (errorValue) { + storage.error = deserializeWorkflowError(errorValue); + storage.flushedError = storage.error; + } + + return storage; +} + +/** + * Load metadata for an entry (lazy loading). + */ +export async function loadMetadata( + storage: Storage, + driver: EngineDriver, + entryId: string, +): Promise { + // Check if already loaded + const existing = storage.entryMetadata.get(entryId); + if (existing) { + return existing; + } + + // Load from driver + const value = await driver.get(buildEntryMetadataKey(entryId)); + if (value) { + const metadata = deserializeEntryMetadata(value); + metadata.dirty = false; + storage.entryMetadata.set(entryId, metadata); + return metadata; + } + + // Create new metadata + return getOrCreateMetadata(storage, entryId); +} + +/** + * Flush all dirty data to the driver. + */ +export async function flush( + storage: Storage, + driver: EngineDriver, +): Promise { + const writes: KVWrite[] = []; + + // Flush only new names (those added since last flush) + for ( + let i = storage.flushedNameCount; + i < storage.nameRegistry.length; + i++ + ) { + const name = storage.nameRegistry[i]; + if (name !== undefined) { + writes.push({ + key: buildNameKey(i), + value: serializeName(name), + }); + } + } + + // Flush dirty entries + for (const [, entry] of storage.history.entries) { + if (entry.dirty) { + writes.push({ + key: buildHistoryKey(entry.location), + value: serializeEntry(entry), + }); + entry.dirty = false; + } + } + + // Flush dirty metadata + for (const [id, metadata] of storage.entryMetadata) { + if (metadata.dirty) { + writes.push({ + key: buildEntryMetadataKey(id), + value: serializeEntryMetadata(metadata), + }); + metadata.dirty = false; + } + } + + // Flush workflow state if changed + if (storage.state !== storage.flushedState) { + writes.push({ + key: buildWorkflowStateKey(), + value: serializeWorkflowState(storage.state), + }); + } + + // Flush output if changed + if ( + storage.output !== undefined && + storage.output !== storage.flushedOutput + ) { + writes.push({ + key: buildWorkflowOutputKey(), + value: serializeWorkflowOutput(storage.output), + }); + } + + // Flush error if changed (compare by message since objects aren't reference-equal) + const errorChanged = + storage.error !== undefined && + (storage.flushedError === undefined || + storage.error.name !== storage.flushedError.name || + storage.error.message !== storage.flushedError.message); + if (errorChanged) { + writes.push({ + key: buildWorkflowErrorKey(), + value: serializeWorkflowError(storage.error!), + }); + } + + if (writes.length > 0) { + await driver.batch(writes); + } + + // Update flushed tracking after successful write + storage.flushedNameCount = storage.nameRegistry.length; + storage.flushedState = storage.state; + storage.flushedOutput = storage.output; + storage.flushedError = storage.error; +} + +/** + * Add a message to the queue. + */ +export async function addMessage( + storage: Storage, + driver: EngineDriver, + name: string, + data: unknown, +): Promise { + const message: Message = { + id: generateId(), + name, + data, + sentAt: Date.now(), + }; + + storage.messages.push(message); + + // Persist immediately using message's unique ID as key + await driver.set(buildMessageKey(message.id), serializeMessage(message)); +} + +/** + * Consume a message from the queue. + * Returns null if no matching message is found. + * Deletes from driver first to prevent duplicates on failure. + */ +export async function consumeMessage( + storage: Storage, + driver: EngineDriver, + messageName: string, +): Promise { + const index = storage.messages.findIndex( + (message) => message.name === messageName, + ); + if (index === -1) { + return null; + } + + const message = storage.messages[index]; + + // Delete from driver first - if this fails, memory is unchanged + await driver.delete(buildMessageKey(message.id)); + + // Only remove from memory after successful driver deletion + storage.messages.splice(index, 1); + + return message; +} + +/** + * Consume up to N messages from the queue. + * + * Uses allSettled to handle partial failures gracefully: + * - If all deletes succeed, messages are removed from memory + * - If some deletes fail, only successfully deleted messages are removed + * - On next load, failed messages will be re-read from KV + */ +export async function consumeMessages( + storage: Storage, + driver: EngineDriver, + messageName: string, + limit: number, +): Promise { + // Find all matching messages up to limit (don't modify memory yet) + const toConsume: { message: Message; index: number }[] = []; + let count = 0; + + for (let i = 0; i < storage.messages.length && count < limit; i++) { + if (storage.messages[i].name === messageName) { + toConsume.push({ message: storage.messages[i], index: i }); + count++; + } + } + + if (toConsume.length === 0) { + return []; + } + + // Delete from driver using allSettled to handle partial failures + const deleteResults = await Promise.allSettled( + toConsume.map(({ message }) => + driver.delete(buildMessageKey(message.id)), + ), + ); + + // Track which messages were successfully deleted + const successfullyDeleted: { message: Message; index: number }[] = []; + for (let i = 0; i < deleteResults.length; i++) { + if (deleteResults[i].status === "fulfilled") { + successfullyDeleted.push(toConsume[i]); + } + } + + // Only remove successfully deleted messages from memory + // Remove in reverse order to preserve indices + for (let i = successfullyDeleted.length - 1; i >= 0; i--) { + const { index } = successfullyDeleted[i]; + storage.messages.splice(index, 1); + } + + return successfullyDeleted.map(({ message }) => message); +} + +/** + * Delete entries with a given location prefix (used for loop forgetting). + * Also cleans up associated metadata from both memory and driver. + */ +export async function deleteEntriesWithPrefix( + storage: Storage, + driver: EngineDriver, + prefixLocation: Location, +): Promise { + // Collect entry IDs for metadata cleanup + const entryIds: string[] = []; + + // Collect entries to delete and their IDs + for (const [key, entry] of storage.history.entries) { + // Check if the entry's location starts with the prefix location + if (isLocationPrefix(prefixLocation, entry.location)) { + entryIds.push(entry.id); + storage.entryMetadata.delete(entry.id); + storage.history.entries.delete(key); + } + } + + // Delete entries from driver using binary prefix + await driver.deletePrefix(buildHistoryPrefix(prefixLocation)); + + // Delete metadata from driver in parallel + await Promise.all( + entryIds.map((id) => driver.delete(buildEntryMetadataKey(id))), + ); +} + +/** + * Get an entry by location. + */ +export function getEntry( + storage: Storage, + location: Location, +): Entry | undefined { + const key = locationToKey(storage, location); + return storage.history.entries.get(key); +} + +/** + * Set an entry by location. + */ +export function setEntry( + storage: Storage, + location: Location, + entry: Entry, +): void { + const key = locationToKey(storage, location); + storage.history.entries.set(key, entry); +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/testing.ts b/rivetkit-typescript/packages/workflow-engine/src/testing.ts new file mode 100644 index 0000000000..3e423a0bd5 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/testing.ts @@ -0,0 +1,134 @@ +import type { EngineDriver, KVEntry, KVWrite } from "./driver.js"; +import { compareKeys, keyStartsWith, keyToHex } from "./keys.js"; +import { sleep } from "./utils.js"; + +/** + * In-memory implementation of EngineDriver for testing. + * Uses binary keys (Uint8Array) with hex encoding for internal Map storage. + */ +export class InMemoryDriver implements EngineDriver { + // Map from hex-encoded key to { originalKey, value } + private kv = new Map(); + private alarms = new Map(); + + /** Simulated latency per operation (ms) */ + latency = 10; + + /** How often the worker polls for work */ + workerPollInterval = 100; + + async get(key: Uint8Array): Promise { + await sleep(this.latency); + const entry = this.kv.get(keyToHex(key)); + return entry?.value ?? null; + } + + async set(key: Uint8Array, value: Uint8Array): Promise { + await sleep(this.latency); + this.kv.set(keyToHex(key), { key, value }); + } + + async delete(key: Uint8Array): Promise { + await sleep(this.latency); + this.kv.delete(keyToHex(key)); + } + + async deletePrefix(prefix: Uint8Array): Promise { + await sleep(this.latency); + for (const [hexKey, entry] of this.kv) { + if (keyStartsWith(entry.key, prefix)) { + this.kv.delete(hexKey); + } + } + } + + async list(prefix: Uint8Array): Promise { + await sleep(this.latency); + const results: KVEntry[] = []; + for (const entry of this.kv.values()) { + if (keyStartsWith(entry.key, prefix)) { + results.push({ key: entry.key, value: entry.value }); + } + } + // Sort by key lexicographically + return results.sort((a, b) => compareKeys(a.key, b.key)); + } + + async batch(writes: KVWrite[]): Promise { + await sleep(this.latency); + for (const { key, value } of writes) { + this.kv.set(keyToHex(key), { key, value }); + } + } + + async setAlarm(workflowId: string, wakeAt: number): Promise { + await sleep(this.latency); + this.alarms.set(workflowId, wakeAt); + } + + async clearAlarm(workflowId: string): Promise { + await sleep(this.latency); + this.alarms.delete(workflowId); + } + + /** + * Get the alarm time for a workflow (for testing). + */ + getAlarm(workflowId: string): number | undefined { + return this.alarms.get(workflowId); + } + + /** + * Check if any alarms are due and return their workflow IDs. + */ + getDueAlarms(): string[] { + const now = Date.now(); + const due: string[] = []; + for (const [workflowId, wakeAt] of this.alarms) { + if (wakeAt <= now) { + due.push(workflowId); + } + } + return due; + } + + /** + * Clear all data (for testing). + */ + clear(): void { + this.kv.clear(); + this.alarms.clear(); + } + + /** + * Get a snapshot of all data (for testing/debugging). + */ + snapshot(): { + kv: Record; + alarms: Record; + } { + const kvSnapshot: Record = {}; + for (const [hexKey, entry] of this.kv) { + kvSnapshot[hexKey] = entry.value; + } + return { + kv: kvSnapshot, + alarms: Object.fromEntries(this.alarms), + }; + } + + /** + * Get all hex-encoded keys (for testing). + */ + keys(): string[] { + return [...this.kv.keys()]; + } +} + +// Export serde functions for testing +export { serializeMessage } from "../schemas/serde.js"; +// Re-export main exports for convenience +export * from "./index.js"; + +// Export key builders for testing +export { buildMessageKey } from "./keys.js"; diff --git a/rivetkit-typescript/packages/workflow-engine/src/types.ts b/rivetkit-typescript/packages/workflow-engine/src/types.ts new file mode 100644 index 0000000000..ea2a530c9e --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/types.ts @@ -0,0 +1,442 @@ +/** + * Index into the entry name registry. + * Names are stored once and referenced by this index to avoid repetition. + */ +export type NameIndex = number; + +/** + * A segment in a location path. + * Either a name index (for named entries) or a loop iteration marker. + */ +export type PathSegment = NameIndex | LoopIterationMarker; + +/** + * Marker for a loop iteration in a location path. + */ +export interface LoopIterationMarker { + loop: NameIndex; + iteration: number; +} + +/** + * Location identifies where an entry exists in the workflow execution tree. + * It forms a path from the root through loops, joins, and branches. + */ +export type Location = PathSegment[]; + +/** + * Current state of a sleep entry. + */ +export type SleepState = "pending" | "completed" | "interrupted"; + +/** + * Status of an entry in the workflow. + */ +export type EntryStatus = + | "pending" + | "running" + | "completed" + | "failed" + | "exhausted"; + +/** + * Status of a branch in join/race. + */ +export type BranchStatusType = + | "pending" + | "running" + | "completed" + | "failed" + | "cancelled"; + +/** + * Current state of the workflow. + */ +export type WorkflowState = + | "pending" + | "running" + | "sleeping" + | "failed" + | "completed" + | "cancelled" + | "rolling_back"; + +/** + * Step entry data. + */ +export interface StepEntry { + output?: unknown; + error?: string; +} + +/** + * Loop entry data. + */ +export interface LoopEntry { + state: unknown; + iteration: number; + output?: unknown; +} + +/** + * Sleep entry data. + */ +export interface SleepEntry { + deadline: number; + state: SleepState; +} + +/** + * Message entry data. + */ +export interface MessageEntry { + name: string; + data: unknown; +} + +/** + * Rollback checkpoint entry data. + */ +export interface RollbackCheckpointEntry { + name: string; +} + +/** + * Branch status for join/race entries. + */ +export interface BranchStatus { + status: BranchStatusType; + output?: unknown; + error?: string; +} + +/** + * Join entry data. + */ +export interface JoinEntry { + branches: Record; +} + +/** + * Race entry data. + */ +export interface RaceEntry { + winner: string | null; + branches: Record; +} + +/** + * Removed entry data - placeholder for removed steps in workflow migrations. + */ +export interface RemovedEntry { + originalType: EntryKindType; + originalName?: string; +} + +/** + * All possible entry kind types. + */ +export type EntryKindType = + | "step" + | "loop" + | "sleep" + | "message" + | "rollback_checkpoint" + | "join" + | "race" + | "removed"; + +/** + * Type-specific entry data. + */ +export type EntryKind = + | { type: "step"; data: StepEntry } + | { type: "loop"; data: LoopEntry } + | { type: "sleep"; data: SleepEntry } + | { type: "message"; data: MessageEntry } + | { type: "rollback_checkpoint"; data: RollbackCheckpointEntry } + | { type: "join"; data: JoinEntry } + | { type: "race"; data: RaceEntry } + | { type: "removed"; data: RemovedEntry }; + +/** + * An entry in the workflow history. + */ +export interface Entry { + id: string; + location: Location; + kind: EntryKind; + dirty: boolean; +} + +/** + * Metadata for an entry (stored separately, lazily loaded). + */ +export interface EntryMetadata { + status: EntryStatus; + error?: string; + attempts: number; + lastAttemptAt: number; + createdAt: number; + completedAt?: number; + rollbackCompletedAt?: number; + rollbackError?: string; + dirty: boolean; +} + +/** + * A message in the queue. + */ +export interface Message { + /** Unique message ID (used as KV key). */ + id: string; + name: string; + data: unknown; + sentAt: number; +} + +/** + * Workflow history - maps location keys to entries. + */ +export interface History { + entries: Map; +} + +/** + * Structured error information for workflow failures. + */ +export interface WorkflowError { + /** Error name/type (e.g., "TypeError", "CriticalError") */ + name: string; + /** Error message */ + message: string; + /** Stack trace if available */ + stack?: string; + /** Custom error properties (for structured errors) */ + metadata?: Record; +} + +/** + * Complete storage state for a workflow. + */ +export interface Storage { + nameRegistry: string[]; + flushedNameCount: number; + history: History; + entryMetadata: Map; + messages: Message[]; + output?: unknown; + state: WorkflowState; + flushedState?: WorkflowState; + error?: WorkflowError; + flushedError?: WorkflowError; + flushedOutput?: unknown; +} + +/** + * Context available to rollback handlers. + */ +export interface RollbackContextInterface { + readonly workflowId: string; + readonly abortSignal: AbortSignal; + isEvicted(): boolean; +} + +/** + * Configuration for a step. + */ +export interface StepConfig { + name: string; + run: () => Promise; + rollback?: (ctx: RollbackContextInterface, output: T) => Promise; + /** If true, step result is not persisted (use for idempotent operations). */ + ephemeral?: boolean; + /** Maximum number of retry attempts (default: 3). */ + maxRetries?: number; + /** Base delay in ms for exponential backoff (default: 100). */ + retryBackoffBase?: number; + /** Maximum delay in ms for exponential backoff (default: 30000). */ + retryBackoffMax?: number; + /** Timeout in ms for step execution (default: 30000). Set to 0 to disable. */ + timeout?: number; +} + +/** + * Result from a loop iteration. + */ +export type LoopResult = + | { continue: true; state: S } + | { break: true; value: T }; + +/** + * Configuration for a loop. + */ +export interface LoopConfig { + name: string; + state?: S; + run: (ctx: WorkflowContextInterface, state: S) => Promise>; + commitInterval?: number; + /** Trim loop history every N iterations. Defaults to commitInterval or 20. */ + historyEvery?: number; + /** Retain the last N iterations of history. Defaults to commitInterval or 20. */ + historyKeep?: number; +} + +/** + * Configuration for a branch in join/race. + */ +export interface BranchConfig { + run: (ctx: WorkflowContextInterface) => Promise; +} + +/** + * Extract the output type from a BranchConfig. + */ +export type BranchOutput = T extends BranchConfig ? O : never; + +/** + * The workflow context interface exposed to workflow functions. + */ +export interface WorkflowContextInterface { + readonly workflowId: string; + readonly abortSignal: AbortSignal; + + step(name: string, run: () => Promise): Promise; + step(config: StepConfig): Promise; + + loop( + name: string, + run: ( + ctx: WorkflowContextInterface, + ) => Promise>, + ): Promise; + loop(config: LoopConfig): Promise; + + sleep(name: string, durationMs: number): Promise; + sleepUntil(name: string, timestampMs: number): Promise; + + listen(name: string, messageName: string): Promise; + listenN(name: string, messageName: string, limit: number): Promise; + listenWithTimeout( + name: string, + messageName: string, + timeoutMs: number, + ): Promise; + listenUntil( + name: string, + messageName: string, + timestampMs: number, + ): Promise; + listenNWithTimeout( + name: string, + messageName: string, + limit: number, + timeoutMs: number, + ): Promise; + listenNUntil( + name: string, + messageName: string, + limit: number, + timestampMs: number, + ): Promise; + + rollbackCheckpoint(name: string): Promise; + + join>>( + name: string, + branches: T, + ): Promise<{ [K in keyof T]: BranchOutput }>; + + race( + name: string, + branches: Array<{ + name: string; + run: (ctx: WorkflowContextInterface) => Promise; + }>, + ): Promise<{ winner: string; value: T }>; + + removed(name: string, originalType: EntryKindType): Promise; + + isEvicted(): boolean; +} + +/** + * Workflow function type. + */ +export type WorkflowRunMode = "yield" | "live"; + +export interface RunWorkflowOptions { + mode?: WorkflowRunMode; +} + +export type WorkflowFunction = ( + ctx: WorkflowContextInterface, + input: TInput, +) => Promise; + +/** + * Result returned when a workflow run completes or yields. + */ +export interface WorkflowResult { + state: WorkflowState; + output?: TOutput; + sleepUntil?: number; + waitingForMessages?: string[]; +} + +/** + * Handle for managing a running workflow. + * + * Returned by `runWorkflow()`. The workflow starts executing immediately. + * Use `.result` to await completion, and other methods to interact with + * the running workflow. + */ +export interface WorkflowHandle { + readonly workflowId: string; + + /** + * Promise that resolves when the workflow completes or yields. + */ + readonly result: Promise>; + + /** + * Send a message to the workflow. + * The message is persisted and will be available on the next run. + * In live mode, this also wakes the workflow if it's waiting. + */ + message(name: string, data: unknown): Promise; + + /** + * Wake the workflow immediately by setting an alarm for now. + */ + wake(): Promise; + + /** + * Reset exhausted retries and schedule the workflow to run again. + */ + recover(): Promise; + + /** + * Request the workflow to stop gracefully. + * The workflow will throw EvictedError at its next yield point, + * flush its state, and resolve the result promise. + */ + evict(): void; + + /** + * Cancel the workflow permanently. + * Sets the workflow state to "cancelled" and clears any pending alarms. + * Unlike evict(), this marks the workflow as permanently stopped. + */ + cancel(): Promise; + + /** + * Get the workflow output if completed. + */ + getOutput(): Promise; + + /** + * Get the current workflow state. + */ + getState(): Promise; +} diff --git a/rivetkit-typescript/packages/workflow-engine/src/utils.ts b/rivetkit-typescript/packages/workflow-engine/src/utils.ts new file mode 100644 index 0000000000..16de15f87d --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/src/utils.ts @@ -0,0 +1,48 @@ +/** + * Sleep for a given number of milliseconds. + */ +export function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +const TIMEOUT_MAX = 2147483647; + +export type LongTimeoutHandle = { abort: () => void }; + +export function setLongTimeout( + listener: () => void, + after: number, +): LongTimeoutHandle { + let timeout: ReturnType | undefined; + + function start(remaining: number) { + if (remaining <= TIMEOUT_MAX) { + timeout = setTimeout(listener, remaining); + } else { + timeout = setTimeout(() => { + start(remaining - TIMEOUT_MAX); + }, TIMEOUT_MAX); + } + } + + start(after); + + return { + abort: () => { + if (timeout !== undefined) clearTimeout(timeout); + }, + }; +} + +/** + * Safely parse JSON with a meaningful error message. + */ +export function safeJsonParse(value: string, context: string): T { + try { + return JSON.parse(value) as T; + } catch (error) { + throw new Error( + `Failed to parse ${context}: ${error instanceof Error ? error.message : String(error)}`, + ); + } +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/driver-kv.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/driver-kv.test.ts new file mode 100644 index 0000000000..4b8ba9af8a --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/driver-kv.test.ts @@ -0,0 +1,120 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + buildMessageKey, + buildMessagePrefix, + buildWorkflowStateKey, + parseMessageKey, +} from "../src/keys.js"; +import { InMemoryDriver } from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +const encoder = new TextEncoder(); + +function encode(value: string): Uint8Array { + return encoder.encode(value); +} + +for (const mode of modes) { + describe( + `Workflow Engine Driver KV (${mode})`, + { sequential: true }, + () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should set and get values", async () => { + const key = encode("key-a"); + const value = encode("value-a"); + + await driver.set(key, value); + + const result = await driver.get(key); + expect(result).toEqual(value); + }); + + it("should return null for missing keys", async () => { + const result = await driver.get(encode("missing")); + expect(result).toBeNull(); + }); + + it("should overwrite existing keys", async () => { + const key = encode("key-b"); + await driver.set(key, encode("first")); + await driver.set(key, encode("second")); + + const result = await driver.get(key); + expect(result).toEqual(encode("second")); + }); + + it("should delete keys", async () => { + const key = encode("key-c"); + await driver.set(key, encode("value")); + await driver.delete(key); + + const result = await driver.get(key); + expect(result).toBeNull(); + }); + + it("should list keys by prefix", async () => { + await driver.set(buildMessageKey("a"), encode("one")); + await driver.set(buildMessageKey("b"), encode("two")); + await driver.set(buildWorkflowStateKey(), encode("state")); + + const entries = await driver.list(buildMessagePrefix()); + const ids = entries.map((entry) => parseMessageKey(entry.key)); + + expect(ids).toEqual(["a", "b"]); + }); + + it("should delete only keys with a prefix", async () => { + const messageKey = buildMessageKey("message"); + const stateKey = buildWorkflowStateKey(); + + await driver.set(messageKey, encode("message")); + await driver.set(stateKey, encode("state")); + + await driver.deletePrefix(buildMessagePrefix()); + + expect(await driver.get(messageKey)).toBeNull(); + expect(await driver.get(stateKey)).not.toBeNull(); + }); + + it("should list messages in sorted order", async () => { + await driver.set(buildMessageKey("b"), encode("two")); + await driver.set(buildMessageKey("a"), encode("one")); + + const entries = await driver.list(buildMessagePrefix()); + const ids = entries.map((entry) => parseMessageKey(entry.key)); + + expect(ids).toEqual(["a", "b"]); + }); + + it("should batch writes", async () => { + const keyA = encode("batch-a"); + const keyB = encode("batch-b"); + + await driver.batch([ + { key: keyA, value: encode("one") }, + { key: keyB, value: encode("two") }, + ]); + + expect(await driver.get(keyA)).toEqual(encode("one")); + expect(await driver.get(keyB)).toEqual(encode("two")); + }); + + it("should batch overwrite existing keys", async () => { + const key = encode("batch-c"); + await driver.set(key, encode("old")); + + await driver.batch([{ key, value: encode("new") }]); + + expect(await driver.get(key)).toEqual(encode("new")); + }); + }, + ); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/driver-scheduling.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/driver-scheduling.test.ts new file mode 100644 index 0000000000..a0296c5b78 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/driver-scheduling.test.ts @@ -0,0 +1,42 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { InMemoryDriver } from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe( + `Workflow Engine Driver Scheduling (${mode})`, + { sequential: true }, + () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should set and clear alarms", async () => { + const wakeAt = Date.now() + 1000; + + await driver.setAlarm("wf-1", wakeAt); + expect(driver.getAlarm("wf-1")).toBe(wakeAt); + + await driver.clearAlarm("wf-1"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + }); + + it("should return due alarms", async () => { + await driver.setAlarm("wf-due", Date.now() - 1); + await driver.setAlarm("wf-later", Date.now() + 1000); + + const due = driver.getDueAlarms(); + expect(due).toContain("wf-due"); + expect(due).not.toContain("wf-later"); + }); + + it("should expose worker poll interval", async () => { + expect(driver.workerPollInterval).toBeGreaterThan(0); + }); + }, + ); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/eviction-cancel.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/eviction-cancel.test.ts new file mode 100644 index 0000000000..e9f0444093 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/eviction-cancel.test.ts @@ -0,0 +1,82 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + EvictedError, + InMemoryDriver, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe( + `Workflow Engine Eviction and Cancellation (${mode})`, + { sequential: true }, + () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should surface eviction through the abort event", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await new Promise((resolve) => { + if (ctx.abortSignal.aborted) { + resolve(); + return; + } + ctx.abortSignal.addEventListener( + "abort", + () => resolve(), + { + once: true, + }, + ); + }); + return ctx.isEvicted(); + }; + + const handle = runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ); + handle.evict(); + + const result = await handle.result; + expect(result.state).toBe("completed"); + expect(result.output).toBe(true); + }); + + it("should cancel workflow and clear alarms", async () => { + const workflow = async (_ctx: WorkflowContextInterface) => { + return "done"; + }; + + const handle = runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ); + await driver.setAlarm("wf-1", Date.now() + 1000); + + await handle.cancel(); + + await expect(handle.result).rejects.toThrow(EvictedError); + expect(await handle.getState()).toBe("cancelled"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(EvictedError); + }); + }, + ); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts new file mode 100644 index 0000000000..dad3433dc5 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/handle.test.ts @@ -0,0 +1,86 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Handle (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should send messages via handle", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listen("wait", "message-name"); + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + + if (mode === "yield") { + await handle.result; + await handle.message("message-name", "payload"); + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("payload"); + return; + } + + await handle.message("message-name", "payload"); + const result = await handle.result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("payload"); + }); + + it("should set alarms with wake", async () => { + const workflow = async (_ctx: WorkflowContextInterface) => { + return "done"; + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + await handle.wake(); + + const alarm = driver.getAlarm("wf-1"); + if (mode === "yield") { + expect(alarm).toBeDefined(); + expect(alarm).toBeLessThanOrEqual(Date.now()); + } else { + expect(alarm).toBeUndefined(); + } + }); + + it("should read output and state", async () => { + const workflow = async (_ctx: WorkflowContextInterface) => { + return "finished"; + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + const result = await handle.result; + + expect(result.state).toBe("completed"); + expect(await handle.getOutput()).toBe("finished"); + expect(await handle.getState()).toBe("completed"); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/join.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/join.test.ts new file mode 100644 index 0000000000..33e223de0f --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/join.test.ts @@ -0,0 +1,149 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + JoinError, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Join (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should execute branches in parallel", async () => { + const order: string[] = []; + + const workflow = async (ctx: WorkflowContextInterface) => { + const results = await ctx.join("parallel", { + a: { + run: async (ctx) => { + order.push("a-start"); + const val = await ctx.step("step-a", async () => 1); + order.push("a-end"); + return val; + }, + }, + b: { + run: async (ctx) => { + order.push("b-start"); + const val = await ctx.step("step-b", async () => 2); + order.push("b-end"); + return val; + }, + }, + }); + + return results.a + results.b; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe(3); + expect(order.indexOf("a-start")).toBeLessThan( + order.indexOf("b-end"), + ); + expect(order.indexOf("b-start")).toBeLessThan( + order.indexOf("a-end"), + ); + }); + + it("should wait for all branches even on error", async () => { + let bCompleted = false; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.join("parallel", { + a: { + run: async () => { + throw new Error("A failed"); + }, + }, + b: { + run: async (ctx) => { + await ctx.step("step-b", async () => { + bCompleted = true; + return "b"; + }); + return "b"; + }, + }, + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(); + + expect(bCompleted).toBe(true); + }); + + it("should surface join errors per branch", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.join("parallel", { + good: { + run: async () => "ok", + }, + bad: { + run: async () => { + throw new Error("boom"); + }, + }, + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(JoinError); + + try { + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + } catch (error) { + const joinError = error as JoinError; + expect(Object.keys(joinError.errors)).toEqual(["bad"]); + } + }); + + it("should replay join results", async () => { + let callCount = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + const result = await ctx.join("replay", { + one: { + run: async () => { + callCount += 1; + return "one"; + }, + }, + two: { + run: async () => "two", + }, + }); + + return result.one; + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + + expect(callCount).toBe(1); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/loops.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/loops.test.ts new file mode 100644 index 0000000000..394b3e9a0b --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/loops.test.ts @@ -0,0 +1,174 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + Loop, + loadStorage, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +function isLoopIteration(segment: unknown): segment is { iteration: number } { + return ( + typeof segment === "object" && + segment !== null && + "iteration" in segment + ); +} + +for (const mode of modes) { + describe(`Workflow Engine Loops (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should execute a simple loop", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.loop({ + name: "count-loop", + state: { count: 0 }, + run: async (ctx, state) => { + if (state.count >= 3) { + return Loop.break(state.count); + } + await ctx.step(`step-${state.count}`, async () => {}); + return Loop.continue({ count: state.count + 1 }); + }, + }); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe(3); + }); + + it("should run a stateless loop", async () => { + let iteration = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.loop("stateless", async () => { + iteration += 1; + if (iteration >= 3) { + return Loop.break("done"); + } + return Loop.continue(undefined); + }); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + }); + + it("should resume loop from saved state", async () => { + let iteration = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.loop({ + name: "resume-loop", + state: { count: 0 }, + commitInterval: 2, + run: async (_ctx, state) => { + iteration++; + + if (state.count >= 5) { + return Loop.break(state.count); + } + + if (state.count === 2 && iteration === 3) { + throw new Error("Simulated crash"); + } + + return Loop.continue({ count: state.count + 1 }); + }, + }); + }; + + try { + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + } catch {} + + iteration = 0; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe(5); + }); + + it("should forget old iterations on history window", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.loop({ + name: "cleanup-loop", + state: { count: 0 }, + commitInterval: 2, + historyEvery: 2, + historyKeep: 2, + run: async (ctx, state) => { + await ctx.step(`step-${state.count}`, async () => {}); + if (state.count >= 4) { + return Loop.break(state.count); + } + return Loop.continue({ count: state.count + 1 }); + }, + }); + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + + const storage = await loadStorage(driver); + const iterations = [...storage.history.entries.values()] + .flatMap((entry) => entry.location) + .flatMap((segment) => + isLoopIteration(segment) ? [segment] : [], + ) + .map((segment) => segment.iteration); + + const minIteration = Math.min(...iterations); + expect(minIteration).toBeGreaterThanOrEqual(2); + }); + + it("should propagate loop errors", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.loop({ + name: "error-loop", + state: { count: 0 }, + run: async () => { + throw new Error("loop failure"); + }, + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow("loop failure"); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts new file mode 100644 index 0000000000..1a5538a6cb --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/messages.test.ts @@ -0,0 +1,316 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + buildMessageKey, + generateId, + InMemoryDriver, + runWorkflow, + serializeMessage, + type WorkflowContextInterface, +} from "../src/testing.js"; + +function buildMessagePayload(name: string, data: string, id = generateId()) { + return { + id, + name, + data, + sentAt: Date.now(), + }; +} + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Messages (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should wait for messages", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const message = await ctx.listen( + "wait-message", + "my-message", + ); + return message; + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + + if (mode === "yield") { + const result1 = await handle.result; + expect(result1.state).toBe("sleeping"); + expect(result1.waitingForMessages).toContain("my-message"); + return; + } + + await handle.message("my-message", "payload"); + const result = await handle.result; + expect(result.state).toBe("completed"); + expect(result.output).toBe("payload"); + }); + + it("should consume pending messages", async () => { + const messageId = generateId(); + await driver.set( + buildMessageKey(messageId), + serializeMessage( + buildMessagePayload("my-message", "hello", messageId), + ), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + const message = await ctx.listen( + "wait-message", + "my-message", + ); + return message; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("hello"); + }); + + it("should collect multiple messages with listenN", async () => { + await driver.set( + buildMessageKey("1"), + serializeMessage(buildMessagePayload("batch", "a", "1")), + ); + await driver.set( + buildMessageKey("2"), + serializeMessage(buildMessagePayload("batch", "b", "2")), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listenN("batch-wait", "batch", 2); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toEqual(["a", "b"]); + }); + + it("should time out listenWithTimeout", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listenWithTimeout( + "timeout", + "missing", + 50, + ); + }; + + if (mode === "yield") { + const result1 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result1.state).toBe("sleeping"); + + await new Promise((r) => setTimeout(r, 80)); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result2.state).toBe("completed"); + expect(result2.output).toBeNull(); + return; + } + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result.state).toBe("completed"); + expect(result.output).toBeNull(); + }); + + it("should return a message before listenUntil deadline", async () => { + const messageId = generateId(); + await driver.set( + buildMessageKey(messageId), + serializeMessage( + buildMessagePayload("deadline", "data", messageId), + ), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listenUntil( + "deadline", + "deadline", + Date.now() + 1000, + ); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("data"); + }); + + it("should wait for listenNWithTimeout messages", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listenNWithTimeout( + "batch", + "batch", + 2, + 5000, + ); + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + + if (mode === "yield") { + const result1 = await handle.result; + expect(result1.state).toBe("sleeping"); + + await handle.message("batch", "first"); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result2.state).toBe("completed"); + expect(result2.output).toEqual(["first"]); + return; + } + + await handle.message("batch", "first"); + const result = await handle.result; + expect(result.state).toBe("completed"); + expect(result.output).toEqual(["first"]); + }); + + it("should respect limits and FIFO ordering in listenNUntil", async () => { + await driver.set( + buildMessageKey("1"), + serializeMessage(buildMessagePayload("fifo", "first", "1")), + ); + await driver.set( + buildMessageKey("2"), + serializeMessage(buildMessagePayload("fifo", "second", "2")), + ); + await driver.set( + buildMessageKey("3"), + serializeMessage(buildMessagePayload("fifo", "third", "3")), + ); + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.listenNUntil( + "fifo", + "fifo", + 2, + Date.now() + 1000, + ); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toEqual(["first", "second"]); + }); + + it("should not see messages sent during execution until next run", async () => { + let resolveGate: (() => void) | null = null; + let resolveStarted: (() => void) | null = null; + const gate = new Promise((resolve) => { + resolveGate = () => resolve(); + }); + const started = new Promise((resolve) => { + resolveStarted = () => resolve(); + }); + + const workflow = async (ctx: WorkflowContextInterface) => { + resolveStarted?.(); + await ctx.step("gate", async () => { + await gate; + return "ready"; + }); + + return await ctx.listen("wait", "mid"); + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + await started; + await handle.message("mid", "value"); + if (resolveGate) { + (resolveGate as () => void)(); + } + + if (mode === "yield") { + const result1 = await handle.result; + expect(result1.state).toBe("sleeping"); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result2.state).toBe("completed"); + expect(result2.output).toBe("value"); + return; + } + + const result = await handle.result; + expect(result.state).toBe("completed"); + expect(result.output).toBe("value"); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/race.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/race.test.ts new file mode 100644 index 0000000000..482d454291 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/race.test.ts @@ -0,0 +1,155 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + RaceError, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Race (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should return first completed branch", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.race("race", [ + { + name: "fast", + run: async (ctx) => { + return await ctx.step( + "fast-step", + async () => "fast", + ); + }, + }, + { + name: "slow", + run: async (ctx) => { + return await ctx.step( + "slow-step", + async () => "slow", + ); + }, + }, + ]); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output?.winner).toBe("fast"); + expect(result.output?.value).toBe("fast"); + }); + + it("should cancel losing branches", async () => { + let aborted = false; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.race("cancel", [ + { + name: "fast", + run: async () => "winner", + }, + { + name: "slow", + run: async (ctx) => { + await new Promise((resolve) => { + if (ctx.abortSignal.aborted) { + aborted = true; + resolve(); + return; + } + ctx.abortSignal.addEventListener( + "abort", + () => { + aborted = true; + resolve(); + }, + { once: true }, + ); + }); + return "aborted"; + }, + }, + ]); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(aborted).toBe(true); + }); + + it("should surface errors when all branches fail", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.race("fail", [ + { + name: "one", + run: async () => { + throw new Error("fail one"); + }, + }, + { + name: "two", + run: async () => { + throw new Error("fail two"); + }, + }, + ]); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(RaceError); + }); + + it("should replay race winner", async () => { + let runs = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + const result = await ctx.race("replay", [ + { + name: "first", + run: async () => { + runs += 1; + return "winner"; + }, + }, + { + name: "second", + run: async () => "loser", + }, + ]); + + return result.value; + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + + expect(runs).toBe(1); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts new file mode 100644 index 0000000000..a6ec0165c7 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/removals.test.ts @@ -0,0 +1,93 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + buildMessageKey, + generateId, + InMemoryDriver, + Loop, + runWorkflow, + serializeMessage, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe( + `Workflow Engine Removed Entries (${mode})`, + { sequential: true }, + () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should skip removed entries of all kinds", async () => { + const messageId = generateId(); + await driver.set( + buildMessageKey(messageId), + serializeMessage({ + id: messageId, + name: "old-message", + data: "message-data", + sentAt: Date.now(), + }), + ); + + const workflow1 = async (ctx: WorkflowContextInterface) => { + await ctx.step("old-step", async () => "old"); + await ctx.loop({ + name: "old-loop", + state: { count: 0 }, + run: async (_ctx, state) => { + if (state.count >= 1) { + return Loop.break("done"); + } + return Loop.continue({ count: state.count + 1 }); + }, + }); + await ctx.sleep("old-sleep", 0); + await ctx.listen("old-listen", "old-message"); + await ctx.join("old-join", { + branch: { + run: async () => "ok", + }, + }); + await ctx.race("old-race", [ + { + name: "fast", + run: async () => "fast", + }, + ]); + return "done"; + }; + + await runWorkflow("wf-1", workflow1, undefined, driver, { + mode, + }).result; + + const workflow2 = async (ctx: WorkflowContextInterface) => { + await ctx.removed("old-step", "step"); + await ctx.removed("old-loop", "loop"); + await ctx.removed("old-sleep", "sleep"); + await ctx.removed("old-listen", "message"); + await ctx.removed("old-join", "join"); + await ctx.removed("old-race", "race"); + return "updated"; + }; + + const result = await runWorkflow( + "wf-1", + workflow2, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("updated"); + }); + }, + ); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/rollback.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/rollback.test.ts new file mode 100644 index 0000000000..ffad1b7538 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/rollback.test.ts @@ -0,0 +1,135 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + RollbackCheckpointError, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Rollback (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should execute rollback steps in reverse order", async () => { + const rollbacks: string[] = []; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.rollbackCheckpoint("checkpoint"); + await ctx.step({ + name: "first", + run: async () => "one", + rollback: async (_ctx, output) => { + rollbacks.push(`first:${output}`); + }, + }); + await ctx.step({ + name: "second", + run: async () => "two", + rollback: async (_ctx, output) => { + rollbacks.push(`second:${output}`); + }, + }); + throw new Error("boom"); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow("boom"); + + expect(rollbacks).toEqual(["second:two", "first:one"]); + }); + + it("should error if rollback checkpoint missing", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.step({ + name: "missing-checkpoint", + run: async () => "value", + rollback: async () => { + return; + }, + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(RollbackCheckpointError); + }); + + it("should resume rollback after eviction", async () => { + const rollbacks: string[] = []; + let unblockRollback: (() => void) | undefined; + let startRollback: (() => void) | undefined; + + const rollbackGate = new Promise((resolve) => { + unblockRollback = () => resolve(); + }); + const rollbackStarted = new Promise((resolve) => { + startRollback = () => resolve(); + }); + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.rollbackCheckpoint("checkpoint"); + await ctx.step({ + name: "first", + run: async () => "one", + rollback: async () => { + rollbacks.push("first"); + }, + }); + await ctx.step({ + name: "second", + run: async () => "two", + rollback: async (rollbackCtx) => { + if (startRollback) { + startRollback(); + } + await rollbackGate; + if (rollbackCtx.abortSignal.aborted) { + return; + } + rollbacks.push("second"); + }, + }); + await ctx.step({ + name: "third", + run: async () => "three", + rollback: async () => { + rollbacks.push("third"); + }, + }); + throw new Error("boom"); + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + + await rollbackStarted; + handle.evict(); + + const result = await handle.result; + expect(result.state).toBe("rolling_back"); + expect(rollbacks).toEqual(["third"]); + + if (unblockRollback) { + unblockRollback(); + } + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow("boom"); + + expect(rollbacks).toEqual(["third", "second", "first"]); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/sleep.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/sleep.test.ts new file mode 100644 index 0000000000..0087474264 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/sleep.test.ts @@ -0,0 +1,206 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Sleep (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should yield on long sleep", async () => { + const durationMs = mode === "live" ? 150 : 10000; + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleep("my-sleep", durationMs); + return "done"; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + if (mode === "yield") { + expect(result.state).toBe("sleeping"); + expect(result.sleepUntil).toBeDefined(); + expect(result.sleepUntil).toBeGreaterThan(Date.now()); + return; + } + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + }); + + it("should complete short sleep in memory", async () => { + driver.workerPollInterval = 1000; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleep("short-sleep", 10); + return "done"; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + }); + + it("should resume after sleep deadline", async () => { + driver.workerPollInterval = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleep("my-sleep", 20); + return "done"; + }; + + if (mode === "yield") { + const result1 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result1.state).toBe("sleeping"); + + await new Promise((r) => setTimeout(r, 30)); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result2.state).toBe("completed"); + expect(result2.output).toBe("done"); + return; + } + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + }); + + it("should complete sleepUntil with past timestamp", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleepUntil("past", Date.now() - 1); + return "done"; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + }); + + it("should keep short sleeps in memory near poll interval", async () => { + driver.workerPollInterval = 50; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleep("near-poll", 25); + return "done"; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { + mode, + }, + ).result; + + expect(result.state).toBe("completed"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + }); + + it("should schedule and clear alarms for long sleep", async () => { + driver.workerPollInterval = 1; + const durationMs = mode === "live" ? 200 : 20; + + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.sleep("alarm-sleep", durationMs); + return "done"; + }; + + const handle = runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }); + + if (mode === "yield") { + const result1 = await handle.result; + expect(result1.state).toBe("sleeping"); + expect(driver.getAlarm("wf-1")).toBe(result1.sleepUntil); + + await new Promise((r) => setTimeout(r, 30)); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result2.state).toBe("completed"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + return; + } + + const alarm = await new Promise((resolve) => { + const start = Date.now(); + const check = () => { + const value = driver.getAlarm("wf-1"); + if (value !== undefined || Date.now() - start > 50) { + resolve(value); + return; + } + setTimeout(check, 1); + }; + check(); + }); + expect(alarm).toBeDefined(); + + const result = await handle.result; + expect(result.state).toBe("completed"); + expect(driver.getAlarm("wf-1")).toBeUndefined(); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/steps.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/steps.test.ts new file mode 100644 index 0000000000..fef0e41336 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/steps.test.ts @@ -0,0 +1,372 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + CriticalError, + EntryInProgressError, + HistoryDivergedError, + InMemoryDriver, + RollbackError, + runWorkflow, + StepExhaustedError, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +class CountingDriver extends InMemoryDriver { + batchCalls = 0; + + async batch( + writes: { key: Uint8Array; value: Uint8Array }[], + ): Promise { + this.batchCalls += 1; + await super.batch(writes); + } +} + +for (const mode of modes) { + describe(`Workflow Engine Steps (${mode})`, { sequential: true }, () => { + let driver: CountingDriver; + + beforeEach(() => { + driver = new CountingDriver(); + driver.latency = 0; + }); + + it("should execute a simple step", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const result = await ctx.step("my-step", async () => { + return "hello world"; + }); + return result; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("hello world"); + }); + + it("should replay step on restart", async () => { + let callCount = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + const result = await ctx.step("my-step", async () => { + callCount++; + return "hello"; + }); + return result; + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + expect(callCount).toBe(1); + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + expect(callCount).toBe(1); + }); + + it("should execute multiple steps in sequence", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const a = await ctx.step("step-a", async () => 1); + const b = await ctx.step("step-b", async () => 2); + const c = await ctx.step("step-c", async () => 3); + return a + b + c; + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe(6); + }); + + it("should retry failed steps", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step({ + name: "flaky-step", + maxRetries: 3, + retryBackoffBase: 1, + retryBackoffMax: 10, + run: async () => { + attempts++; + if (attempts < 3) { + throw new Error("Transient failure"); + } + return "success"; + }, + }); + }; + + try { + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + } catch {} + + try { + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + } catch {} + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("success"); + expect(attempts).toBe(3); + }); + + it("should yield during backoff retries", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step({ + name: "always-fails", + maxRetries: 3, + retryBackoffBase: 50, + retryBackoffMax: 100, + run: async () => { + attempts++; + throw new Error("Failure"); + }, + }); + }; + + if (mode === "yield") { + const result1 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result1.state).toBe("sleeping"); + expect(result1.sleepUntil).toBeDefined(); + + const result2 = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + expect(result2.state).toBe("sleeping"); + expect(result2.sleepUntil).toBeDefined(); + expect(result2.sleepUntil).toBeGreaterThan(Date.now()); + expect(driver.getAlarm("wf-1")).toBe(result2.sleepUntil); + expect(attempts).toBe(1); + return; + } + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(StepExhaustedError); + expect(attempts).toBe(3); + }); + + it("should not retry CriticalError", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step("critical-step", async () => { + attempts++; + throw new CriticalError("Unrecoverable"); + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(CriticalError); + + expect(attempts).toBe(1); + }); + + it("should not retry RollbackError", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step("rollback-step", async () => { + attempts++; + throw new RollbackError("Rollback now"); + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(RollbackError); + + expect(attempts).toBe(1); + }); + + it("should exhaust retries", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step({ + name: "always-fails", + maxRetries: 2, + retryBackoffBase: 1, + run: async () => { + attempts++; + throw new Error("Always fails"); + }, + }); + }; + + if (mode === "yield") { + for (let i = 0; i < 3; i++) { + try { + await runWorkflow("wf-1", workflow, undefined, driver, { + mode, + }).result; + } catch {} + } + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(StepExhaustedError); + return; + } + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(StepExhaustedError); + }); + + it("should recover exhausted retries", async () => { + let attempts = 0; + + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step({ + name: "recoverable", + maxRetries: 1, + retryBackoffBase: 1, + run: async () => { + attempts++; + throw new Error("Always fails"); + }, + }); + }; + + const runOnce = () => + runWorkflow("wf-1", workflow, undefined, driver, { mode }); + + if (mode === "yield") { + await runOnce().result; + } + + const exhaustedHandle = runOnce(); + await expect(exhaustedHandle.result).rejects.toThrow( + StepExhaustedError, + ); + + const attemptsAfterExhaust = attempts; + await exhaustedHandle.recover(); + + if (mode === "yield") { + await runOnce().result; + } else { + await expect(runOnce().result).rejects.toThrow( + StepExhaustedError, + ); + } + + expect(attempts).toBeGreaterThan(attemptsAfterExhaust); + }); + + it("should fail steps that exceed timeout", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step({ + name: "timeout-step", + timeout: 5, + run: async () => { + await new Promise((resolve) => setTimeout(resolve, 25)); + return "late"; + }, + }); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(CriticalError); + }); + + it("should fail when a step is not awaited", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + const first = ctx.step("step-a", async () => "a"); + await ctx.step("step-b", async () => "b"); + return await first; + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(EntryInProgressError); + }); + + it("should reject duplicate entry names", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.step("dup", async () => "first"); + await ctx.step("dup", async () => "second"); + return "done"; + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow(HistoryDivergedError); + }); + + it("should batch ephemeral steps until a durable flush", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + await ctx.step({ + name: "ephemeral-a", + ephemeral: true, + run: async () => "a", + }); + await ctx.step({ + name: "ephemeral-b", + ephemeral: true, + run: async () => "b", + }); + return await ctx.step("durable", async () => "done"); + }; + + const result = await runWorkflow( + "wf-1", + workflow, + undefined, + driver, + { mode }, + ).result; + + expect(result.state).toBe("completed"); + expect(result.output).toBe("done"); + expect(driver.batchCalls).toBe(2); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tests/storage.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/storage.test.ts new file mode 100644 index 0000000000..280582bbd0 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/storage.test.ts @@ -0,0 +1,66 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { + InMemoryDriver, + loadMetadata, + loadStorage, + runWorkflow, + type WorkflowContextInterface, +} from "../src/testing.js"; + +const modes = ["yield", "live"] as const; + +for (const mode of modes) { + describe(`Workflow Engine Storage (${mode})`, { sequential: true }, () => { + let driver: InMemoryDriver; + + beforeEach(() => { + driver = new InMemoryDriver(); + driver.latency = 0; + }); + + it("should persist workflow output and state", async () => { + const workflow = async (_ctx: WorkflowContextInterface) => { + return "value"; + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + + const storage = await loadStorage(driver); + expect(storage.state).toBe("completed"); + expect(storage.output).toBe("value"); + }); + + it("should persist workflow errors", async () => { + const workflow = async (_ctx: WorkflowContextInterface) => { + throw new Error("boom"); + }; + + await expect( + runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result, + ).rejects.toThrow("boom"); + + const storage = await loadStorage(driver); + expect(storage.state).toBe("failed"); + expect(storage.error?.message).toBe("boom"); + }); + + it("should persist entry metadata and names", async () => { + const workflow = async (ctx: WorkflowContextInterface) => { + return await ctx.step("named-step", async () => "ok"); + }; + + await runWorkflow("wf-1", workflow, undefined, driver, { mode }) + .result; + + const storage = await loadStorage(driver); + expect(storage.nameRegistry).toContain("named-step"); + + const entry = [...storage.history.entries.values()][0]; + const metadata = await loadMetadata(storage, driver, entry.id); + expect(metadata.status).toBe("completed"); + expect(metadata.attempts).toBe(1); + }); + }); +} diff --git a/rivetkit-typescript/packages/workflow-engine/tsconfig.json b/rivetkit-typescript/packages/workflow-engine/tsconfig.json new file mode 100644 index 0000000000..2482b0d391 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../../tsconfig.base.json", + "compilerOptions": { + "types": ["node"], + "outDir": "./dist" + }, + "include": ["src/**/*", "schemas/**/*", "dist/schemas/**/*", "tests/**/*"] +} diff --git a/rivetkit-typescript/packages/workflow-engine/tsup.config.ts b/rivetkit-typescript/packages/workflow-engine/tsup.config.ts new file mode 100644 index 0000000000..d8652c0151 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tsup.config.ts @@ -0,0 +1,7 @@ +import { defineConfig } from "tsup"; +import defaultConfig from "../../../tsup.base.ts"; + +export default defineConfig({ + ...defaultConfig, + outDir: "dist/tsup/", +}); diff --git a/rivetkit-typescript/packages/workflow-engine/vitest.config.ts b/rivetkit-typescript/packages/workflow-engine/vitest.config.ts new file mode 100644 index 0000000000..4afc2b39dd --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/vitest.config.ts @@ -0,0 +1,6 @@ +import { defineConfig } from "vitest/config"; +import defaultConfig from "../../../vitest.base.ts"; + +export default defineConfig({ + ...defaultConfig, +});