bitfab 0.54.8 → 0.54.10

package/dist/{chunk-XPCVDQDS.js → chunk-A3OBQOLV.js}

@@ -11,6 +11,7 @@ import {
 } from "./chunk-EGWJYCNT.js";
 import {
   NO_MOCK_OVERRIDE,
+  createIntegrationFixtureContext,
   deserializeInputs,
   deserializeValue,
   randomUuid,
@@ -18,7 +19,7 @@ import {
   serializeValue,
   toJsonSafe,
   toJsonSafeReport
-} from "./chunk-RUCVJQBI.js";
+} from "./chunk-PRXW4B5X.js";
 import {
   BitfabError,
   DEFAULT_SERVICE_URL,
@@ -35,7 +36,7 @@ import {
   recordCallerTraceMetadata,
   retireTraceMetadata,
   warnOnce
-} from "./chunk-5H3D7J4K.js";
+} from "./chunk-B2O33DBH.js";
 import {
   __privateAdd,
   __privateGet,
@@ -2485,6 +2486,527 @@ var ReplayBranch = class {
 _url = new WeakMap();
 _context = new WeakMap();
 
+// src/rivetContext.ts
+var storage = null;
+var browserBinding = null;
+var STORAGE_SYMBOL = /* @__PURE__ */ Symbol.for("bitfab.rivetContextStorage");
+function initializeStorage() {
+  if (storage) {
+    return;
+  }
+  const shared = globalThis;
+  const existing = shared[STORAGE_SYMBOL];
+  if (existing) {
+    storage = existing;
+    return;
+  }
+  const created = createAsyncLocalStorage();
+  if (created) {
+    shared[STORAGE_SYMBOL] = created;
+    storage = created;
+  }
+}
+var rivetContextReady = asyncStorageReady.then(() => {
+  initializeStorage();
+});
+function getRivetContextBinding() {
+  initializeStorage();
+  return storage?.getStore() ?? browserBinding ?? void 0;
+}
+function runWithBrowserBinding(binding, fn) {
+  const previous = browserBinding;
+  browserBinding = binding;
+  try {
+    const result = fn();
+    if (result instanceof Promise) {
+      return result.finally(() => {
+        browserBinding = previous;
+      });
+    }
+    browserBinding = previous;
+    return result;
+  } catch (error) {
+    browserBinding = previous;
+    throw error;
+  }
+}
+function runWithRivetContext(binding, fn) {
+  initializeStorage();
+  if (!storage && !isAsyncStorageInitDone()) {
+    return rivetContextReady.then(() => runWithRivetContext(binding, fn));
+  }
+  if (storage) {
+    return storage.run(binding, fn);
+  }
+  return runWithBrowserBinding(binding, fn);
+}
+
+// src/rivetIntegration.ts
+var RIVET_DOCS_URL = "https://docs.bitfab.ai/frameworks/rivet";
+var RIVET_GROUP_EFFECT_METHODS = {
+  db: ["execute", "nativeMetrics", "close"],
+  kv: [
+    "get",
+    "put",
+    "delete",
+    "batchPut",
+    "batchGet",
+    "batchDelete",
+    "deleteRange",
+    "listPrefix",
+    "listRange",
+    "list"
+  ],
+  schedule: ["after", "at", "cancel", "get", "list"],
+  cron: ["set", "every", "get", "list", "delete", "history"],
+  queue: [
+    "send",
+    "next",
+    "nextBatch",
+    "waitForNames",
+    "waitForAvailable",
+    "complete",
+    "enqueueAndWait",
+    "tryNext",
+    "tryNextBatch"
+  ],
+  run: ["setWakeAt"]
+};
+var RIVET_TOP_LEVEL_EFFECTS = [
+  "broadcast",
+  "saveState",
+  "sleep",
+  "destroy"
+];
+var TOP_LEVEL_UNAVAILABLE_CAPABILITIES = /* @__PURE__ */ new Set([
+  "conn",
+  "conns",
+  "log",
+  "request",
+  "actorRuntimeSocket",
+  "keepAwake",
+  "waitUntil",
+  "client"
+]);
+var SYNC_VOID_EFFECTS = /* @__PURE__ */ new Set([
+  "broadcast",
+  "sleep",
+  "destroy"
+]);
+var QUEUE_MESSAGE_EFFECTS = /* @__PURE__ */ new Set([
+  "queue.next",
+  "queue.nextBatch",
+  "queue.waitForNames",
+  "queue.tryNext",
+  "queue.tryNextBatch"
+]);
+function isRecord2(value) {
+  return typeof value === "object" && value !== null;
+}
+function readProperty(value, property) {
+  try {
+    return Reflect.get(value, property, value);
+  } catch {
+    return void 0;
+  }
+}
+function snapshotContext(context) {
+  const source = context;
+  const snapshot = {};
+  for (const property of [
+    "state",
+    "vars",
+    "actorId",
+    "name",
+    "key",
+    "region",
+    "aborted"
+  ]) {
+    const value = readProperty(source, property);
+    if (value !== void 0) {
+      snapshot[property] = value;
+    }
+  }
+  return snapshot;
+}
+function unavailableCapability(path) {
+  return new BitfabError(
+    `Rivet context capability "${path}" is unavailable during replay. Add first-class support for that public context path before replaying code that uses it.`,
+    RIVET_DOCS_URL
+  );
+}
+function effectNode(traceFunctionKey, effect, mockSpanId) {
+  return {
+    traceFunctionKey,
+    spanName: effect,
+    type: "function",
+    originalSpanId: mockSpanId
+  };
+}
+function replayEffectStatus(replay, traceFunctionKey, effect, marked) {
+  const counterKey = `${traceFunctionKey}:${effect}`;
+  const callIndex = replay.callCounters?.get(counterKey) ?? 0;
+  const mockSpan = replay.mockTree?.spans.get(`${counterKey}:${callIndex}`);
+  const hasOverride = replay.mockOverrides?.some(
+    (override) => override.match(
+      effectNode(traceFunctionKey, effect, mockSpan?.sourceSpanId)
+    )
+  ) ?? false;
+  return {
+    willMock: hasOverride || replay.mockStrategy === "all" || replay.mockStrategy === "marked" && marked,
+    hasRecordedSpan: mockSpan !== void 0,
+    hasOverride,
+    callIndex
+  };
+}
+function assertReplayEffectAvailable(replay, traceFunctionKey, effect, marked) {
+  if (!replay) {
+    return;
+  }
+  const status = replayEffectStatus(replay, traceFunctionKey, effect, marked);
+  if (!status.willMock) {
+    throw new BitfabError(
+      `Rivet effect "${effect}" has no live actor context during replay. Add "${effect}" to mockOnReplay or provide a replay override.`,
+      RIVET_DOCS_URL
+    );
+  }
+  if (!status.hasRecordedSpan && !status.hasOverride) {
+    throw new BitfabError(
+      `No recorded Rivet effect "${effect}" exists at call ${status.callIndex + 1}; the live effect was not executed.`,
+      RIVET_DOCS_URL
+    );
+  }
+}
+function fixtureName(traceFunctionKey) {
+  return `rivet:${traceFunctionKey}`;
+}
+function snapshotFromReplay(traceFunctionKey, replay) {
+  const fixture = replay.integrationFixtures?.get(fixtureName(traceFunctionKey));
+  if (!isRecord2(fixture)) {
+    throw new BitfabError(
+      `Trace "${traceFunctionKey}" has no Rivet context fixture. Capture a fresh trace with the Rivet integration before replaying it.`,
+      RIVET_DOCS_URL
+    );
+  }
+  return fixture;
+}
+var BitfabRivetIntegration = class {
+  /** @internal Constructed by `Bitfab.getRivetIntegration()`. */
+  constructor(traceFunctionKey, runtime, options = {}) {
+    this.traceFunctionKey = traceFunctionKey;
+    this.runtime = runtime;
+    this.replayProxies = /* @__PURE__ */ new WeakMap();
+    this.markedEffects = new Set(options.mockOnReplay ?? []);
+  }
+  /**
+   * Run a callback with a live Rivet context bound to the current async chain.
+   * This binds context only; it creates no trace or span.
+   */
+  withContext(context, fn) {
+    if (!isRecord2(context)) {
+      throw new BitfabError("Rivet withContext requires a context object");
+    }
+    const binding = {
+      context,
+      fixtureContext: createIntegrationFixtureContext(
+        fixtureName(this.traceFunctionKey),
+        snapshotContext(context)
+      ),
+      recordedRoots: /* @__PURE__ */ new Set()
+    };
+    return runWithRivetContext(binding, fn);
+  }
+  /**
+   * Adapt a traced function into a Rivet handler whose first argument is the
+   * live actor context. The context is not added to the traced root's inputs.
+   */
+  bind(fn) {
+    const integration = this;
+    return function(context, ...args) {
+      return integration.withContext(context, () => fn.apply(this, args));
+    };
+  }
+  /**
+   * Return the current Rivet context proxy. In production it delegates to the
+   * bound live context. During replay it hydrates recorded actor state and
+   * intercepts supported effects without requiring a live Rivet actor.
+   */
+  getContext() {
+    const replay = getReplayContext();
+    if (replay) {
+      const existing = this.replayProxies.get(replay);
+      if (existing) {
+        return existing;
+      }
+      const proxy = this.createContextProxy(
+        snapshotFromReplay(this.traceFunctionKey, replay),
+        void 0
+      );
+      this.replayProxies.set(replay, proxy);
+      return proxy;
+    }
+    const binding = getRivetContextBinding();
+    if (!binding) {
+      throw new BitfabError(
+        "No Rivet context is bound. Call rivet.bind(handler) or rivet.withContext(context, callback) at the Rivet boundary.",
+        RIVET_DOCS_URL
+      );
+    }
+    this.runtime.recordRootFixture(
+      binding.fixtureContext,
+      binding.recordedRoots
+    );
+    if (!binding.proxy) {
+      binding.proxy = this.createContextProxy(
+        binding.context,
+        binding
+      );
+    }
+    return binding.proxy;
+  }
+  createContextProxy(target, binding) {
+    const replay = getReplayContext();
+    const dbProxy = this.createDbProxy(binding, replay);
+    const groupProxies = /* @__PURE__ */ new Map();
+    for (const group of Object.keys(
+      RIVET_GROUP_EFFECT_METHODS
+    )) {
+      if (group !== "db") {
+        groupProxies.set(group, this.createGroupProxy(group, binding, replay));
+      }
+    }
+    const topLevelEffects = /* @__PURE__ */ new Map();
+    for (const effect of RIVET_TOP_LEVEL_EFFECTS) {
+      topLevelEffects.set(
+        effect,
+        this.createEffect(
+          effect,
+          () => binding?.context,
+          effect,
+          replay
+        )
+      );
+    }
+    const replayTarget = replay ? { ...target } : target;
+    if (replay) {
+      const controller = new AbortController();
+      if (target.aborted === true) {
+        controller.abort();
+      }
+      replayTarget.abortSignal = controller.signal;
+    }
+    return new Proxy(replayTarget, {
+      get: (source, property) => {
+        if (property === "db") {
+          return dbProxy;
+        }
+        if (typeof property === "string") {
+          const group = groupProxies.get(property);
+          if (group) {
+            return group;
+          }
+          const effect = topLevelEffects.get(property);
+          if (effect) {
+            return effect;
+          }
+        }
+        if (property === "setPreventSleep" && replay) {
+          return (_preventSleep) => void 0;
+        }
+        if (replay && typeof property === "string" && TOP_LEVEL_UNAVAILABLE_CAPABILITIES.has(property)) {
+          throw unavailableCapability(property);
+        }
+        const value = Reflect.get(source, property, source);
+        return typeof value === "function" ? value.bind(source) : value;
+      }
+    });
+  }
+  createDbProxy(binding, replay) {
+    const proxy = this.createGroupProxy("db", binding, replay);
+    return new Proxy(proxy, {
+      get: (source, property) => property === "transaction" ? this.createTransaction(binding, replay) : Reflect.get(source, property, source)
+    });
+  }
+  createGroupProxy(group, binding, replay) {
+    const targetProvider = () => {
+      if (!binding) {
+        return void 0;
+      }
+      const context = binding.context;
+      const target = readProperty(context, group);
+      return isRecord2(target) ? target : void 0;
+    };
+    const proxy = this.createMethodProxy(
+      targetProvider,
+      group,
+      RIVET_GROUP_EFFECT_METHODS[group],
+      replay
+    );
+    if (group !== "queue") {
+      return proxy;
+    }
+    return new Proxy(proxy, {
+      get: (source, property) => property === "iter" ? (options) => this.createQueueIterator(source, options) : Reflect.get(source, property, source)
+    });
+  }
+  createMethodProxy(targetProvider, prefix, methods, replay) {
+    const effects = /* @__PURE__ */ new Map();
+    for (const method of methods) {
+      const effect = `${prefix}.${method}`;
+      effects.set(
+        method,
+        this.createEffect(effect, targetProvider, method, replay)
+      );
+    }
+    return new Proxy(
+      {},
+      {
+        get: (_source, property) => {
+          if (typeof property === "string") {
+            const effect = effects.get(property);
+            if (effect) {
+              return effect;
+            }
+            if (replay) {
+              throw unavailableCapability(`${prefix}.${property}`);
+            }
+          }
+          const target = targetProvider();
+          if (!target) {
+            return void 0;
+          }
+          const value = Reflect.get(target, property, target);
+          return typeof value === "function" ? value.bind(target) : value;
+        }
+      }
+    );
+  }
+  createEffect(effect, targetProvider, methodName, replay) {
+    const marked = this.markedEffects.has(effect);
+    const invoke = (...args) => {
+      assertReplayEffectAvailable(
+        getReplayContext(),
+        this.traceFunctionKey,
+        effect,
+        marked
+      );
+      const target = targetProvider();
+      const method = target ? readProperty(target, methodName) : void 0;
+      if (typeof method !== "function") {
+        throw unavailableCapability(effect);
+      }
+      return method.apply(target, args);
+    };
+    const wrapped = SYNC_VOID_EFFECTS.has(effect) ? this.runtime.wrapEffect(
+      effect,
+      marked,
+      invoke,
+      replay ? () => void 0 : void 0
+    ) : this.runtime.wrapEffect(effect, marked, async (...args) => {
+      return await invoke(...args);
+    });
+    return (...args) => {
+      if (replay) {
+        assertReplayEffectAvailable(
+          replay,
+          this.traceFunctionKey,
+          effect,
+          marked
+        );
+      }
+      const result = wrapped(...args);
+      if (!QUEUE_MESSAGE_EFFECTS.has(effect) || !(result instanceof Promise)) {
+        return result;
+      }
+      return result.then(
+        (value) => this.decorateQueueMessages(value, effect, args, replay)
+      );
+    };
+  }
+  createTransaction(binding, replay) {
+    return async (callback, options) => {
+      if (typeof callback !== "function") {
+        throw new BitfabError("Rivet db.transaction requires a callback");
+      }
+      if (replay) {
+        return await callback(
+          this.createMethodProxy(
+            () => void 0,
+            "db.transaction",
+            ["execute"],
+            replay
+          )
+        );
+      }
+      const context = binding?.context;
+      const db = context ? readProperty(context, "db") : void 0;
+      const transaction = isRecord2(db) ? readProperty(db, "transaction") : void 0;
+      if (!isRecord2(db) || typeof transaction !== "function") {
+        throw unavailableCapability("db.transaction");
+      }
+      return await transaction.call(
+        db,
+        async (transactionDb) => {
+          if (!isRecord2(transactionDb)) {
+            throw unavailableCapability("db.transaction.execute");
+          }
+          return await callback(
+            this.createMethodProxy(
+              () => transactionDb,
+              "db.transaction",
+              ["execute"],
+              replay
+            )
+          );
+        },
+        options
+      );
+    };
+  }
+  decorateQueueMessages(value, effect, args, replay) {
+    if (!this.queueCallIsCompletable(effect, args)) {
+      return value;
+    }
+    if (Array.isArray(value)) {
+      return value.map((message) => this.decorateQueueMessage(message, replay));
+    }
+    return this.decorateQueueMessage(value, replay);
+  }
+  decorateQueueMessage(value, replay) {
+    if (!isRecord2(value)) {
+      return value;
+    }
+    const originalComplete = readProperty(value, "complete");
+    const complete = this.createEffect(
+      "queue.message.complete",
+      () => typeof originalComplete === "function" ? { complete: originalComplete.bind(value) } : void 0,
+      "complete",
+      replay
+    );
+    return { ...value, complete };
+  }
+  queueCallIsCompletable(effect, args) {
+    const options = effect === "queue.waitForNames" ? args[1] : args[0];
+    return isRecord2(options) && options.completable === true;
+  }
+  createQueueIterator(queue, options) {
+    return {
+      async *[Symbol.asyncIterator]() {
+        for (; ; ) {
+          const next = Reflect.get(queue, "next", queue);
+          if (typeof next !== "function") {
+            throw unavailableCapability("queue.next");
+          }
+          const message = await next.call(queue, options);
+          if (message !== void 0 && message !== null) {
+            yield message;
+          }
+        }
+      }
+    };
+  }
+};
+
 // src/seedContext.ts
 var seedContextStorage = null;
 var SEED_CONTEXT_STORAGE_SYMBOL = /* @__PURE__ */ Symbol.for("bitfab.seedContextStorage");
@@ -4180,6 +4702,47 @@ var Bitfab = class {
       options
     });
   }
+  /**
+   * Get the first-class Rivet actor-context integration for an existing trace
+   * root. The integration binds context and records effect spans; it never
+   * creates another root.
+   *
+   * @param traceFunctionKey - Key used by the enclosing `withTrace()` root
+   * @param options - Explicit per-context-path replay mocking policy
+   * @experimental This API may change before it is stable.
+   */
+  getRivetIntegration(traceFunctionKey, options) {
+    return new BitfabRivetIntegration(
+      traceFunctionKey,
+      {
+        wrapEffect: (effect, mockOnReplay, fn, replayOutput) => {
+          const spanOptions = {
+            name: effect,
+            type: "function",
+            captureWhen: "nested",
+            mockOnReplay,
+            surface: "inherit",
+            instrumentation: "rivet",
+            ...replayOutput !== void 0 && { replayOutput }
+          };
+          return this.withSpan(traceFunctionKey, spanOptions, fn);
+        },
+        recordRootFixture: (context, recordedRoots) => {
+          const root = getSpanStack()[0];
+          if (!root) {
+            return;
+          }
+          const rootId = `${root.traceId}:${root.spanId}`;
+          if (recordedRoots.has(rootId)) {
+            return;
+          }
+          root.contexts.push(context);
+          recordedRoots.add(rootId);
+        }
+      },
+      options
+    );
+  }
   /**
    * Get a Claude Agent SDK handler for tracing.
    *
@@ -4678,6 +5241,7 @@ var Bitfab = class {
             }
             return output;
           };
+          const resolveBaseOutput = () => options.replayOutput !== void 0 ? options.replayOutput() : resolveRecordedOutput();
           const shouldMockWithBaseStrategy = replayCtxForMock.mockStrategy === "all" || replayCtxForMock.mockStrategy === "marked" && options.mockOnReplay === true;
           if (replayCtxForMock.mockOverrides?.length) {
             const nodeMeta = {
@@ -4734,7 +5298,7 @@ var Bitfab = class {
                 }
                 if (shouldMockWithBaseStrategy) {
                   const output = await resolveMockOrAbort(
-                    () => resolveRecordedOutput()
+                    () => resolveBaseOutput()
                   );
                   void sendSpan({
                     result: output,
@@ -4757,7 +5321,7 @@ var Bitfab = class {
             );
           }
           if (shouldMockWithBaseStrategy) {
-            const recorded = resolveRecordedOutput();
+            const recorded = resolveBaseOutput();
             if (recorded instanceof Promise) {
               return emitMockAsync(recorded, "recorded");
             }
@@ -5237,7 +5801,7 @@ var Bitfab = class {
       await runWithSeedContext({ traceId }, async () => target(...args));
     } finally {
       try {
-        const { flushTraces: flushTraces2 } = await import("./http-7PL32JLX.js");
+        const { flushTraces: flushTraces2 } = await import("./http-AVWWAMFV.js");
         await flushTraces2(3e4);
       } finally {
         unrecorded = activeTraceStates.delete(traceId);
@@ -5296,7 +5860,7 @@ var Bitfab = class {
         `Function is wrapped with trace function key '${wrappedKey}' but replay was called with '${traceFunctionKey}'. Pass matching keys, or pass the unwrapped function to replay it under the explicit key.`
       );
     }
-    const { replay: doReplay } = await import("./replay-KWBZ2JQV.js");
+    const { replay: doReplay } = await import("./replay-BCWIDCKH.js");
     return doReplay(
       this.httpClient,
       this.serviceUrl,
@@ -5359,6 +5923,28 @@ var BitfabFunction = class {
   span(options = {}) {
     return this.client.span(this.traceFunctionKey, options);
   }
+  withTrace(optionsOrFn, maybeFn) {
+    if (typeof optionsOrFn === "function") {
+      return this.client.withTrace(this.traceFunctionKey, optionsOrFn);
+    }
+    if (!maybeFn) {
+      throw new BitfabError("BitfabFunction.withTrace requires a function");
+    }
+    return this.client.withTrace(this.traceFunctionKey, optionsOrFn, maybeFn);
+  }
+  /**
+   * Create a trace method decorator bound to this function's key.
+   *
+   * Equivalent to `client.trace(key, options)` but reuses the key bound on
+   * this handle.
+   *
+   * @param options - Trace configuration applied to the decorated method
+   * @returns A method decorator configured for this key
+   * @experimental Automatic child-call instrumentation is experimental.
+   */
+  trace(options = {}) {
+    return this.client.trace(this.traceFunctionKey, options);
+  }
   /**
    * Get a Vercel AI SDK language-model middleware bound to this function's key.
    *
@@ -5456,6 +6042,20 @@ var BitfabFunction = class {
   getLangGraphIntegration(options) {
     return this.client.getLangGraphIntegration(this.traceFunctionKey, options);
   }
+  /**
+   * Get the first-class Rivet actor-context integration bound to this key.
+   * Use its `bind()` or `withContext()` outside this handle's `withTrace()`
+   * root, then call `getContext()` inside the root.
+   *
+   * @param options - Explicit per-context-path replay mocking policy
+   * @experimental This API may change before it is stable.
+   */
+  getRivetIntegration(options) {
+    return this.client.getRivetIntegration(
+      this.traceFunctionKey,
+      options
+    );
+  }
   /**
    * Wrap a BAML client method to automatically capture prompt and LLM metadata.
    * Delegates to the parent client's wrapBAML method.
@@ -5580,7 +6180,7 @@ async function seedFromRegistry(registry, pipeline, cases, options = {}) {
       sessionId: seedCase.sessionId
     })
   );
-  const { flushTraces: flushTraces2 } = await import("./http-7PL32JLX.js");
+  const { flushTraces: flushTraces2 } = await import("./http-AVWWAMFV.js");
   await flushTraces2(3e4);
   return { pipeline, traceFunctionKey, traceIds };
 }
@@ -5605,6 +6205,7 @@ export {
   BitfabLangGraphCallbackHandler,
   BitfabLangGraphIntegration,
   BitfabOpenAIAgentHandler,
+  BitfabRivetIntegration,
   TracesClient,
   BitfabOpenAITracingProcessor,
   BitfabVercelAiHandler,
@@ -5618,4 +6219,4 @@ export {
   reseedFromRegistry,
   seedFromRegistry
 };
-//# sourceMappingURL=chunk-XPCVDQDS.js.map
+//# sourceMappingURL=chunk-A3OBQOLV.js.map