rivetkit 2.0.2 → 2.0.4-rc.1

package/dist/tsup/chunk-WMPW7JYC.js

@@ -0,0 +1,642 @@
+import {
+  encodeWorkflowHistoryTransport
+} from "./chunk-Z7HNQ2WF.js";
+import {
+  ACTOR_CONTEXT_INTERNAL_SYMBOL,
+  RUN_FUNCTION_CONFIG_SYMBOL,
+  makeWorkflowKey,
+  workflowStoragePrefix
+} from "./chunk-RTOCTWME.js";
+import {
+  assertUnreachable,
+  bufferToArrayBuffer,
+  stringifyError
+} from "./chunk-325TLXJT.js";
+
+// src/workflow/mod.ts
+import { runWorkflow } from "@rivetkit/workflow-engine";
+import invariant from "invariant";
+
+// src/workflow/constants.ts
+var WORKFLOW_GUARD_KV_KEY = "__rivet_actor_workflow_guard_triggered";
+
+// src/workflow/context.ts
+var ActorWorkflowContext = class _ActorWorkflowContext {
+  #inner;
+  #runCtx;
+  #actorAccessDepth = 0;
+  #allowActorAccess = false;
+  #guardViolation = false;
+  constructor(inner, runCtx) {
+    this.#inner = inner;
+    this.#runCtx = runCtx;
+  }
+  get workflowId() {
+    return this.#inner.workflowId;
+  }
+  get abortSignal() {
+    return this.#inner.abortSignal;
+  }
+  get queue() {
+    const self = this;
+    async function next(name, opts) {
+      const message = await self.#inner.queue.next(name, opts);
+      return self.#toActorQueueMessage(message);
+    }
+    async function nextBatch(name, opts) {
+      const messages = await self.#inner.queue.nextBatch(name, opts);
+      return messages.map((message) => self.#toActorQueueMessage(message));
+    }
+    async function send(name, body) {
+      self.#ensureActorAccess("queue.send");
+      await self.#runCtx.queue.send(name, body);
+    }
+    return {
+      next,
+      nextBatch,
+      send
+    };
+  }
+  async step(nameOrConfig, run) {
+    if (typeof nameOrConfig === "string") {
+      if (!run) {
+        throw new Error("Step run function missing");
+      }
+      return await this.#wrapActive(
+        () => this.#inner.step(nameOrConfig, () => this.#withActorAccess(run))
+      );
+    }
+    const stepConfig = nameOrConfig;
+    const config = {
+      ...stepConfig,
+      run: () => this.#withActorAccess(stepConfig.run)
+    };
+    return await this.#wrapActive(() => this.#inner.step(config));
+  }
+  async loop(nameOrConfig, run) {
+    if (typeof nameOrConfig === "string") {
+      if (!run) {
+        throw new Error("Loop run function missing");
+      }
+      return await this.#wrapActive(
+        () => this.#inner.loop(
+          nameOrConfig,
+          async (ctx) => run(this.#createChildContext(ctx))
+        )
+      );
+    }
+    const wrapped = {
+      ...nameOrConfig,
+      run: async (ctx, state) => nameOrConfig.run(this.#createChildContext(ctx), state)
+    };
+    return await this.#wrapActive(() => this.#inner.loop(wrapped));
+  }
+  sleep(name, durationMs) {
+    return this.#inner.sleep(name, durationMs);
+  }
+  sleepUntil(name, timestampMs) {
+    return this.#inner.sleepUntil(name, timestampMs);
+  }
+  async rollbackCheckpoint(name) {
+    await this.#wrapActive(() => this.#inner.rollbackCheckpoint(name));
+  }
+  async join(name, branches) {
+    const wrappedBranches = Object.fromEntries(
+      Object.entries(branches).map(([key, branch]) => [
+        key,
+        {
+          run: async (ctx) => branch.run(this.#createChildContext(ctx))
+        }
+      ])
+    );
+    return await this.#wrapActive(
+      () => this.#inner.join(name, wrappedBranches)
+    );
+  }
+  async race(name, branches) {
+    const wrappedBranches = branches.map((branch) => ({
+      name: branch.name,
+      run: (ctx) => branch.run(this.#createChildContext(ctx))
+    }));
+    return await this.#wrapActive(
+      () => this.#inner.race(name, wrappedBranches)
+    );
+  }
+  async removed(name, originalType) {
+    await this.#wrapActive(() => this.#inner.removed(name, originalType));
+  }
+  isEvicted() {
+    return this.#inner.isEvicted();
+  }
+  get state() {
+    this.#ensureActorAccess("state");
+    return this.#runCtx.state;
+  }
+  get vars() {
+    this.#ensureActorAccess("vars");
+    return this.#runCtx.vars;
+  }
+  client() {
+    this.#ensureActorAccess("client");
+    return this.#runCtx.client();
+  }
+  get db() {
+    this.#ensureActorAccess("db");
+    return this.#runCtx.db;
+  }
+  get log() {
+    return this.#runCtx.log;
+  }
+  keepAwake(promise) {
+    this.#ensureActorAccess("keepAwake");
+    return this.#runCtx.keepAwake(promise);
+  }
+  waitUntil(promise) {
+    this.#ensureActorAccess("waitUntil");
+    this.#runCtx.waitUntil(promise);
+  }
+  get actorId() {
+    return this.#runCtx.actorId;
+  }
+  broadcast(name, ...args) {
+    this.#ensureActorAccess("broadcast");
+    this.#runCtx.broadcast(
+      name,
+      ...args
+    );
+  }
+  #toActorQueueMessage(message) {
+    let id;
+    try {
+      id = BigInt(message.id);
+    } catch {
+      throw new Error(`Invalid queue message id "${message.id}"`);
+    }
+    return {
+      id,
+      name: message.name,
+      body: message.body,
+      createdAt: message.createdAt,
+      ...message.complete ? { complete: message.complete } : {}
+    };
+  }
+  async #wrapActive(run) {
+    return await this.#runCtx.keepAwake(run());
+  }
+  async #withActorAccess(run) {
+    this.#actorAccessDepth++;
+    if (this.#actorAccessDepth === 1) {
+      this.#allowActorAccess = true;
+    }
+    try {
+      return await run();
+    } finally {
+      this.#actorAccessDepth--;
+      if (this.#actorAccessDepth === 0) {
+        this.#allowActorAccess = false;
+      }
+    }
+  }
+  #ensureActorAccess(feature) {
+    if (!this.#allowActorAccess) {
+      this.#guardViolation = true;
+      this.#markGuardTriggered();
+      throw new Error(
+        `${feature} is only available inside workflow steps`
+      );
+    }
+  }
+  consumeGuardViolation() {
+    const violated = this.#guardViolation;
+    this.#guardViolation = false;
+    return violated;
+  }
+  #markGuardTriggered() {
+    try {
+      const state = this.#runCtx.state;
+      if (state && typeof state === "object" && "guardTriggered" in state) {
+        state.guardTriggered = true;
+      }
+    } catch {
+    }
+    this.#runCtx.waitUntil(
+      (async () => {
+        try {
+          await this.#runCtx.kv.put(WORKFLOW_GUARD_KV_KEY, "true");
+        } catch (error) {
+          this.#runCtx.log.error({
+            msg: "failed to persist workflow guard flag",
+            error
+          });
+        }
+      })()
+    );
+  }
+  #createChildContext(ctx) {
+    return new _ActorWorkflowContext(ctx, this.#runCtx);
+  }
+};
+
+// src/workflow/driver.ts
+var WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix();
+function stripWorkflowKey(prefixed) {
+  return prefixed.slice(WORKFLOW_STORAGE_PREFIX.length);
+}
+var ActorWorkflowMessageDriver = class {
+  #actor;
+  #runCtx;
+  constructor(actor, runCtx) {
+    this.#actor = actor;
+    this.#runCtx = runCtx;
+  }
+  async addMessage(message) {
+    await this.#runCtx.keepAwake(
+      this.#actor.queueManager.enqueue(message.name, message.data)
+    );
+  }
+  async receiveMessages(opts) {
+    const messages = await this.#runCtx.keepAwake(
+      this.#actor.queueManager.receive(
+        opts.names && opts.names.length > 0 ? [...opts.names] : void 0,
+        opts.count,
+        0,
+        void 0,
+        opts.completable
+      )
+    );
+    return messages.map((message) => ({
+      id: message.id.toString(),
+      name: message.name,
+      data: message.body,
+      sentAt: message.createdAt,
+      ...opts.completable ? {
+        complete: async (response) => {
+          await this.#runCtx.keepAwake(
+            this.#actor.queueManager.completeMessage(
+              message,
+              response
+            )
+          );
+        }
+      } : {}
+    }));
+  }
+  async completeMessage(messageId, response) {
+    let parsedId;
+    try {
+      parsedId = BigInt(messageId);
+    } catch {
+      return;
+    }
+    await this.#runCtx.keepAwake(
+      this.#actor.queueManager.completeMessageById(parsedId, response)
+    );
+  }
+};
+var ActorWorkflowDriver = class {
+  workerPollInterval = 100;
+  messageDriver;
+  #actor;
+  #runCtx;
+  constructor(actor, runCtx) {
+    this.#actor = actor;
+    this.#runCtx = runCtx;
+    this.messageDriver = new ActorWorkflowMessageDriver(actor, runCtx);
+  }
+  async get(key) {
+    const [value] = await this.#runCtx.keepAwake(
+      this.#actor.driver.kvBatchGet(this.#actor.id, [
+        makeWorkflowKey(key)
+      ])
+    );
+    return value ?? null;
+  }
+  async set(key, value) {
+    await this.#runCtx.keepAwake(
+      this.#actor.driver.kvBatchPut(this.#actor.id, [
+        [makeWorkflowKey(key), value]
+      ])
+    );
+  }
+  async delete(key) {
+    await this.#runCtx.keepAwake(
+      this.#actor.driver.kvBatchDelete(this.#actor.id, [
+        makeWorkflowKey(key)
+      ])
+    );
+  }
+  async deletePrefix(prefix) {
+    const entries = await this.#runCtx.keepAwake(
+      this.#actor.driver.kvListPrefix(
+        this.#actor.id,
+        makeWorkflowKey(prefix)
+      )
+    );
+    if (entries.length === 0) {
+      return;
+    }
+    await this.#runCtx.keepAwake(
+      this.#actor.driver.kvBatchDelete(
+        this.#actor.id,
+        entries.map(([key]) => key)
+      )
+    );
+  }
+  async list(prefix) {
+    const entries = await this.#runCtx.keepAwake(
+      this.#actor.driver.kvListPrefix(
+        this.#actor.id,
+        makeWorkflowKey(prefix)
+      )
+    );
+    return entries.map(([key, value]) => ({
+      key: stripWorkflowKey(key),
+      value
+    }));
+  }
+  async batch(writes) {
+    if (writes.length === 0) return;
+    await this.#runCtx.keepAwake(
+      Promise.all([
+        this.#actor.driver.kvBatchPut(
+          this.#actor.id,
+          writes.map(({ key, value }) => [makeWorkflowKey(key), value])
+        ),
+        this.#actor.stateManager.saveState({
+          immediate: true,
+          allowStoppingState: true
+        })
+      ])
+    );
+  }
+  async setAlarm(_workflowId, wakeAt) {
+    await this.#runCtx.keepAwake(
+      this.#actor.driver.setAlarm(this.#actor, wakeAt)
+    );
+  }
+  async clearAlarm(_workflowId) {
+    return;
+  }
+  waitForMessages(messageNames, abortSignal) {
+    return this.#actor.queueManager.waitForNames(
+      messageNames.length > 0 ? messageNames : void 0,
+      abortSignal
+    );
+  }
+};
+
+// src/workflow/inspector.ts
+import * as cbor from "cbor-x";
+import { createNanoEvents } from "nanoevents";
+function createWorkflowInspectorAdapter() {
+  const emitter = createNanoEvents();
+  let history = null;
+  const adapter = {
+    getHistory: () => history,
+    onHistoryUpdated: (listener) => emitter.on("updated", listener)
+  };
+  const update = (snapshot) => {
+    const transportHistory = toWorkflowHistory(snapshot);
+    const next = encodeWorkflowHistoryTransport(transportHistory);
+    history = next;
+    emitter.emit("updated", next);
+  };
+  return { adapter, update };
+}
+function encodeCbor(value) {
+  return bufferToArrayBuffer(cbor.encode(value));
+}
+function encodeOptionalCbor(value) {
+  if (value === void 0) {
+    return null;
+  }
+  return encodeCbor(value);
+}
+function toU64(value) {
+  return BigInt(Math.max(0, Math.floor(value)));
+}
+function toWorkflowLocation(location) {
+  return location.map((segment) => {
+    if (typeof segment === "number") {
+      return { tag: "WorkflowNameIndex", val: segment };
+    }
+    return {
+      tag: "WorkflowLoopIterationMarker",
+      val: {
+        loop: segment.loop,
+        iteration: segment.iteration
+      }
+    };
+  });
+}
+function toWorkflowEntryKind(kind) {
+  switch (kind.type) {
+    case "step":
+      return {
+        tag: "WorkflowStepEntry",
+        val: {
+          output: encodeOptionalCbor(kind.data.output),
+          error: kind.data.error ?? null
+        }
+      };
+    case "loop":
+      return {
+        tag: "WorkflowLoopEntry",
+        val: {
+          state: encodeCbor(kind.data.state),
+          iteration: kind.data.iteration,
+          output: encodeOptionalCbor(kind.data.output)
+        }
+      };
+    case "sleep":
+      return {
+        tag: "WorkflowSleepEntry",
+        val: {
+          deadline: toU64(kind.data.deadline),
+          state: toWorkflowSleepState(kind.data.state)
+        }
+      };
+    case "message":
+      return {
+        tag: "WorkflowMessageEntry",
+        val: {
+          name: kind.data.name,
+          messageData: encodeCbor(kind.data.data)
+        }
+      };
+    case "rollback_checkpoint":
+      return {
+        tag: "WorkflowRollbackCheckpointEntry",
+        val: { name: kind.data.name }
+      };
+    case "join":
+      return {
+        tag: "WorkflowJoinEntry",
+        val: { branches: toWorkflowBranchStatusMap(kind.data.branches) }
+      };
+    case "race":
+      return {
+        tag: "WorkflowRaceEntry",
+        val: {
+          winner: kind.data.winner ?? null,
+          branches: toWorkflowBranchStatusMap(kind.data.branches)
+        }
+      };
+    case "removed":
+      return {
+        tag: "WorkflowRemovedEntry",
+        val: {
+          originalType: kind.data.originalType,
+          originalName: kind.data.originalName ?? null
+        }
+      };
+    default:
+      assertUnreachable(kind);
+  }
+}
+function toWorkflowEntry(entry) {
+  return {
+    id: entry.id,
+    location: toWorkflowLocation(entry.location),
+    kind: toWorkflowEntryKind(entry.kind)
+  };
+}
+function toWorkflowEntryStatus(status) {
+  switch (status) {
+    case "pending":
+      return "PENDING" /* PENDING */;
+    case "running":
+      return "RUNNING" /* RUNNING */;
+    case "completed":
+      return "COMPLETED" /* COMPLETED */;
+    case "failed":
+      return "FAILED" /* FAILED */;
+    case "exhausted":
+      return "EXHAUSTED" /* EXHAUSTED */;
+    default:
+      assertUnreachable(status);
+  }
+}
+function toWorkflowSleepState(state) {
+  switch (state) {
+    case "pending":
+      return "PENDING" /* PENDING */;
+    case "completed":
+      return "COMPLETED" /* COMPLETED */;
+    case "interrupted":
+      return "INTERRUPTED" /* INTERRUPTED */;
+    default:
+      assertUnreachable(state);
+  }
+}
+function toWorkflowBranchStatusType(status) {
+  switch (status) {
+    case "pending":
+      return "PENDING" /* PENDING */;
+    case "running":
+      return "RUNNING" /* RUNNING */;
+    case "completed":
+      return "COMPLETED" /* COMPLETED */;
+    case "failed":
+      return "FAILED" /* FAILED */;
+    case "cancelled":
+      return "CANCELLED" /* CANCELLED */;
+    default:
+      assertUnreachable(status);
+  }
+}
+function toWorkflowBranchStatus(status) {
+  return {
+    status: toWorkflowBranchStatusType(status.status),
+    output: encodeOptionalCbor(status.output),
+    error: status.error ?? null
+  };
+}
+function toWorkflowBranchStatusMap(branches) {
+  return new Map(
+    Object.entries(branches).map(([name, status]) => [
+      name,
+      toWorkflowBranchStatus(status)
+    ])
+  );
+}
+function toWorkflowEntryMetadata(metadata) {
+  return {
+    status: toWorkflowEntryStatus(metadata.status),
+    error: metadata.error ?? null,
+    attempts: metadata.attempts,
+    lastAttemptAt: toU64(metadata.lastAttemptAt),
+    createdAt: toU64(metadata.createdAt),
+    completedAt: metadata.completedAt === void 0 ? null : toU64(metadata.completedAt),
+    rollbackCompletedAt: metadata.rollbackCompletedAt === void 0 ? null : toU64(metadata.rollbackCompletedAt),
+    rollbackError: metadata.rollbackError ?? null
+  };
+}
+function toWorkflowHistory(snapshot) {
+  const entryMetadata = /* @__PURE__ */ new Map();
+  for (const [id, metadata] of snapshot.entryMetadata) {
+    entryMetadata.set(id, toWorkflowEntryMetadata(metadata));
+  }
+  return {
+    nameRegistry: snapshot.nameRegistry,
+    entries: snapshot.entries.map((entry) => toWorkflowEntry(entry)),
+    entryMetadata
+  };
+}
+
+// src/workflow/mod.ts
+import { Loop } from "@rivetkit/workflow-engine";
+function workflow(fn) {
+  const workflowInspector = createWorkflowInspectorAdapter();
+  async function run(runCtx) {
+    const actor = runCtx[ACTOR_CONTEXT_INTERNAL_SYMBOL];
+    invariant(actor, "workflow() requires an actor instance");
+    const driver = new ActorWorkflowDriver(actor, runCtx);
+    const handle = runWorkflow(
+      actor.id,
+      async (ctx) => await fn(new ActorWorkflowContext(ctx, runCtx)),
+      void 0,
+      driver,
+      {
+        mode: "live",
+        logger: runCtx.log,
+        onHistoryUpdated: workflowInspector.update
+      }
+    );
+    const onAbort = () => {
+      handle.evict();
+    };
+    if (runCtx.abortSignal.aborted) {
+      onAbort();
+    } else {
+      runCtx.abortSignal.addEventListener("abort", onAbort, {
+        once: true
+      });
+    }
+    try {
+      await handle.result;
+    } catch (error) {
+      runCtx.log.error({
+        msg: "workflow run failed",
+        error: stringifyError(error)
+      });
+      throw error;
+    } finally {
+      runCtx.abortSignal.removeEventListener("abort", onAbort);
+    }
+  }
+  const runWithConfig = run;
+  runWithConfig[RUN_FUNCTION_CONFIG_SYMBOL] = {
+    icon: "diagram-project",
+    inspector: { workflow: workflowInspector.adapter }
+  };
+  return runWithConfig;
+}
+
+export {
+  WORKFLOW_GUARD_KV_KEY,
+  ActorWorkflowContext,
+  workflow,
+  Loop
+};
+//# sourceMappingURL=chunk-WMPW7JYC.js.map