rivetkit 2.3.3 → 2.3.5-rc.1

package/dist/tsup/workflow/mod.js

@@ -1,20 +1,18 @@
 import {
   createWorkflowInspectorAdapter
-} from "../chunk-BZM4PMFS.js";
-import "../chunk-KEKWJECX.js";
+} from "../chunk-P4OGQ4JL.js";
+import "../chunk-MGTSVSUC.js";
 import {
   ACTOR_CONTEXT_INTERNAL_SYMBOL,
-  RAW_STATE_SYMBOL,
   RUN_FUNCTION_CONFIG_SYMBOL
-} from "../chunk-XIX5DOZN.js";
+} from "../chunk-SHBE77IT.js";
 import {
   makeWorkflowKey,
   workflowStoragePrefix
-} from "../chunk-HXUEHHJF.js";
-import "../chunk-DHXWD4DF.js";
+} from "../chunk-JI6GZ2C2.js";
 import {
   stringifyError
-} from "../chunk-BLDO2FKO.js";
+} from "../chunk-B4IZU537.js";
 import {
   RivetError,
   isActorAbortedError
@@ -289,32 +287,13 @@ var WorkflowContext = class _WorkflowContext {
   isEvicted() {
     return this.#inner.isEvicted();
   }
-  // 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.
+  // Runs a user step body inside a fresh step context and deactivates the step
+  // context once the body settles. Workflows are roll-forward only, so state and
+  // vars mutations made before a failure always remain visible.
   async #runStep(run) {
     const stepCtx = new WorkflowStepContext(this.#runCtx, () => this.#markGuardTriggered());
-    let stateSnapshot = null;
-    try {
-      stateSnapshot = { state: this.#runCtx[RAW_STATE_SYMBOL]() };
-    } catch (error) {
-      this.#runCtx.log.debug({
-        msg: "failed to get state, likely due to being stateless workflow",
-        error
-      });
-    }
-    if (stateSnapshot) {
-      stateSnapshot.state = structuredClone(stateSnapshot.state);
-    }
-    const varsSnapshot = structuredClone(this.#runCtx.vars);
     try {
       return await run(stepCtx);
-    } catch (error) {
-      if (stateSnapshot) {
-        this.#runCtx.state = stateSnapshot.state;
-      }
-      this.#runCtx.vars = varsSnapshot;
-      throw error;
     } finally {
       stepCtx[DEACTIVATE_STEP]();
     }
@@ -334,8 +313,7 @@ var WorkflowContext = class _WorkflowContext {
       ...message.complete ? { complete: message.complete } : {}
     };
   }
-  // Runs a step rollback compensation with an active step context. Rollbacks
-  // intentionally mutate actor state, so their writes are not snapshotted.
+  // Runs a step rollback compensation with an active step context.
   async #runRollback(rollback, output) {
     const stepCtx = new WorkflowStepContext(this.#runCtx, () => this.#markGuardTriggered());
     try {
@@ -377,6 +355,8 @@ var WorkflowContext = class _WorkflowContext {
 
 // src/workflow/driver.ts
 var WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix();
+var WORKFLOW_SQLITE_MAX_BATCH_ROWS = 128;
+var WORKFLOW_SQLITE_MAX_BATCH_BYTES = 512 * 1024;
 function stripWorkflowKey(prefixed) {
   return prefixed.slice(WORKFLOW_STORAGE_PREFIX.length);
 }
@@ -390,6 +370,159 @@ function computeUpperBound(prefix) {
   }
   return null;
 }
+function normalizeSqlBlob(value) {
+  if (value instanceof Uint8Array) {
+    return value;
+  }
+  if (value instanceof ArrayBuffer) {
+    return new Uint8Array(value);
+  }
+  if (ArrayBuffer.isView(value)) {
+    return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
+  }
+  if (Array.isArray(value)) {
+    const bytes = new Uint8Array(value.length);
+    for (const [i, item] of value.entries()) {
+      if (!Number.isInteger(item) || item < 0 || item > 255) {
+        throw new Error("workflow sqlite value was not a byte array");
+      }
+      bytes[i] = item;
+    }
+    return bytes;
+  }
+  throw new Error("workflow sqlite value was not a blob");
+}
+function runtimeSqlFromContext(runCtx) {
+  const sql = runCtx == null ? void 0 : runCtx.sql;
+  if (sql && typeof sql === "object" && "query" in sql && "execute" in sql && "executeBatch" in sql && "run" in sql) {
+    return sql;
+  }
+  return void 0;
+}
+var WorkflowStorage = class {
+  #sql;
+  constructor(sql) {
+    if (!sql) {
+      throw new Error(
+        "workflow storage requires embedded SQLite; actor KV workflow storage is no longer supported"
+      );
+    }
+    if (!sql.executeBatch) {
+      throw new Error(
+        "workflow storage requires a SQLite database with executeBatch support"
+      );
+    }
+    this.#sql = sql;
+  }
+  async get(key) {
+    var _a;
+    const prefixed = makeWorkflowKey(key);
+    const result = await this.#sql.query(
+      "SELECT value FROM _rivet_wf_kv WHERE key = ?",
+      [prefixed]
+    );
+    return ((_a = result.rows[0]) == null ? void 0 : _a[0]) == null ? null : normalizeSqlBlob(result.rows[0][0]);
+  }
+  async set(key, value) {
+    await this.batch([{ key, value }]);
+  }
+  async delete(key) {
+    const prefixed = makeWorkflowKey(key);
+    await this.#sql.run("DELETE FROM _rivet_wf_kv WHERE key = ?", [
+      prefixed
+    ]);
+  }
+  async deletePrefix(prefix) {
+    const start = makeWorkflowKey(prefix);
+    const end = computeUpperBound(start);
+    if (end) {
+      await this.#sql.run(
+        "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?",
+        [start, end]
+      );
+      return;
+    }
+    const entries = await this.listRaw(start);
+    await this.deleteRawKeys(
+      entries.filter(([key]) => keyStartsWith(key, start)).map(([key]) => key)
+    );
+  }
+  async deleteRange(start, end) {
+    const prefixedStart = makeWorkflowKey(start);
+    const prefixedEnd = makeWorkflowKey(end);
+    await this.#sql.run(
+      "DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?",
+      [prefixedStart, prefixedEnd]
+    );
+  }
+  async list(prefix) {
+    const prefixed = makeWorkflowKey(prefix);
+    const entries = await this.listRaw(prefixed);
+    return entries.filter(([key]) => keyStartsWith(key, prefixed)).map(([key, value]) => ({
+      key: stripWorkflowKey(key),
+      value
+    }));
+  }
+  async batch(writes) {
+    if (writes.length === 0) return;
+    for (const chunk of chunkWorkflowWrites(writes)) {
+      await this.#sql.executeBatch(
+        chunk.map(({ key, value }) => ({
+          sql: "INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+          params: [makeWorkflowKey(key), value]
+        }))
+      );
+    }
+  }
+  async listRaw(prefix) {
+    const end = computeUpperBound(prefix);
+    const result = end ? await this.#sql.query(
+      "SELECT key, value FROM _rivet_wf_kv WHERE key >= ? AND key < ? ORDER BY key ASC",
+      [prefix, end]
+    ) : await this.#sql.query(
+      "SELECT key, value FROM _rivet_wf_kv WHERE key >= ? ORDER BY key ASC",
+      [prefix]
+    );
+    return result.rows.map((row) => [
+      normalizeSqlBlob(row[0]),
+      normalizeSqlBlob(row[1])
+    ]);
+  }
+  async deleteRawKeys(keys) {
+    for (let start = 0; start < keys.length; start += WORKFLOW_SQLITE_MAX_BATCH_ROWS) {
+      await this.#sql.executeBatch(
+        keys.slice(start, start + WORKFLOW_SQLITE_MAX_BATCH_ROWS).map((key) => ({
+          sql: "DELETE FROM _rivet_wf_kv WHERE key = ?",
+          params: [key]
+        }))
+      );
+    }
+  }
+};
+function chunkWorkflowWrites(writes) {
+  const chunks = [];
+  let chunk = [];
+  let chunkBytes = 0;
+  for (const write of writes) {
+    const writeBytes = WORKFLOW_STORAGE_PREFIX.length + write.key.length + write.value.length;
+    if (chunk.length > 0 && (chunk.length >= WORKFLOW_SQLITE_MAX_BATCH_ROWS || chunkBytes + writeBytes > WORKFLOW_SQLITE_MAX_BATCH_BYTES)) {
+      chunks.push(chunk);
+      chunk = [];
+      chunkBytes = 0;
+    }
+    chunk.push(write);
+    chunkBytes += writeBytes;
+  }
+  if (chunk.length > 0) chunks.push(chunk);
+  return chunks;
+}
+function keyStartsWith(key, prefix) {
+  if (key.length < prefix.length) return false;
+  for (let i = 0; i < prefix.length; i++) {
+    if (key[i] !== prefix[i]) return false;
+  }
+  return true;
+}
 var ActorWorkflowMessageDriver = class {
   #actor;
   #runCtx;
@@ -442,95 +575,44 @@ var ActorWorkflowMessageDriver = class {
   }
 };
 var ActorWorkflowDriver = class {
+  atomicBatch = true;
   workerPollInterval = 100;
   messageDriver;
   #actor;
   #runCtx;
+  #storage;
   constructor(actor, runCtx) {
     this.#actor = actor;
     this.#runCtx = runCtx;
     this.messageDriver = new ActorWorkflowMessageDriver(actor, runCtx);
+    this.#storage = new WorkflowStorage(runtimeSqlFromContext(runCtx));
   }
   async get(key) {
-    const [value] = await this.#runCtx.internalKeepAwake(
-      this.#actor.driver.kvBatchGet(this.#actor.id, [
-        makeWorkflowKey(key)
-      ])
-    );
-    return value ?? null;
+    return await this.#runCtx.internalKeepAwake(this.#storage.get(key));
   }
   async set(key, value) {
-    await this.#runCtx.internalKeepAwake(
-      this.#actor.driver.kvBatchPut(this.#actor.id, [
-        [makeWorkflowKey(key), value]
-      ])
-    );
+    await this.#runCtx.internalKeepAwake(this.#storage.set(key, value));
   }
   async delete(key) {
-    await this.#runCtx.internalKeepAwake(
-      this.#actor.driver.kvBatchDelete(this.#actor.id, [
-        makeWorkflowKey(key)
-      ])
-    );
+    await this.#runCtx.internalKeepAwake(this.#storage.delete(key));
   }
   async deletePrefix(prefix) {
-    const start = makeWorkflowKey(prefix);
-    const end = computeUpperBound(start);
-    if (end) {
-      await this.#runCtx.internalKeepAwake(
-        this.#actor.driver.kvDeleteRange(this.#actor.id, start, end)
-      );
-    } else {
-      const entries = await this.#runCtx.internalKeepAwake(
-        this.#actor.driver.kvListPrefix(this.#actor.id, start)
-      );
-      if (entries.length === 0) {
-        return;
-      }
-      await this.#runCtx.internalKeepAwake(
-        this.#actor.driver.kvBatchDelete(
-          this.#actor.id,
-          entries.map(([key]) => key)
-        )
-      );
-    }
+    await this.#runCtx.internalKeepAwake(
+      this.#storage.deletePrefix(prefix)
+    );
   }
   async deleteRange(start, end) {
     await this.#runCtx.internalKeepAwake(
-      this.#actor.driver.kvDeleteRange(
-        this.#actor.id,
-        makeWorkflowKey(start),
-        makeWorkflowKey(end)
-      )
+      this.#storage.deleteRange(start, end)
     );
   }
   async list(prefix) {
-    const entries = await this.#runCtx.internalKeepAwake(
-      this.#actor.driver.kvListPrefix(
-        this.#actor.id,
-        makeWorkflowKey(prefix)
-      )
-    );
-    return entries.map(([key, value]) => ({
-      key: stripWorkflowKey(key),
-      value
-    }));
+    return await this.#runCtx.internalKeepAwake(this.#storage.list(prefix));
   }
   async batch(writes) {
     if (writes.length === 0) return;
     await this.#runCtx.internalKeepAwake(
-      Promise.all([
-        this.#actor.driver.kvBatchPut(
-          this.#actor.id,
-          writes.map(({ key, value }) => [
-            makeWorkflowKey(key),
-            value
-          ])
-        ),
-        this.#actor.stateManager.saveState({
-          immediate: true
-        })
-      ])
+      this.#actor.stateManager.saveStateAndWorkflowBatch(writes)
     );
   }
   async setAlarm(_workflowId, wakeAt) {
@@ -563,69 +645,34 @@ var ActorWorkflowControlDriver = class {
   workerPollInterval = 100;
   messageDriver = new NoopWorkflowMessageDriver();
   #actor;
-  constructor(actor) {
+  #storage;
+  constructor(actor, runCtx) {
     this.#actor = actor;
+    this.#storage = new WorkflowStorage(runtimeSqlFromContext(runCtx));
   }
   async get(key) {
-    const [value] = await this.#actor.driver.kvBatchGet(this.#actor.id, [
-      makeWorkflowKey(key)
-    ]);
-    return value ?? null;
+    return await this.#storage.get(key);
   }
   async set(key, value) {
-    await this.#actor.driver.kvBatchPut(this.#actor.id, [
-      [makeWorkflowKey(key), value]
-    ]);
+    await this.#storage.set(key, value);
   }
   async delete(key) {
-    await this.#actor.driver.kvBatchDelete(this.#actor.id, [
-      makeWorkflowKey(key)
-    ]);
+    await this.#storage.delete(key);
   }
   async deletePrefix(prefix) {
-    const start = makeWorkflowKey(prefix);
-    const end = computeUpperBound(start);
-    if (end) {
-      await this.#actor.driver.kvDeleteRange(this.#actor.id, start, end);
-      return;
-    }
-    const entries = await this.#actor.driver.kvListPrefix(
-      this.#actor.id,
-      start
-    );
-    if (entries.length === 0) {
-      return;
-    }
-    await this.#actor.driver.kvBatchDelete(
-      this.#actor.id,
-      entries.map(([key]) => key)
-    );
+    await this.#storage.deletePrefix(prefix);
   }
   async deleteRange(start, end) {
-    await this.#actor.driver.kvDeleteRange(
-      this.#actor.id,
-      makeWorkflowKey(start),
-      makeWorkflowKey(end)
-    );
+    await this.#storage.deleteRange(start, end);
   }
   async list(prefix) {
-    const entries = await this.#actor.driver.kvListPrefix(
-      this.#actor.id,
-      makeWorkflowKey(prefix)
-    );
-    return entries.map(([key, value]) => ({
-      key: stripWorkflowKey(key),
-      value
-    }));
+    return await this.#storage.list(prefix);
   }
   async batch(writes) {
     if (writes.length === 0) {
       return;
     }
-    await this.#actor.driver.kvBatchPut(
-      this.#actor.id,
-      writes.map(({ key, value }) => [makeWorkflowKey(key), value])
-    );
+    await this.#storage.batch(writes);
   }
   async setAlarm(_workflowId, wakeAt) {
     await this.#actor.driver.setAlarm(this.#actor, wakeAt);
@@ -679,7 +726,7 @@ function workflow(fn, options = {}) {
     invariant(actor, "workflow() requires an actor instance");
     const workflowInspector = getWorkflowInspector(actor.id);
     const driver = new ActorWorkflowDriver(actor, runCtx);
-    const controlDriver = new ActorWorkflowControlDriver(actor);
+    const controlDriver = new ActorWorkflowControlDriver(actor, runCtx);
     workflowInspector.setReplayFromStep(async (entryId) => {
       const workflowState = await workflowInspector.adapter.getState();
       if (actor.isRunHandlerActive() || workflowState === "pending" || workflowState === "running") {