bitfab 0.54.9 → 0.54.10

package/dist/index.cjs

@@ -51,7 +51,7 @@ var __version__, __packageName__;
 var init_version_generated = __esm({
   "src/version.generated.ts"() {
     "use strict";
-    __version__ = "0.54.9";
+    __version__ = "0.54.10";
     __packageName__ = "bitfab";
   }
 });
@@ -836,8 +836,8 @@ function encodePayloadBody(payload) {
       const marker = { error: `payload_serialize_failed: ${message}` };
       return { body: JSON.stringify(marker), dropped, value: marker };
     }
-    const isRecord4 = typeof sanitized === "object" && sanitized !== null && !Array.isArray(sanitized);
-    if (dropped.length > 0 && isRecord4) {
+    const isRecord6 = typeof sanitized === "object" && sanitized !== null && !Array.isArray(sanitized);
+    if (dropped.length > 0 && isRecord6) {
       const obj = sanitized;
       const existing = Array.isArray(obj.errors) ? obj.errors : [];
       obj.errors = [
@@ -854,7 +854,7 @@ function encodePayloadBody(payload) {
     return {
       body: JSON.stringify(sanitized),
       dropped,
-      value: isRecord4 ? sanitized : void 0
+      value: isRecord6 ? sanitized : void 0
     };
   }
 }
@@ -956,13 +956,15 @@ var init_spanOrigin = __esm({
       "openai-agents",
       "langgraph",
       "claude-agent-sdk",
-      "vercel-ai"
+      "vercel-ai",
+      "rivet"
     ];
     FRAMEWORK_INSTRUMENTATIONS = [
       "openai-agents",
       "langgraph",
       "claude-agent-sdk",
-      "vercel-ai"
+      "vercel-ai",
+      "rivet"
     ];
     FRAMEWORK_INSTRUMENTATION_SET = new Set(
       FRAMEWORK_INSTRUMENTATIONS
@@ -3466,6 +3468,55 @@ var init_codeChange = __esm({
   }
 });
 
+// src/integrationFixture.ts
+function isRecord4(value) {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+function createIntegrationFixtureContext(name, value) {
+  const serialized = serializeValue(value);
+  return {
+    [INTEGRATION_FIXTURE_CONTEXT_KEY]: {
+      version: INTEGRATION_FIXTURE_VERSION,
+      name,
+      value: serialized.json,
+      ...serialized.meta !== void 0 && { valueMeta: serialized.meta }
+    }
+  };
+}
+function extractIntegrationFixtures(spanData) {
+  const contexts = spanData.contexts;
+  if (!Array.isArray(contexts)) {
+    return void 0;
+  }
+  const fixtures = /* @__PURE__ */ new Map();
+  for (const context of contexts) {
+    if (!isRecord4(context)) {
+      continue;
+    }
+    const fixture = context[INTEGRATION_FIXTURE_CONTEXT_KEY];
+    if (!isRecord4(fixture) || fixture.version !== INTEGRATION_FIXTURE_VERSION || typeof fixture.name !== "string") {
+      continue;
+    }
+    fixtures.set(
+      fixture.name,
+      deserializeValue({
+        json: fixture.value,
+        ...fixture.valueMeta !== void 0 && { meta: fixture.valueMeta }
+      })
+    );
+  }
+  return fixtures.size > 0 ? fixtures : void 0;
+}
+var INTEGRATION_FIXTURE_CONTEXT_KEY, INTEGRATION_FIXTURE_VERSION;
+var init_integrationFixture = __esm({
+  "src/integrationFixture.ts"() {
+    "use strict";
+    init_serialize();
+    INTEGRATION_FIXTURE_CONTEXT_KEY = "__bitfabIntegrationFixture";
+    INTEGRATION_FIXTURE_VERSION = 1;
+  }
+});
+
 // src/replay.ts
 var replay_exports = {};
 __export(replay_exports, {
@@ -3690,6 +3741,7 @@ async function processItem(httpClient, serverItem, fn, testRunId, mockStrategy,
     const spanData = span.rawData?.span_data ?? {};
     inputs = deserializeInputs(spanData);
     originalOutput = deserializeOutput(spanData);
+    const integrationFixtures = extractIntegrationFixtures(spanData);
     if (adaptInputs) {
       const originalMetadata = serverItem.originalMetadata;
       inputs = adaptInputs(inputs, {
@@ -3781,6 +3833,7 @@ async function processItem(httpClient, serverItem, fn, testRunId, mockStrategy,
           mockStrategy,
           mockOverrides: hasOverrides ? resolvedOverrides : void 0,
           fetchSpanOutput,
+          integrationFixtures,
           dbBranchLease: lease,
           dbBranchTimings
         },
@@ -4261,6 +4314,7 @@ var init_replay = __esm({
     init_codeChange();
     init_errors();
     init_http();
+    init_integrationFixture();
     init_mockOverride();
     init_randomUuid();
     init_replayContext();
@@ -4304,6 +4358,7 @@ __export(index_exports, {
   BitfabLangGraphIntegration: () => BitfabLangGraphIntegration,
   BitfabOpenAIAgentHandler: () => BitfabOpenAIAgentHandler,
   BitfabOpenAITracingProcessor: () => BitfabOpenAITracingProcessor,
+  BitfabRivetIntegration: () => BitfabRivetIntegration,
   BitfabVercelAiHandler: () => BitfabVercelAiHandler,
   DEFAULT_SERVICE_URL: () => DEFAULT_SERVICE_URL,
   DatasetsClient: () => DatasetsClient,
@@ -6978,6 +7033,533 @@ _context = new WeakMap();
 // src/client.ts
 init_replayContext();
 
+// src/rivetIntegration.ts
+init_errors();
+init_integrationFixture();
+init_replayContext();
+
+// src/rivetContext.ts
+init_asyncStorage();
+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 isRecord5(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(replay2, traceFunctionKey, effect, marked) {
+  const counterKey = `${traceFunctionKey}:${effect}`;
+  const callIndex = replay2.callCounters?.get(counterKey) ?? 0;
+  const mockSpan = replay2.mockTree?.spans.get(`${counterKey}:${callIndex}`);
+  const hasOverride = replay2.mockOverrides?.some(
+    (override) => override.match(
+      effectNode(traceFunctionKey, effect, mockSpan?.sourceSpanId)
+    )
+  ) ?? false;
+  return {
+    willMock: hasOverride || replay2.mockStrategy === "all" || replay2.mockStrategy === "marked" && marked,
+    hasRecordedSpan: mockSpan !== void 0,
+    hasOverride,
+    callIndex
+  };
+}
+function assertReplayEffectAvailable(replay2, traceFunctionKey, effect, marked) {
+  if (!replay2) {
+    return;
+  }
+  const status = replayEffectStatus(replay2, 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, replay2) {
+  const fixture = replay2.integrationFixtures?.get(fixtureName(traceFunctionKey));
+  if (!isRecord5(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 (!isRecord5(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 replay2 = getReplayContext();
+    if (replay2) {
+      const existing = this.replayProxies.get(replay2);
+      if (existing) {
+        return existing;
+      }
+      const proxy = this.createContextProxy(
+        snapshotFromReplay(this.traceFunctionKey, replay2),
+        void 0
+      );
+      this.replayProxies.set(replay2, 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 replay2 = getReplayContext();
+    const dbProxy = this.createDbProxy(binding, replay2);
+    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, replay2));
+      }
+    }
+    const topLevelEffects = /* @__PURE__ */ new Map();
+    for (const effect of RIVET_TOP_LEVEL_EFFECTS) {
+      topLevelEffects.set(
+        effect,
+        this.createEffect(
+          effect,
+          () => binding?.context,
+          effect,
+          replay2
+        )
+      );
+    }
+    const replayTarget = replay2 ? { ...target } : target;
+    if (replay2) {
+      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" && replay2) {
+          return (_preventSleep) => void 0;
+        }
+        if (replay2 && 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, replay2) {
+    const proxy = this.createGroupProxy("db", binding, replay2);
+    return new Proxy(proxy, {
+      get: (source, property) => property === "transaction" ? this.createTransaction(binding, replay2) : Reflect.get(source, property, source)
+    });
+  }
+  createGroupProxy(group, binding, replay2) {
+    const targetProvider = () => {
+      if (!binding) {
+        return void 0;
+      }
+      const context = binding.context;
+      const target = readProperty(context, group);
+      return isRecord5(target) ? target : void 0;
+    };
+    const proxy = this.createMethodProxy(
+      targetProvider,
+      group,
+      RIVET_GROUP_EFFECT_METHODS[group],
+      replay2
+    );
+    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, replay2) {
+    const effects = /* @__PURE__ */ new Map();
+    for (const method of methods) {
+      const effect = `${prefix}.${method}`;
+      effects.set(
+        method,
+        this.createEffect(effect, targetProvider, method, replay2)
+      );
+    }
+    return new Proxy(
+      {},
+      {
+        get: (_source, property) => {
+          if (typeof property === "string") {
+            const effect = effects.get(property);
+            if (effect) {
+              return effect;
+            }
+            if (replay2) {
+              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, replay2) {
+    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,
+      replay2 ? () => void 0 : void 0
+    ) : this.runtime.wrapEffect(effect, marked, async (...args) => {
+      return await invoke(...args);
+    });
+    return (...args) => {
+      if (replay2) {
+        assertReplayEffectAvailable(
+          replay2,
+          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, replay2)
+      );
+    };
+  }
+  createTransaction(binding, replay2) {
+    return async (callback, options) => {
+      if (typeof callback !== "function") {
+        throw new BitfabError("Rivet db.transaction requires a callback");
+      }
+      if (replay2) {
+        return await callback(
+          this.createMethodProxy(
+            () => void 0,
+            "db.transaction",
+            ["execute"],
+            replay2
+          )
+        );
+      }
+      const context = binding?.context;
+      const db = context ? readProperty(context, "db") : void 0;
+      const transaction = isRecord5(db) ? readProperty(db, "transaction") : void 0;
+      if (!isRecord5(db) || typeof transaction !== "function") {
+        throw unavailableCapability("db.transaction");
+      }
+      return await transaction.call(
+        db,
+        async (transactionDb) => {
+          if (!isRecord5(transactionDb)) {
+            throw unavailableCapability("db.transaction.execute");
+          }
+          return await callback(
+            this.createMethodProxy(
+              () => transactionDb,
+              "db.transaction",
+              ["execute"],
+              replay2
+            )
+          );
+        },
+        options
+      );
+    };
+  }
+  decorateQueueMessages(value, effect, args, replay2) {
+    if (!this.queueCallIsCompletable(effect, args)) {
+      return value;
+    }
+    if (Array.isArray(value)) {
+      return value.map((message) => this.decorateQueueMessage(message, replay2));
+    }
+    return this.decorateQueueMessage(value, replay2);
+  }
+  decorateQueueMessage(value, replay2) {
+    if (!isRecord5(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",
+      replay2
+    );
+    return { ...value, complete };
+  }
+  queueCallIsCompletable(effect, args) {
+    const options = effect === "queue.waitForNames" ? args[1] : args[0];
+    return isRecord5(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
 init_asyncStorage();
 var seedContextStorage = null;
@@ -8687,6 +9269,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.
    *
@@ -9185,6 +9808,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 = {
@@ -9241,7 +9865,7 @@ var Bitfab = class {
                 }
                 if (shouldMockWithBaseStrategy) {
                   const output = await resolveMockOrAbort(
-                    () => resolveRecordedOutput()
+                    () => resolveBaseOutput()
                   );
                   void sendSpan({
                     result: output,
@@ -9264,7 +9888,7 @@ var Bitfab = class {
             );
           }
           if (shouldMockWithBaseStrategy) {
-            const recorded = resolveRecordedOutput();
+            const recorded = resolveBaseOutput();
             if (recorded instanceof Promise) {
               return emitMockAsync(recorded, "recorded");
             }
@@ -9866,6 +10490,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.
    *
@@ -9963,6 +10609,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.
@@ -10124,6 +10784,7 @@ function resolveTraceFunctionKey(registration) {
   BitfabLangGraphIntegration,
   BitfabOpenAIAgentHandler,
   BitfabOpenAITracingProcessor,
+  BitfabRivetIntegration,
   BitfabVercelAiHandler,
   DEFAULT_SERVICE_URL,
   DatasetsClient,