rivetkit 2.3.0 → 2.3.1

package/dist/tsup/workflow/mod.js

@@ -1,23 +1,24 @@
 import {
   createWorkflowInspectorAdapter
-} from "../chunk-P2GNQ4RN.js";
-import "../chunk-UMZVD6DQ.js";
+} from "../chunk-H6VVZMWN.js";
+import "../chunk-NERUIBOT.js";
 import {
   ACTOR_CONTEXT_INTERNAL_SYMBOL,
   RAW_STATE_SYMBOL,
   RUN_FUNCTION_CONFIG_SYMBOL
-} from "../chunk-DPIMKYNB.js";
+} from "../chunk-XIX5DOZN.js";
 import {
   makeWorkflowKey,
   workflowStoragePrefix
 } from "../chunk-HXUEHHJF.js";
-import "../chunk-2OTRTA3J.js";
+import "../chunk-R6KPN5EW.js";
 import {
   stringifyError
-} from "../chunk-VTTFNQQI.js";
+} from "../chunk-L2X3YFER.js";
 import {
-  RivetError
-} from "../chunk-BRP62GZC.js";
+  RivetError,
+  isActorAbortedError
+} from "../chunk-ZZ3WBRPD.js";
 
 // src/workflow/mod.ts
 import {
@@ -38,12 +39,97 @@ import invariant from "invariant";
 var WORKFLOW_GUARD_KV_KEY = "__rivet_actor_workflow_guard_triggered";
 
 // src/workflow/context.ts
-var ActorWorkflowContext = class _ActorWorkflowContext {
+var DEACTIVATE_STEP = /* @__PURE__ */ Symbol("workflow.step.deactivate");
+var WorkflowStepContext = class {
+  #runCtx;
+  #active = true;
+  #onGuardViolation;
+  constructor(runCtx, onGuardViolation) {
+    this.#runCtx = runCtx;
+    this.#onGuardViolation = onGuardViolation;
+  }
+  [DEACTIVATE_STEP]() {
+    this.#active = false;
+  }
+  #ensureActive(feature) {
+    if (!this.#active) {
+      this.#onGuardViolation();
+      throw new Error(
+        `${feature} is only available inside workflow steps`
+      );
+    }
+  }
+  get actorId() {
+    return this.#runCtx.actorId;
+  }
+  get name() {
+    return this.#runCtx.name;
+  }
+  get key() {
+    return this.#runCtx.key;
+  }
+  get log() {
+    return this.#runCtx.log;
+  }
+  get abortSignal() {
+    return this.#runCtx.abortSignal;
+  }
+  get state() {
+    this.#ensureActive("state");
+    return this.#runCtx.state;
+  }
+  get vars() {
+    this.#ensureActive("vars");
+    return this.#runCtx.vars;
+  }
+  get db() {
+    this.#ensureActive("db");
+    return this.#runCtx.db;
+  }
+  client() {
+    this.#ensureActive("client");
+    return this.#runCtx.client();
+  }
+  broadcast(name, ...args) {
+    this.#ensureActive("broadcast");
+    this.#runCtx.broadcast(
+      name,
+      ...args
+    );
+  }
+  get queue() {
+    const self = this;
+    async function send(name, body) {
+      self.#ensureActive("queue.send");
+      await self.#runCtx.queue.send(name, body);
+    }
+    return { send };
+  }
+  /**
+   * Holds the actor awake for the duration of the provided promise. The actor
+   * cannot idle-sleep or finalize the sleep grace period until the promise
+   * settles.
+   */
+  keepAwake(promise) {
+    this.#ensureActive("keepAwake");
+    return this.#runCtx.keepAwake(promise);
+  }
+  /**
+   * Registers a promise that the sleep grace period will wait on. Use this for
+   * best-effort flush/cleanup work that may complete inside the grace window.
+   */
+  waitUntil(promise) {
+    this.#ensureActive("waitUntil");
+    this.#runCtx.waitUntil(promise);
+  }
+  destroy() {
+    this.#ensureActive("destroy");
+    this.#runCtx.destroy();
+  }
+};
+var WorkflowContext = class _WorkflowContext {
   #inner;
   #runCtx;
-  #actorAccessDepth = 0;
-  #allowActorAccess = false;
-  #guardViolation = false;
   constructor(inner, runCtx) {
     this.#inner = inner;
     this.#runCtx = runCtx;
@@ -54,6 +140,18 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
   get abortSignal() {
     return this.#inner.abortSignal;
   }
+  get actorId() {
+    return this.#runCtx.actorId;
+  }
+  get name() {
+    return this.#runCtx.name;
+  }
+  get key() {
+    return this.#runCtx.key;
+  }
+  get log() {
+    return this.#runCtx.log;
+  }
   get queue() {
     const self = this;
     async function next(name, opts) {
@@ -66,14 +164,9 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
         (message) => self.#toActorQueueMessage(message)
       );
     }
-    async function send(name, body) {
-      self.#ensureActorAccess("queue.send");
-      await self.#runCtx.queue.send(name, body);
-    }
     return {
       next,
-      nextBatch,
-      send
+      nextBatch
     };
   }
   async step(nameOrConfig, run) {
@@ -81,17 +174,17 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
       if (!run) {
         throw new Error("Step run function missing");
       }
+      const stepRun = run;
       return await this.#wrapActive(
-        () => this.#inner.step(
-          nameOrConfig,
-          () => this.#withActorAccessAndStateRollback(run)
-        )
+        () => this.#inner.step(nameOrConfig, () => this.#runStep(stepRun))
       );
     }
     const stepConfig = nameOrConfig;
+    const rollback = stepConfig.rollback;
     const config = {
       ...stepConfig,
-      run: () => this.#withActorAccessAndStateRollback(stepConfig.run)
+      run: () => this.#runStep(stepConfig.run),
+      rollback: rollback ? (_ctx, output) => this.#runRollback(rollback, output) : void 0
     };
     return await this.#wrapActive(() => this.#inner.step(config));
   }
@@ -100,17 +193,17 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
       if (!run) {
         throw new Error("Step run function missing");
       }
+      const stepRun = run;
       return await this.#wrapActive(
-        () => this.#inner.tryStep(
-          nameOrConfig,
-          () => this.#withActorAccessAndStateRollback(run)
-        )
+        () => this.#inner.tryStep(nameOrConfig, () => this.#runStep(stepRun))
       );
     }
     const stepConfig = nameOrConfig;
+    const rollback = stepConfig.rollback;
     const config = {
       ...stepConfig,
-      run: () => this.#withActorAccessAndStateRollback(stepConfig.run)
+      run: () => this.#runStep(stepConfig.run),
+      rollback: rollback ? (_ctx, output) => this.#runRollback(rollback, output) : void 0
     };
     return await this.#wrapActive(() => this.#inner.tryStep(config));
   }
@@ -119,10 +212,11 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
       if (!run) {
         throw new Error("Try run function missing");
       }
+      const tryRun = run;
       return await this.#wrapActive(
         () => this.#inner.try(
           nameOrConfig,
-          async (ctx) => run(this.#createChildContext(ctx))
+          async (ctx) => tryRun(this.#createChildContext(ctx))
         )
       );
     }
@@ -138,16 +232,20 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
       if (!run) {
         throw new Error("Loop run function missing");
       }
+      const loopRun = run;
       return await this.#wrapActive(
         () => this.#inner.loop(
           nameOrConfig,
-          async (ctx) => run(this.#createChildContext(ctx))
+          // A void return (no explicit Loop result) is undefined at
+          // runtime, which the engine treats as continue.
+          async (ctx) => await loopRun(this.#createChildContext(ctx)) ?? void 0
         )
       );
     }
+    const loopConfig = nameOrConfig;
     const wrapped = {
-      ...nameOrConfig,
-      run: async (ctx, state) => nameOrConfig.run(this.#createChildContext(ctx), state)
+      ...loopConfig,
+      run: (async (ctx, state) => await loopConfig.run(this.#createChildContext(ctx), state) ?? void 0)
     };
     return await this.#wrapActive(() => this.#inner.loop(wrapped));
   }
@@ -157,10 +255,6 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
   sleepUntil(name, timestampMs) {
     return this.#inner.sleepUntil(name, timestampMs);
   }
-  destroy() {
-    this.#ensureActorAccess("destroy");
-    this.#runCtx.destroy();
-  }
   async rollbackCheckpoint(name) {
     await this.#wrapActive(() => this.#inner.rollbackCheckpoint(name));
   }
@@ -192,96 +286,11 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
   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;
-  }
-  /** @deprecated No-op. Use `keepAwake(promise)` or `waitUntil(promise)` instead. */
-  setPreventSleep(_prevent) {
-    this.#ensureActorAccess("setPreventSleep");
-  }
-  /** @deprecated No-op. Always returns `false`. */
-  get preventSleep() {
-    this.#ensureActorAccess("preventSleep");
-    return false;
-  }
-  /**
-   * Holds the actor awake for the duration of the provided promise. The
-   * actor cannot idle-sleep or finalize the sleep grace period until the
-   * promise settles.
-   */
-  keepAwake(promise) {
-    this.#ensureActorAccess("keepAwake");
-    return this.#runCtx.keepAwake(promise);
-  }
-  /**
-   * Registers a promise that the sleep grace period will wait on. Use this
-   * for best-effort flush/cleanup work that may complete inside the grace
-   * window. For work the actor must stay running through, prefer
-   * `c.keepAwake(promise)` which also blocks idle sleep.
-   */
-  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.internalKeepAwake(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;
-      }
-    }
-  }
-  async #withActorAccessAndStateRollback(run) {
+  // Runs a user step body inside a fresh step context, snapshotting actor
+  // state/vars so a thrown step rolls back its mutations, and deactivating the
+  // step context once the body settles so it cannot be used after the step.
+  async #runStep(run) {
+    const stepCtx = new WorkflowStepContext(this.#runCtx, () => this.#markGuardTriggered());
     let stateSnapshot = null;
     try {
       stateSnapshot = { state: this.#runCtx[RAW_STATE_SYMBOL]() };
@@ -296,29 +305,47 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
     }
     const varsSnapshot = structuredClone(this.#runCtx.vars);
     try {
-      return await this.#withActorAccess(run);
+      return await run(stepCtx);
     } catch (error) {
       if (stateSnapshot) {
         this.#runCtx.state = stateSnapshot.state;
       }
       this.#runCtx.vars = varsSnapshot;
       throw error;
+    } finally {
+      stepCtx[DEACTIVATE_STEP]();
     }
   }
-  #ensureActorAccess(feature) {
-    if (!this.#allowActorAccess) {
-      this.#guardViolation = true;
-      this.#markGuardTriggered();
-      throw new Error(
-        `${feature} is only available inside workflow steps`
-      );
+  #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 } : {}
+    };
   }
-  consumeGuardViolation() {
-    const violated = this.#guardViolation;
-    this.#guardViolation = false;
-    return violated;
+  // Runs a step rollback compensation with an active step context. Rollbacks
+  // intentionally mutate actor state, so their writes are not snapshotted.
+  async #runRollback(rollback, output) {
+    const stepCtx = new WorkflowStepContext(this.#runCtx, () => this.#markGuardTriggered());
+    try {
+      await rollback(stepCtx, output);
+    } finally {
+      stepCtx[DEACTIVATE_STEP]();
+    }
+  }
+  async #wrapActive(run) {
+    return await this.#runCtx.internalKeepAwake(run);
   }
+  // Records that a step context was used outside its step. Mirrors the value
+  // onto actor state and a KV flag so callers can observe the violation.
   #markGuardTriggered() {
     try {
       const state = this.#runCtx.state;
@@ -341,7 +368,7 @@ var ActorWorkflowContext = class _ActorWorkflowContext {
     );
   }
   #createChildContext(ctx) {
-    return new _ActorWorkflowContext(ctx, this.#runCtx);
+    return new _WorkflowContext(ctx, this.#runCtx);
   }
 };
 
@@ -675,11 +702,13 @@ function workflow(fn, options = {}) {
     });
     const handle = runWorkflow(
       actor.id,
-      async (ctx) => await fn(new ActorWorkflowContext(ctx, runCtx)),
+      async (ctx) => await fn(new WorkflowContext(ctx, runCtx)),
       void 0,
       driver,
       {
         mode: "live",
+        // The actor logger and the engine's pino logger are runtime
+        // compatible but not structurally assignable.
         logger: runCtx.log,
         onHistoryUpdated: workflowInspector.update,
         onError: onError ? async (event) => await onError(runCtx, event) : void 0
@@ -699,7 +728,7 @@ function workflow(fn, options = {}) {
     try {
       await handle.result;
     } catch (error) {
-      if (runCtx.abortSignal.aborted) {
+      if (runCtx.abortSignal.aborted || isActorAbortedError(error)) {
         return;
       }
       if (shouldRethrowWorkflowError(error)) {
@@ -720,6 +749,12 @@ function workflow(fn, options = {}) {
   const runWithConfig = run;
   runWithConfig[RUN_FUNCTION_CONFIG_SYMBOL] = {
     icon: "diagram-project",
+    // Drop the per-actor inspector when the actor is destroyed so this map
+    // does not retain one inspector (and its encoded history) per actor id
+    // for the process lifetime.
+    disposeInspector: (actorId) => {
+      workflowInspectors.delete(actorId);
+    },
     inspectorFactory: (actor) => {
       const actorId = resolveWorkflowInspectorActorId(actor);
       return {
@@ -751,8 +786,9 @@ function resolveWorkflowInspectorActorId(actor) {
   return void 0;
 }
 export {
-  ActorWorkflowContext,
   Loop,
+  WorkflowContext,
+  WorkflowStepContext,
   workflow
 };
 //# sourceMappingURL=mod.js.map