yieldstar 0.5.0 → 0.6.1

package/src/internal/step-runner.ts

@@ -1,7 +1,15 @@
 import {
+  assertSynchronousClaim,
+  assertSynchronousUpdater,
   cloneStoreState,
   defineStore,
   type Draft,
+  type DurableObject,
+  type DurableStoreSchema,
+  type IfDurable,
+  type ExecutionStartRequest,
+  type ExecutionStartResult,
+  InfrastructureError,
   type StandardSchemaV1,
   StepResponse,
   StepKey,
@@ -10,7 +18,6 @@ import {
   StepDelay,
   StepCacheCheck,
   StepStoreWait,
-  type StoreClient,
   type StoreDefinition,
   type StoreDeleteFromResult,
   type StoreKey,
@@ -20,11 +27,32 @@ import {
   type StoreTakeResult,
   type StoreUpdateFromResult,
   type StoreUpdateResult,
+  type StoreWaiter,
   trackStoreSelector,
+  type Durable,
+  type StepSchemaVerdict,
+  durableValueCodec,
+  StepEffectWait,
+  validateEffectRequest,
+  WorkflowEffectError,
+  WorkflowStartError,
+  WorkflowValueError,
 } from "@yieldstar/core";
 import { RetryableError } from "../exports/errors";
 import { getCallSiteHash } from "./utils";
+import {
+  StepRuntimeCall,
+  workflowCallback,
+  type RuntimeCallResult,
+} from "./runtime-call";
 import type { WorkflowEvent } from "@yieldstar/core";
+import {
+  StepEffectCall,
+  type EffectCallOutcome,
+  type EffectLookup,
+} from "./effect-call";
+
+export type { EffectLookup };
 
 /**
  * @description A library of step generators, each of which:
@@ -43,7 +71,7 @@ export type WorkflowStore<T> = {
   get(key?: string): AsyncGenerator<StepResponse, StoreSnapshot<T>>;
   select<R>(
     key: string,
-    selector: StoreSelector<T, R>
+    selector: StoreSelector<T, R> & DurableSelector<T, R>
   ): AsyncGenerator<StepResponse, R>;
   update(
     key: string,
@@ -59,158 +87,318 @@ export type WorkflowStore<T> = {
     snapshot: StoreSnapshot<T>
   ): AsyncGenerator<StepResponse, StoreDeleteFromResult>;
   when<R>(
-    selector: StoreSelector<T, R | undefined | null | false>
+    selector: StoreSelector<T, R | undefined | null | false> &
+      DurableMatchSelector<T, R>
   ): AsyncGenerator<StepResponse, NonNullable<R>>;
   when<R>(
     key: string,
-    selector: StoreSelector<T, R | undefined | null | false>
+    selector: StoreSelector<T, R | undefined | null | false> &
+      DurableMatchSelector<T, R>
   ): AsyncGenerator<StepResponse, NonNullable<R>>;
   take<R>(
     key: string,
-    selector: StoreSelector<T, R | undefined | null | false>,
+    selector: StoreSelector<T, R | undefined | null | false> &
+      DurableMatchSelector<T, R>,
     claim: (draft: Draft<T>, selected: NonNullable<R>) => void
   ): AsyncGenerator<StepResponse, NonNullable<R>>;
 };
 
+/** A store selector whose result, which the step records, is durable. */
+type DurableSelector<T, R> = IfDurable<R, StoreSelector<T, Durable<R>>>;
+
+/**
+ * A `when` or `take` selector whose match is durable. `false` is excluded
+ * before the check so that `R` is inferred as it would be without it.
+ */
+type DurableMatchSelector<T, R> = IfDurable<
+  Exclude<R, false>,
+  StoreSelector<T, Durable<Exclude<R, false>> | undefined | null | false>
+>;
+
 export type StepRunner = ReturnType<typeof createStepRunner>;
 
 export function createStepRunner(params: {
   event: WorkflowEvent;
-  storeClient: StoreClient;
+  clock: () => number;
 }) {
-  const { event, storeClient } = params;
-
+  const { event, clock } = params;
+
+  /**
+   * Returns a handle to the store `definition.name` + `id`, where `id`
+   * defaults to this execution's ID. Obtaining a handle reads, creates and
+   * records nothing, so any caller may obtain the same handle again. Each
+   * operation on a missing store first creates it from `definition.initial`.
+   */
   function store<Schema extends StandardSchemaV1>(
-    definition: StoreDefinition<Schema>,
-    params?: {
-      id?: string;
-      initial?:
-        | StoreState<Schema>
-        | (() => StoreState<Schema> | Promise<StoreState<Schema>>);
-    }
-  ): AsyncGenerator<StepResponse, WorkflowStore<StoreState<Schema>>> {
-    const id = params?.id ?? event.executionId;
-    const key = `store:${definition.name}:${id}`;
-
-    return storeStep({
-      key,
+    definition: StoreDefinition<Schema> & {
+      schema: DurableStoreSchema<Schema>;
+    },
+    params?: { id?: string }
+  ): WorkflowStore<StoreState<Schema>> {
+    return createWorkflowStore({
       definition,
-      id,
+      id: params?.id ?? event.executionId,
       event,
-      storeClient,
-      initial: params?.initial,
     });
   }
 
-  return { run, delay, poll, store };
+  /**
+   * Starts another workflow as a durable step and returns its execution ID.
+   * The start is identified by this execution's ID and `stepKey`, so every
+   * replay of the step names the same target. A rejected start throws a
+   * `WorkflowStartError`.
+   */
+  function start<Params extends Record<string, any>>(
+    stepKey: string,
+    target: { workflowId: string; params?: Params & Durable<Params> }
+  ): AsyncGenerator<
+    StepResponse,
+    { executionId: string },
+    StepResult | StepError
+  > {
+    if (stepKey === "") {
+      throw new Error("step.start requires a non-empty step key");
+    }
+    const request: ExecutionStartRequest = {
+      requestId: { executionId: event.executionId, stepKey },
+      workflowId: target.workflowId,
+    };
+    // A missing params and an empty object are different requests.
+    if (target.params !== undefined) {
+      request.params = target.params as DurableObject;
+    }
+    return startStep(stepKey, request);
+  }
+
+  /**
+   * Runs an external operation identified by `descriptor.effectId`, at most
+   * once in effect: after an uncertain attempt, `lookup` decides whether the
+   * operation already happened before `perform` runs again (RFC 005).
+   */
+  function effect<T, Input extends Record<string, any> = DurableObject>(
+    stepKey: string,
+    descriptor: EffectDescriptor<T, Input> & DurableEffectResult<T>
+  ): AsyncGenerator<StepResponse, T, StepResult | StepError> {
+    if (typeof stepKey !== "string" || stepKey === "") {
+      throw new Error("step.effect requires a non-empty step key");
+    }
+    return effectStep(stepKey, descriptor) as AsyncGenerator<
+      StepResponse,
+      T,
+      StepResult | StepError
+    >;
+  }
+
+  return { run, delay: createDelay(clock), poll, store, start, effect };
 }
 
-async function* storeStep<Schema extends StandardSchemaV1>(params: {
-  key: string;
-  definition: StoreDefinition<Schema>;
-  id: string;
-  event: WorkflowEvent;
-  storeClient: StoreClient;
-  initial?:
-    | StoreState<Schema>
-    | (() => StoreState<Schema> | Promise<StoreState<Schema>>);
-}): AsyncGenerator<
-  StepResponse,
-  WorkflowStore<StoreState<Schema>>,
-  StepResult | StepError
-> {
-  const { key, definition, id, event, storeClient, initial } = params;
+/**
+ * Describes one external operation for `step.effect`. `effectId` names the
+ * operation in the external system, so every caller that means it, in any
+ * execution or runtime, uses the same ID, and no caller reuses it for another
+ * operation. The external system must guarantee that repeated `perform` calls
+ * for one ID make one operation with one result, that `not-found` proves no
+ * call exists or can still finish, and that `settled` never changes.
+ */
+export type EffectDescriptor<T, Input = DurableObject> = {
+  effectId: string;
+  /** Names the version of the behaviour of `input`, `lookup` and `perform`. */
+  operation: string;
+  input: Input & Durable<Input>;
+  /** Milliseconds before an uncertain outcome is looked up again. */
+  retryIn: number;
+  /** Reports what the external system knows about `effectId`. */
+  lookup(effectId: string): Promise<EffectLookup<T>>;
+  /** Performs the operation under `effectId`. */
+  perform(effectId: string): Promise<T>;
+};
 
+/**
+ * Checks the effect's result type without taking part in its inference. A
+ * `perform` that resolves with nothing settles the effect with `undefined`.
+ */
+type DurableEffectResult<T> = [T] extends [void]
+  ? unknown
+  : IfDurable<T, { perform(effectId: string): Promise<Durable<T>> }>;
+
+async function* effectStep(
+  key: string,
+  descriptor: EffectDescriptor<unknown, unknown>
+): AsyncGenerator<StepResponse, unknown, any> {
   yield new StepKey(key);
 
   const cached = yield new StepCacheCheck();
 
-  if (cached) {
-    yield cached;
-    if (cached instanceof StepError) {
-      throw cached.err;
+  if (cached) return yield* replay(cached);
+
+  const { effectId, operation, input, retryIn, lookup, perform } =
+    (descriptor ?? {}) as Partial<EffectDescriptor<unknown, unknown>>;
+  const invalid = invalidEffectRequest({
+    effectId,
+    operation,
+    input,
+    retryIn,
+    lookup,
+    perform,
+  });
+  if (invalid) {
+    const err = new WorkflowEffectError(
+      "INVALID_EFFECT_REQUEST",
+      `step.effect "${key}": ${invalid}`
+    );
+    yield new StepError(err);
+    // unreachable – consumer calls throw() on the generator first
+    throw err;
+  }
+
+  const outcome: EffectCallOutcome = yield new StepEffectCall({
+    stepKey: key,
+    effectId: effectId!,
+    operation: operation!,
+    input: input as DurableObject,
+    retryIn: retryIn!,
+    lookup: (id) => lookup!.call(descriptor, id),
+    perform: (id) => perform!.call(descriptor, id),
+  });
+
+  switch (outcome.status) {
+    case "settled":
+      return yield* record(outcome.result);
+    case "rejected": {
+      const err = new WorkflowEffectError(
+        outcome.code,
+        `Effect "${effectId}" was begun with a different operation or input`
+      );
+      yield new StepError(err);
+      throw err;
     }
-    return createWorkflowStore({
-      definition,
-      id,
-      event,
-      storeClient,
-    });
+    case "suspended":
+      yield new StepEffectWait(outcome.wakeAt);
+      throw new Error("Effect wait yielded control unexpectedly");
   }
+}
 
-  try {
-    await storeClient.getOrCreateStore({
-      definition,
-      id,
-      initial,
-    });
-    yield new StepResult({ storeName: definition.name, storeId: id });
-    return createWorkflowStore({
-      definition,
-      id,
-      event,
-      storeClient,
+function invalidEffectRequest(request: {
+  effectId: unknown;
+  operation: unknown;
+  input: unknown;
+  retryIn: unknown;
+  lookup: unknown;
+  perform: unknown;
+}): string | undefined {
+  const validated = validateEffectRequest(request);
+  if (!validated.ok) return validated.message;
+  const { retryIn } = request;
+  if (typeof retryIn !== "number" || !Number.isFinite(retryIn) || retryIn < 0) {
+    return "retryIn must be a finite, non-negative number";
+  }
+  if (typeof request.lookup !== "function") return "lookup must be a function";
+  if (typeof request.perform !== "function") return "perform must be a function";
+  return undefined;
+}
+
+async function* startStep(
+  key: string,
+  request: ExecutionStartRequest
+): AsyncGenerator<StepResponse, { executionId: string }, any> {
+  yield new StepKey(key);
+
+  const cached = yield new StepCacheCheck();
+
+  if (cached) return yield* replay(cached);
+
+  const started: RuntimeCallResult<ExecutionStartResult> =
+    yield new StepRuntimeCall(async ({ executions }) => {
+      // The client's failures are never the workflow's to catch: whatever it
+      // throws, the delivery ends and the next attempt calls it again.
+      try {
+        return await executions.start(request);
+      } catch (err) {
+        throw InfrastructureError.is(err) ? err : new InfrastructureError(err);
+      }
     });
-  } catch (err: unknown) {
+
+  if (!started.ok) {
+    throw new Error("ExecutionClient.start failed without an infrastructure error");
+  }
+
+  const outcome = started.value;
+
+  if (outcome.status === "rejected") {
+    const err = new WorkflowStartError(outcome.code, outcome.message);
     yield new StepError(err);
+    // unreachable – consumer calls throw() on the generator first
     throw err;
   }
+
+  return yield* record({ executionId: outcome.executionId });
 }
 
 function createWorkflowStore<T>(params: {
   definition: StoreDefinition<any>;
   id: string;
   event: WorkflowEvent;
-  storeClient: StoreClient;
 }): WorkflowStore<T> {
-  const { definition, id, event, storeClient } = params;
-  const key = { storeName: definition.name, storeId: id };
+  const { id, event } = params;
+  const key = { storeName: params.definition.name, storeId: id };
+  // Any operation may run `initial` on a missing store, so its errors must
+  // reach workflow code instead of ending the delivery as infrastructure.
+  const definition: StoreDefinition<any> = {
+    ...params.definition,
+    initial: workflowCallback(params.definition.initial),
+  };
 
   return {
-    definition,
+    definition: params.definition,
     id,
     key,
     get(stepKey?: string) {
-      return durableStep<StoreSnapshot<T>>(stepKey ?? getCallSiteHash(this.get), async () =>
-        (await storeClient.getStore({ definition, id })) as StoreSnapshot<T>
+      return runtimeStep(
+        stepKey ?? getCallSiteHash(this.get),
+        async ({ stores }) =>
+          (await stores.getStore({ definition, id })) as StoreSnapshot<T>
       );
     },
     select<R>(stepKey: string, selector: StoreSelector<T, R>) {
-      return durableStep(stepKey, async () => {
-        const snapshot = await storeClient.getStore({ definition, id });
+      return runtimeStep(stepKey, async ({ stores }) => {
+        const snapshot = await stores.getStore({ definition, id });
         // Run the selector against a clone so accidental mutation cannot
         // corrupt shared state (consistent with when, which throws via
         // its tracking proxy).
-        return selector(cloneStoreState(snapshot.state as T));
+        return workflowCallback(selector)(cloneStoreState(snapshot.state as T));
       });
     },
     update(stepKey, updater) {
-      return durableStep<StoreUpdateResult<T>>(stepKey, async () =>
-        (await storeClient.updateStore({
-          definition,
-          id,
-          updater: updater as any,
-          // Exactly-once across the store-commit/heap-write gap: on replay
-          // after a crash, the store's applied-steps ledger returns the
-          // recorded result instead of re-running the updater.
-          stepId: { executionId: event.executionId, stepKey },
-        })) as StoreUpdateResult<T>
+      return runtimeStep(
+        stepKey,
+        async ({ stores }) =>
+          (await stores.updateStore({
+            definition,
+            id,
+            updater: workflowCallback(updater, assertSynchronousUpdater) as any,
+            // Exactly-once across the store-commit/heap-write gap: on replay
+            // after a crash, the store's applied-steps ledger returns the
+            // recorded result instead of re-running the updater.
+            stepId: { executionId: event.executionId, stepKey },
+          })) as StoreUpdateResult<T>
       );
     },
     updateFrom(stepKey, snapshot, updater) {
-      return durableStep<StoreUpdateFromResult<T>>(stepKey, async () =>
-        (await storeClient.updateStoreFrom({
-          definition,
-          id,
-          snapshot,
-          updater: updater as any,
-          stepId: { executionId: event.executionId, stepKey },
-        })) as StoreUpdateFromResult<T>
+      return runtimeStep(
+        stepKey,
+        async ({ stores }) =>
+          (await stores.updateStoreFrom({
+            definition,
+            id,
+            snapshot,
+            updater: workflowCallback(updater, assertSynchronousUpdater) as any,
+            stepId: { executionId: event.executionId, stepKey },
+          })) as StoreUpdateFromResult<T>
       );
     },
     deleteFrom(stepKey, snapshot) {
-      return durableStep<StoreDeleteFromResult>(stepKey, () =>
-        storeClient.deleteStoreFrom({
+      return runtimeStep<StoreDeleteFromResult>(stepKey, ({ stores }) =>
+        stores.deleteStoreFrom({
           definition,
           id,
           snapshot,
@@ -234,7 +422,6 @@ function createWorkflowStore<T>(params: {
         definition,
         id,
         event,
-        storeClient,
         stepKey,
         selector,
       });
@@ -248,7 +435,6 @@ function createWorkflowStore<T>(params: {
         definition,
         id,
         event,
-        storeClient,
         stepKey,
         selector,
         claim,
@@ -257,14 +443,45 @@ function createWorkflowStore<T>(params: {
   };
 }
 
-function run<T extends any>(
-  fn: () => T | Promise<T>,
-): AsyncGenerator<StepResponse, T, StepResult | StepError>;
+export type StepRunOptions<S extends StandardSchemaV1> = {
+  /**
+   * Validates the step's decoded value, on the first run and on every
+   * replay; the step returns the schema's output. The step function must
+   * return the schema's input as a durable value.
+   */
+  schema: S;
+};
 
-function run<T extends any>(
+/**
+ * A step function whose result, or the value its promise resolves to, is a
+ * durable value. Checking the whole function type keeps the result's type as
+ * TypeScript infers it, and names the offending property in the error.
+ */
+export type DurableStepFn<F> = F extends () => infer R
+  ? () => R extends PromiseLike<infer U> ? PromiseLike<DurableUnion<U>> : DurableUnion<R>
+  : never;
+
+type DurableUnion<T> = T extends unknown ? Durable<T> : never;
+
+function run<F extends () => unknown>(
+  fn: F & DurableStepFn<F>
+): AsyncGenerator<StepResponse, Awaited<ReturnType<F>>, StepResponse>;
+
+function run<F extends () => unknown>(
+  key: string,
+  fn: F & DurableStepFn<F>
+): AsyncGenerator<StepResponse, Awaited<ReturnType<F>>, StepResponse>;
+
+function run<S extends StandardSchemaV1>(
+  fn: () => StepSchemaInput<S> | PromiseLike<StepSchemaInput<S>>,
+  options: StepRunOptions<S>
+): AsyncGenerator<StepResponse, StandardSchemaV1.InferOutput<S>, StepResponse>;
+
+function run<S extends StandardSchemaV1>(
   key: string,
-  fn: () => T | Promise<T>,
-): AsyncGenerator<StepResponse, T, StepResult | StepError>;
+  fn: () => StepSchemaInput<S> | PromiseLike<StepSchemaInput<S>>,
+  options: StepRunOptions<S>
+): AsyncGenerator<StepResponse, StandardSchemaV1.InferOutput<S>, StepResponse>;
 
 // NOTE: run/delay/poll are plain (synchronous) functions that delegate to
 // async generators. This matters for keyless steps: the call-site hash must be
@@ -272,59 +489,72 @@ function run<T extends any>(
 // generator body doesn't execute until its first .next(), which happens inside
 // the workflow driver, so hashing there would give every keyless step the same
 // key.
-function run<T extends any>(
-  arg1: string | (() => T | Promise<T>),
-  arg2?: (() => T | Promise<T>) | string,
-): AsyncGenerator<StepResponse, T, StepResult | StepError> {
+function run(
+  arg1: string | (() => unknown),
+  arg2?: (() => unknown) | StepRunOptions<StandardSchemaV1>,
+  arg3?: StepRunOptions<StandardSchemaV1>
+): AsyncGenerator<StepResponse, any, StepResponse> {
   let key: string | null = null;
-  let fn: () => T | Promise<T>;
+  let fn: () => unknown;
+  let options: StepRunOptions<StandardSchemaV1> | undefined;
 
   if (typeof arg1 === "string") {
     key = arg1;
-    fn = arg2 as () => T | Promise<T>;
+    fn = arg2 as () => unknown;
+    options = arg3;
   } else {
     fn = arg1;
+    options = arg2 as StepRunOptions<StandardSchemaV1> | undefined;
   }
 
   if (!key) key = getCallSiteHash(run);
 
-  return durableStep(key, fn);
+  return durableStep(key, fn, options?.schema);
 }
 
-function delay(retryInterval: number): any;
-function delay(key: string, retryInterval: number): any;
+type StepSchemaInput<S extends StandardSchemaV1> = Durable<
+  StandardSchemaV1.InferInput<S>
+>;
+
+/** `step.delay`, reading its wake time from the runtime clock. */
+function createDelay(clock: () => number) {
+  function delay(retryInterval: number): any;
+  function delay(key: string, retryInterval: number): any;
+  function delay(
+    arg1: string | number,
+    arg2?: number
+  ): AsyncGenerator<any, void, StepDelay> {
+    let key: string | null = null;
+    let retryInterval: number;
+
+    if (typeof arg1 === "string") {
+      key = arg1;
+      retryInterval = arg2!;
+    } else {
+      retryInterval = arg1;
+    }
 
-function delay(
-  arg1: string | number,
-  arg2?: number
-): AsyncGenerator<any, void, StepDelay> {
-  let key: string | null = null;
-  let retryInterval: number;
+    if (!key) key = getCallSiteHash(delay);
 
-  if (typeof arg1 === "string") {
-    key = arg1;
-    retryInterval = arg2!;
-  } else {
-    retryInterval = arg1;
+    return delayStep(key, retryInterval, clock);
   }
 
-  if (!key) key = getCallSiteHash(delay);
-
-  return delayStep(key, retryInterval);
+  return delay;
 }
 
 async function* delayStep(
   key: string,
-  retryInterval: number
+  retryInterval: number,
+  clock: () => number
 ): AsyncGenerator<any, void, StepDelay> {
   yield new StepKey(key);
 
   const cached = yield new StepCacheCheck();
 
   if (cached) {
-    yield new StepDelay(cached.resumeIn);
+    yield new StepDelay(cached.wakeAt);
   } else {
-    yield new StepDelay(Date.now() + retryInterval);
+    yield new StepDelay(clock() + retryInterval);
   }
 }
 
@@ -366,27 +596,32 @@ function poll(
   return run(key, task);
 }
 
-async function* durableStep<T extends any>(
+async function* durableStep(
   key: string,
-  fn: () => T | Promise<T>
-): AsyncGenerator<StepResponse, T, StepResult | StepError> {
+  fn: () => unknown,
+  schema: StandardSchemaV1 | undefined
+): AsyncGenerator<StepResponse, any, any> {
   yield new StepKey(key);
 
   const cached = yield new StepCacheCheck();
 
   if (cached) {
-    yield cached;
-    if (cached instanceof StepError) {
-      // unreachable – consumer calls throw() on the generator first
-      throw cached.err;
+    const value = yield* replay(cached);
+    if (!schema) return value;
+    const validated = await validateStepValue(key, schema, value);
+    if (cached.schema === "rejected" && validated.ok) {
+      throw new WorkflowValueError(
+        "STEP_SCHEMA_INVALID",
+        key,
+        `Step "${key}" schema rejected its value on the first run`
+      );
     }
-    return cached.result;
+    return unwrapValidated(validated);
   }
 
+  let result: unknown;
   try {
-    const result = await fn();
-    yield new StepResult(result);
-    return result;
+    result = await fn();
   } catch (err: unknown) {
     if (err instanceof RetryableError) {
       yield new StepError(err, {
@@ -399,47 +634,139 @@ async function* durableStep<T extends any>(
     // unreachable – consumer calls throw() on the generator first
     throw err;
   }
+
+  if (!schema) return yield* record(result);
+
+  // The schema validates the value as replay will decode it. A value the
+  // codec rejects is left for the driver to record as INVALID_STEP_VALUE.
+  const encoding = durableValueCodec.encode(result);
+  if (encoding.status === "invalid") return yield* record(result);
+  const decoding = durableValueCodec.decode(encoding.canonical);
+  if (decoding.status === "invalid") throw new Error(decoding.message);
+  const validated = await validateStepValue(key, schema, decoding.value);
+  yield* record(result, validated.ok ? "accepted" : "rejected");
+  return unwrapValidated(validated);
+}
+
+type ValidatedStepValue =
+  | { ok: true; value: unknown }
+  | { ok: false; error: WorkflowValueError };
+
+async function validateStepValue(
+  key: string,
+  schema: StandardSchemaV1,
+  value: unknown
+): Promise<ValidatedStepValue> {
+  const result = await schema["~standard"].validate(value);
+  if (result.issues) {
+    const messages = result.issues.map((issue) => issue.message).join("; ");
+    return {
+      ok: false,
+      error: new WorkflowValueError(
+        "STEP_SCHEMA_INVALID",
+        key,
+        `Step "${key}" value failed its schema: ${messages}`
+      ),
+    };
+  }
+  return { ok: true, value: result.value };
+}
+
+function unwrapValidated(validated: ValidatedStepValue) {
+  if (!validated.ok) throw validated.error;
+  return validated.value;
+}
+
+/**
+ * Yields a step's result for the driver to record, and returns the decoded
+ * record the driver resumes the step with, so that the first run returns
+ * what replay will.
+ */
+async function* record<T>(
+  result: T,
+  schema?: StepSchemaVerdict
+): AsyncGenerator<StepResponse, T, any> {
+  const recorded: StepResult = yield new StepResult(result, schema);
+  return recorded.result;
+}
+
+/** Yields a recorded response; the driver throws a recorded error into the step. */
+async function* replay(cached: StepResponse): AsyncGenerator<StepResponse, any, any> {
+  const recorded: StepResult = yield cached;
+  return recorded.result;
+}
+
+/**
+ * Yields a store step's error for the driver to record and throw. A value
+ * error raised by a store client is given this step's key.
+ */
+async function* fail(key: string, error: unknown): AsyncGenerator<StepResponse, never, any> {
+  const err =
+    error instanceof WorkflowValueError && error.stepKey !== key
+      ? new WorkflowValueError(error.code, key, error.message)
+      : error;
+  yield new StepError(err);
+  // unreachable – consumer calls throw() on the generator first
+  throw err;
+}
+
+/**
+ * A durable step whose work is a runtime client call. The driver makes the
+ * call and resumes the step with its result.
+ */
+async function* runtimeStep<T>(
+  key: string,
+  call: StepRuntimeCall<T>["call"]
+): AsyncGenerator<StepResponse, T, any> {
+  yield new StepKey(key);
+
+  const cached = yield new StepCacheCheck();
+
+  if (cached) return yield* replay(cached);
+
+  const outcome: RuntimeCallResult<T> = yield new StepRuntimeCall(call);
+
+  if (!outcome.ok) return yield* fail(key, outcome.error);
+
+  return yield* record(outcome.value);
 }
 
 async function* whenStep<T, R>(params: {
   definition: StoreDefinition<any>;
   id: string;
   event: WorkflowEvent;
-  storeClient: StoreClient;
   stepKey: string;
   selector: StoreSelector<T, R | undefined | null | false>;
-}): AsyncGenerator<StepResponse, NonNullable<R>, StepResult | StepStoreWait> {
-  const { definition, id, event, storeClient, stepKey, selector } = params;
+}): AsyncGenerator<StepResponse, NonNullable<R>, any> {
+  const { definition, id, event, stepKey, selector } = params;
 
   yield new StepKey(stepKey);
 
   const cached = yield new StepCacheCheck();
 
-  if (cached) {
-    yield cached;
-    if (!(cached instanceof StepResult)) {
-      throw new Error("Store wait step cache must resolve to a step result");
-    }
-    return cached.result;
-  }
+  if (cached) return yield* replay(cached);
 
-  const snapshot = await storeClient.getStore({ definition, id });
+  const read: RuntimeCallResult<StoreSnapshot<unknown>> =
+    yield new StepRuntimeCall(({ stores }) =>
+      stores.getStore({ definition, id })
+    );
+
+  if (!read.ok) return yield* fail(stepKey, read.error);
+
+  const snapshot = read.value;
   const { result, readPaths } = trackStoreSelector(
     cloneStoreState(snapshot.state as T),
     selector
   );
 
   if (result !== undefined && result !== null && result !== false) {
-    const clonedResult = cloneStoreState(result);
-    yield new StepResult(clonedResult);
-    return clonedResult as NonNullable<R>;
+    return yield* record(result as NonNullable<R>);
   }
 
-  await storeClient.registerWaiter({
+  yield* registerWaiter({
     workflowId: event.workflowId,
     executionId: event.executionId,
     stepKey,
-    event,
     storeName: definition.name,
     storeId: id,
     instanceId: snapshot.instanceId,
@@ -447,8 +774,6 @@ async function* whenStep<T, R>(params: {
     readPaths,
   });
 
-  yield new StepStoreWait();
-
   throw new Error("Store wait yielded control unexpectedly");
 }
 
@@ -456,54 +781,42 @@ async function* takeStep<T, R>(params: {
   definition: StoreDefinition<any>;
   id: string;
   event: WorkflowEvent;
-  storeClient: StoreClient;
   stepKey: string;
   selector: StoreSelector<T, R | undefined | null | false>;
   claim: (draft: Draft<T>, selected: NonNullable<R>) => void;
-}): AsyncGenerator<StepResponse, NonNullable<R>, StepResult | StepStoreWait> {
-  const { definition, id, event, storeClient, stepKey, selector, claim } =
-    params;
+}): AsyncGenerator<StepResponse, NonNullable<R>, any> {
+  const { definition, id, event, stepKey, selector, claim } = params;
 
   yield new StepKey(stepKey);
 
   const cached = yield new StepCacheCheck();
 
-  if (cached) {
-    yield cached;
-    if (!(cached instanceof StepResult)) {
-      throw new Error("Store take step cache must resolve to a step result");
-    }
-    return cached.result;
-  }
+  if (cached) return yield* replay(cached);
 
-  let outcome: StoreTakeResult<R>;
-  try {
-    outcome = await storeClient.takeFromStore({
-      definition,
-      id,
-      selector: selector as any,
-      claim: claim as any,
-      // Exactly-once across the store-commit/heap-write gap: on replay
-      // after a crash, the store's applied-steps ledger returns the
-      // recorded (matched) outcome instead of re-running selector/claim.
-      stepId: { executionId: event.executionId, stepKey },
-    });
-  } catch (err: unknown) {
-    yield new StepError(err);
-    // unreachable – consumer calls throw() on the generator first
-    throw err;
-  }
+  const taken: RuntimeCallResult<StoreTakeResult<R>> =
+    yield new StepRuntimeCall(({ stores }) =>
+      stores.takeFromStore({
+        definition,
+        id,
+        selector: workflowCallback(selector) as any,
+        claim: workflowCallback(claim, assertSynchronousClaim) as any,
+        // Exactly-once across the store-commit/heap-write gap: on replay
+        // after a crash, the store's applied-steps ledger returns the
+        // recorded (matched) outcome instead of re-running selector/claim.
+        stepId: { executionId: event.executionId, stepKey },
+      })
+    );
 
-  if (outcome.matched) {
-    yield new StepResult(outcome.selected);
-    return outcome.selected as NonNullable<R>;
-  }
+  if (!taken.ok) return yield* fail(stepKey, taken.error);
+
+  const outcome = taken.value;
 
-  await storeClient.registerWaiter({
+  if (outcome.matched) return yield* record(outcome.selected as NonNullable<R>);
+
+  yield* registerWaiter({
     workflowId: event.workflowId,
     executionId: event.executionId,
     stepKey,
-    event,
     storeName: definition.name,
     storeId: id,
     instanceId: outcome.instanceId,
@@ -511,7 +824,18 @@ async function* takeStep<T, R>(params: {
     readPaths: outcome.readPaths,
   });
 
-  yield new StepStoreWait();
-
   throw new Error("Store take yielded control unexpectedly");
 }
+
+/** Registers a waiter, then suspends the step until a store change wakes it. */
+async function* registerWaiter(
+  waiter: StoreWaiter
+): AsyncGenerator<StepResponse, void, any> {
+  const registered: RuntimeCallResult<void> = yield new StepRuntimeCall(
+    ({ stores }) => stores.registerWaiter(waiter)
+  );
+
+  if (!registered.ok) return yield* fail(waiter.stepKey, registered.error);
+
+  yield new StepStoreWait();
+}