yieldstar 0.4.5 → 0.5.0

package/src/internal/step-runner.ts

@@ -1,13 +1,30 @@
 import {
+  cloneStoreState,
+  defineStore,
+  type Draft,
+  type StandardSchemaV1,
   StepResponse,
   StepKey,
   StepError,
   StepResult,
   StepDelay,
   StepCacheCheck,
+  StepStoreWait,
+  type StoreClient,
+  type StoreDefinition,
+  type StoreDeleteFromResult,
+  type StoreKey,
+  type StoreSelector,
+  type StoreSnapshot,
+  type StoreState,
+  type StoreTakeResult,
+  type StoreUpdateFromResult,
+  type StoreUpdateResult,
+  trackStoreSelector,
 } from "@yieldstar/core";
 import { RetryableError } from "../exports/errors";
 import { getCallSiteHash } from "./utils";
+import type { WorkflowEvent } from "@yieldstar/core";
 
 /**
  * @description A library of step generators, each of which:
@@ -17,9 +34,228 @@ import { getCallSiteHash } from "./utils";
  * @throws any error caught when running the user-defined function,
  * once retries have been exhausted, or if there are no retry semantics
  */
-export const stepRunner = { run, delay, poll };
+export { defineStore };
 
-export type StepRunner = typeof stepRunner;
+export type WorkflowStore<T> = {
+  readonly definition: StoreDefinition<any>;
+  readonly id: string;
+  readonly key: StoreKey;
+  get(key?: string): AsyncGenerator<StepResponse, StoreSnapshot<T>>;
+  select<R>(
+    key: string,
+    selector: StoreSelector<T, R>
+  ): AsyncGenerator<StepResponse, R>;
+  update(
+    key: string,
+    updater: (draft: Draft<T>) => void | T
+  ): AsyncGenerator<StepResponse, StoreUpdateResult<T>>;
+  updateFrom(
+    key: string,
+    snapshot: StoreSnapshot<T>,
+    updater: (draft: Draft<T>) => void | T
+  ): AsyncGenerator<StepResponse, StoreUpdateFromResult<T>>;
+  deleteFrom(
+    key: string,
+    snapshot: StoreSnapshot<T>
+  ): AsyncGenerator<StepResponse, StoreDeleteFromResult>;
+  when<R>(
+    selector: StoreSelector<T, R | undefined | null | false>
+  ): AsyncGenerator<StepResponse, NonNullable<R>>;
+  when<R>(
+    key: string,
+    selector: StoreSelector<T, R | undefined | null | false>
+  ): AsyncGenerator<StepResponse, NonNullable<R>>;
+  take<R>(
+    key: string,
+    selector: StoreSelector<T, R | undefined | null | false>,
+    claim: (draft: Draft<T>, selected: NonNullable<R>) => void
+  ): AsyncGenerator<StepResponse, NonNullable<R>>;
+};
+
+export type StepRunner = ReturnType<typeof createStepRunner>;
+
+export function createStepRunner(params: {
+  event: WorkflowEvent;
+  storeClient: StoreClient;
+}) {
+  const { event, storeClient } = params;
+
+  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,
+      id,
+      event,
+      storeClient,
+      initial: params?.initial,
+    });
+  }
+
+  return { run, delay, poll, store };
+}
+
+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;
+
+  yield new StepKey(key);
+
+  const cached = yield new StepCacheCheck();
+
+  if (cached) {
+    yield cached;
+    if (cached instanceof StepError) {
+      throw cached.err;
+    }
+    return createWorkflowStore({
+      definition,
+      id,
+      event,
+      storeClient,
+    });
+  }
+
+  try {
+    await storeClient.getOrCreateStore({
+      definition,
+      id,
+      initial,
+    });
+    yield new StepResult({ storeName: definition.name, storeId: id });
+    return createWorkflowStore({
+      definition,
+      id,
+      event,
+      storeClient,
+    });
+  } catch (err: unknown) {
+    yield new StepError(err);
+    throw err;
+  }
+}
+
+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 };
+
+  return {
+    definition,
+    id,
+    key,
+    get(stepKey?: string) {
+      return durableStep<StoreSnapshot<T>>(stepKey ?? getCallSiteHash(this.get), async () =>
+        (await storeClient.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 });
+        // 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));
+      });
+    },
+    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>
+      );
+    },
+    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>
+      );
+    },
+    deleteFrom(stepKey, snapshot) {
+      return durableStep<StoreDeleteFromResult>(stepKey, () =>
+        storeClient.deleteStoreFrom({
+          definition,
+          id,
+          snapshot,
+          stepId: { executionId: event.executionId, stepKey },
+        })
+      );
+    },
+    when<R>(
+      arg1:
+        | string
+        | StoreSelector<T, R | undefined | null | false>,
+      arg2?: StoreSelector<T, R | undefined | null | false>
+    ) {
+      const stepKey =
+        typeof arg1 === "string" ? arg1 : getCallSiteHash(this.when);
+      const selector = (
+        typeof arg1 === "string" ? arg2 : arg1
+      ) as StoreSelector<T, R | undefined | null | false>;
+
+      return whenStep({
+        definition,
+        id,
+        event,
+        storeClient,
+        stepKey,
+        selector,
+      });
+    },
+    take<R>(
+      stepKey: string,
+      selector: StoreSelector<T, R | undefined | null | false>,
+      claim: (draft: Draft<T>, selected: NonNullable<R>) => void
+    ) {
+      return takeStep({
+        definition,
+        id,
+        event,
+        storeClient,
+        stepKey,
+        selector,
+        claim,
+      });
+    },
+  };
+}
 
 function run<T extends any>(
   fn: () => T | Promise<T>,
@@ -30,7 +266,13 @@ function run<T extends any>(
   fn: () => T | Promise<T>,
 ): AsyncGenerator<StepResponse, T, StepResult | StepError>;
 
-async function* run<T extends any>(
+// NOTE: run/delay/poll are plain (synchronous) functions that delegate to
+// async generators. This matters for keyless steps: the call-site hash must be
+// captured while the user's call site is still on the stack. An async
+// 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> {
@@ -46,41 +288,13 @@ async function* run<T extends any>(
 
   if (!key) key = getCallSiteHash(run);
 
-  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;
-    }
-    return cached.result;
-  }
-
-  try {
-    const result = await fn();
-    yield new StepResult(result);
-    return result;
-  } catch (err: unknown) {
-    if (err instanceof RetryableError) {
-      yield new StepError(err, {
-        maxAttempts: err.maxAttempts,
-        retryInterval: err.retryInterval,
-      });
-    } else {
-      yield new StepError(err);
-    }
-    // unreachable – consumer calls throw() on the generator first
-    throw err;
-  }
+  return durableStep(key, fn);
 }
 
 function delay(retryInterval: number): any;
 function delay(key: string, retryInterval: number): any;
 
-async function* delay(
+function delay(
   arg1: string | number,
   arg2?: number
 ): AsyncGenerator<any, void, StepDelay> {
@@ -96,6 +310,13 @@ async function* delay(
 
   if (!key) key = getCallSiteHash(delay);
 
+  return delayStep(key, retryInterval);
+}
+
+async function* delayStep(
+  key: string,
+  retryInterval: number
+): AsyncGenerator<any, void, StepDelay> {
   yield new StepKey(key);
 
   const cached = yield new StepCacheCheck();
@@ -113,7 +334,7 @@ type PollPredicate = () => boolean | Promise<boolean>;
 function poll(opts: PollOpts, predicate: PollPredicate): any;
 function poll(key: string, opts: PollOpts, predicate: PollPredicate): any;
 
-async function* poll(
+function poll(
   arg1: string | PollOpts,
   arg2: PollOpts | PollPredicate,
   arg3?: PollPredicate
@@ -142,5 +363,155 @@ async function* poll(
     }
   };
 
-  yield* run(key, task, );
+  return run(key, task);
+}
+
+async function* durableStep<T extends any>(
+  key: string,
+  fn: () => T | Promise<T>
+): AsyncGenerator<StepResponse, T, StepResult | StepError> {
+  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;
+    }
+    return cached.result;
+  }
+
+  try {
+    const result = await fn();
+    yield new StepResult(result);
+    return result;
+  } catch (err: unknown) {
+    if (err instanceof RetryableError) {
+      yield new StepError(err, {
+        maxAttempts: err.maxAttempts,
+        retryInterval: err.retryInterval,
+      });
+    } else {
+      yield new StepError(err);
+    }
+    // unreachable – consumer calls throw() on the generator first
+    throw err;
+  }
+}
+
+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;
+
+  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;
+  }
+
+  const snapshot = await storeClient.getStore({ definition, id });
+  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>;
+  }
+
+  await storeClient.registerWaiter({
+    workflowId: event.workflowId,
+    executionId: event.executionId,
+    stepKey,
+    event,
+    storeName: definition.name,
+    storeId: id,
+    instanceId: snapshot.instanceId,
+    sinceVersion: snapshot.version,
+    readPaths,
+  });
+
+  yield new StepStoreWait();
+
+  throw new Error("Store wait yielded control unexpectedly");
+}
+
+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;
+
+  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;
+  }
+
+  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;
+  }
+
+  if (outcome.matched) {
+    yield new StepResult(outcome.selected);
+    return outcome.selected as NonNullable<R>;
+  }
+
+  await storeClient.registerWaiter({
+    workflowId: event.workflowId,
+    executionId: event.executionId,
+    stepKey,
+    event,
+    storeName: definition.name,
+    storeId: id,
+    instanceId: outcome.instanceId,
+    sinceVersion: outcome.version,
+    readPaths: outcome.readPaths,
+  });
+
+  yield new StepStoreWait();
+
+  throw new Error("Store take yielded control unexpectedly");
 }

package/src/internal/utils.ts

@@ -7,19 +7,33 @@ export const isIterable = <T>(i: T | Iterable<T>): i is Iterable<T> => {
 };
 
 
-export function getCallSiteHash(stepFn: Function) {
-  const originalStack = (stepFn as any).stack;
+/**
+ * Derives a replay-stable cache key from the call site of a keyless step.
+ *
+ * IMPORTANT: `stepFn` must be the synchronously-executing function that the
+ * user called directly (e.g. the `step.run` wrapper) – NOT an async generator
+ * body. Async generator bodies don't run until the first `.next()`, by which
+ * point the user's frame is no longer on the stack. `Error.captureStackTrace`
+ * omits every frame above (and including) `stepFn`, so the first captured
+ * frame is the user's call site (file:line:column), which is deterministic
+ * across executions of the same workflow code.
+ */
+export function getCallSiteHash(stepFn: Function): string {
+  const holder: { stack?: string } = {};
 
-  Error.captureStackTrace(stepFn, getCallSiteHash);
+  Error.captureStackTrace(holder, stepFn);
 
-  const stack = (stepFn as any).stack
-  const stepFnCallSite = stack.split(')')[1]
+  const frames = (holder.stack ?? "")
+    .split("\n")
+    .filter((line) => /^\s*at\s/.test(line));
 
-  if (originalStack === undefined) {
-    delete (stepFn as any).stack;
-  } else {
-    (stepFn as any).stack = originalStack;
+  const callSite = frames[0]?.trim();
+
+  if (!callSite) {
+    throw new Error(
+      "Could not derive a cache key from the call site. Pass an explicit key to this step."
+    );
   }
 
-  return createHash("sha1").update(stepFnCallSite).digest("hex");
+  return createHash("sha1").update(callSite).digest("hex");
 }