@bitfab/sdk 0.28.11 → 0.29.1

package/dist/index.cjs

@@ -282,6 +282,22 @@ var init_asyncStorage = __esm({
   }
 });
 
+// src/mockOverride.ts
+function resolveMockValue(value, ctx) {
+  return typeof value === "function" ? value(ctx) : value;
+}
+function normalizeMockOverrides(mockOverride) {
+  if (mockOverride === void 0) {
+    return [];
+  }
+  return Array.isArray(mockOverride) ? mockOverride : [mockOverride];
+}
+var init_mockOverride = __esm({
+  "src/mockOverride.ts"() {
+    "use strict";
+  }
+});
+
 // src/replayContext.ts
 function getReplayContext() {
   return replayContextStorage?.getStore() ?? null;
@@ -357,6 +373,7 @@ function buildMockTree(rootNode) {
       counters.set(counterKey, index + 1);
       spans.set(`${counterKey}:${index}`, {
         sourceSpanId: node.sourceSpanId,
+        externalSpanId: node.externalSpanId,
         output: node.output,
         outputMeta: node.outputMeta
       });
@@ -370,57 +387,79 @@ function buildMockTree(rootNode) {
   }
   return { spans };
 }
-async function processItem(httpClient, serverItem, fn, testRunId, mockStrategy, environment, adaptInputs) {
+async function processItem(httpClient, serverItem, fn, testRunId, mockStrategy, resolvedOverrides, replayedTraceId, environment, adaptInputs) {
   const lease = environment ? serverItem.dbBranchLease : void 0;
   let inputs = [];
   let originalOutput;
   let result;
   let error = null;
-  const replayedTraceId = randomUuid();
   const pendingPersistence = [];
+  const originalTraceId = serverItem.originalTraceId ?? serverItem.sourceTraceId;
+  const originalSpanId = serverItem.originalSpanId ?? serverItem.sourceSpanId;
   try {
-    const span = await httpClient.getExternalSpan(serverItem.externalSpanId);
+    const span = await httpClient.getExternalSpan(originalSpanId);
     const spanData = span.rawData?.span_data ?? {};
     inputs = deserializeInputs(spanData);
     originalOutput = deserializeOutput(spanData);
     if (adaptInputs) {
       inputs = adaptInputs(inputs, {
-        traceId: serverItem.traceId,
-        sourceSpanId: serverItem.externalSpanId
+        originalTraceId,
+        originalSpanId,
+        // Deprecated aliases for originalTraceId/originalSpanId.
+        sourceTraceId: originalTraceId,
+        sourceSpanId: originalSpanId
       });
     }
+    const hasOverrides = resolvedOverrides.length > 0;
+    const needTree = mockStrategy === "all" || mockStrategy === "marked" || hasOverrides;
+    const includeOutputs = mockStrategy === "all";
     let mockTree;
-    if (mockStrategy === "all" || mockStrategy === "marked") {
+    if (needTree) {
       try {
-        const treeResponse = await httpClient.getSpanTree(
-          serverItem.externalSpanId
-        );
+        const treeResponse = await httpClient.getSpanTree(originalSpanId, {
+          includeOutputs
+        });
         if (treeResponse.root) {
           mockTree = buildMockTree(treeResponse.root);
-        } else if (mockStrategy === "all") {
+        } else if (mockStrategy === "all" || hasOverrides) {
           throw new BitfabError(
-            `Replay mock strategy "all" requires a span tree root for source span ${serverItem.externalSpanId}.`
+            `Replay mock strategy "${mockStrategy}"${hasOverrides ? " with overrides" : ""} requires a span tree root for original span ${originalSpanId}.`
           );
         } else {
           mockTree = void 0;
         }
       } catch (e) {
-        if (mockStrategy === "all") {
+        if (mockStrategy === "all" || hasOverrides) {
           throw e;
         }
         mockTree = void 0;
       }
     }
+    const outputCache = /* @__PURE__ */ new Map();
+    const fetchSpanOutput = mockTree && !includeOutputs ? (externalSpanId) => {
+      let pending = outputCache.get(externalSpanId);
+      if (!pending) {
+        pending = httpClient.getExternalSpan(externalSpanId).then(
+          (s) => deserializeOutput(
+            s.rawData?.span_data ?? {}
+          )
+        );
+        outputCache.set(externalSpanId, pending);
+      }
+      return pending;
+    } : void 0;
     const maybePromise = runWithReplayContext(
       {
         testRunId,
         traceId: replayedTraceId,
         inputSourceSpanId: span.id,
         inputSourceTraceId: span.externalTraceId,
-        sourceBitfabTraceId: serverItem.traceId,
+        sourceBitfabTraceId: originalTraceId,
         mockTree,
         callCounters: mockTree ? /* @__PURE__ */ new Map() : void 0,
         mockStrategy,
+        mockOverrides: hasOverrides ? resolvedOverrides : void 0,
+        fetchSpanOutput,
         dbBranchLease: lease,
         pendingPersistence
       },
@@ -445,7 +484,15 @@ async function processItem(httpClient, serverItem, fn, testRunId, mockStrategy,
     }
   }
   return {
-    traceId: replayedTraceId,
+    // Written in by replay() from the complete-replay response once the server
+    // has minted this replay trace's row. Null until then: the client-side
+    // correlation id (replayedTraceId) is never surfaced as the item's traceId.
+    traceId: null,
+    originalTraceId,
+    originalSpanId,
+    // Deprecated aliases for originalTraceId/originalSpanId.
+    sourceTraceId: originalTraceId,
+    sourceSpanId: originalSpanId,
     input: inputs,
     result,
     originalOutput,
@@ -477,7 +524,7 @@ async function mapWithConcurrency(tasks, maxConcurrency, onSettled) {
   await Promise.all(workers);
   return results;
 }
-async function replay(httpClient, serviceUrl, traceFunctionKey, fn, options) {
+async function replay(httpClient, serviceUrl, traceFunctionKey, fn, options, registeredOverrides = []) {
   if (options?.traceIds !== void 0) {
     if (options.traceIds.length === 0) {
       throw new BitfabError("traceIds must contain at least one trace ID.");
@@ -517,13 +564,20 @@ async function replay(httpClient, serviceUrl, traceFunctionKey, fn, options) {
   );
   const mockStrategy = options?.mock ?? "marked";
   const maxConcurrency = options?.maxConcurrency ?? 10;
+  const resolvedOverrides = [
+    ...normalizeMockOverrides(options?.mockOverride),
+    ...registeredOverrides
+  ];
+  const replayedTraceIds = serverItems.map(() => randomUuid());
   const tasks = serverItems.map(
-    (serverItem) => () => processItem(
+    (serverItem, index) => () => processItem(
       httpClient,
       serverItem,
       fn,
       testRunId,
       mockStrategy,
+      resolvedOverrides,
+      replayedTraceIds[index],
       options?.environment,
       options?.adaptInputs
     )
@@ -550,11 +604,17 @@ async function replay(httpClient, serviceUrl, traceFunctionKey, fn, options) {
           succeeded,
           errored,
           item: {
-            // Source (historical) trace id, so a UI can identify the trace
-            // that just settled. The item's own traceId is the new replay
-            // trace and is assigned later (below), so use the server item.
-            traceId: serverItems[index]?.traceId ?? null,
-            replayTraceId: item.traceId,
+            // The server replay trace id isn't known until completeReplay
+            // runs (below), so it can't be reported mid-run and we never
+            // emit the client-side placeholder. originalTraceId (the
+            // historical trace) is known now and is what a UI keys on to
+            // identify what just settled.
+            traceId: null,
+            originalTraceId: item.originalTraceId ?? null,
+            originalSpanId: item.originalSpanId ?? null,
+            // Deprecated aliases for originalTraceId/originalSpanId.
+            sourceTraceId: item.originalTraceId ?? null,
+            sourceSpanId: item.originalSpanId ?? null,
             input: item.input,
             result: item.result,
             originalOutput: item.originalOutput,
@@ -572,56 +632,58 @@ async function replay(httpClient, serviceUrl, traceFunctionKey, fn, options) {
   const completeResult = await httpClient.completeReplay(testRunId);
   const serverTraceIds = completeResult.traceIds;
   const replayTokens = completeResult.tokens;
-  if (serverTraceIds === void 0) {
-    try {
-      console.warn(
-        "Bitfab: server did not return replay trace IDs; item.traceId will be null (server upgrade required for verdict persistence)"
-      );
-    } catch {
-    }
-    for (const item of resultItems) {
-      item.traceId = null;
-    }
-  } else {
+  if (serverTraceIds !== void 0) {
     const missing = [];
     let completedCount = 0;
-    for (const item of resultItems) {
-      if (item.traceId) {
-        const mapped = serverTraceIds[item.traceId];
-        if (item.error === null) {
-          completedCount += 1;
-          if (mapped === void 0) {
-            missing.push(item.traceId);
-          }
-        }
-        if (mapped !== void 0) {
-          item.tokens = replayTokens?.[mapped] ?? null;
+    for (let index = 0; index < resultItems.length; index += 1) {
+      const item = resultItems[index];
+      const localId = replayedTraceIds[index];
+      const mapped = localId ? serverTraceIds[localId] : void 0;
+      item.traceId = mapped ?? null;
+      if (item.error === null) {
+        completedCount += 1;
+        if (mapped === void 0) {
+          missing.push(localId ?? item.originalTraceId);
         }
-        item.traceId = mapped ?? null;
+      }
+      if (mapped !== void 0) {
+        item.tokens = replayTokens?.[mapped] ?? null;
       }
     }
-    if (missing.length > 0) {
+    if (completedCount > 0 && missing.length === completedCount) {
       const serverCount = completeResult.traceCount !== void 0 ? ` The server persisted ${completeResult.traceCount} trace(s) for this run.` : "";
-      if (missing.length === completedCount) {
-        throw new BitfabError(
-          `Replay completed but the server has no persisted trace for any of the ${completedCount} completed item(s) (testRunId ${testRunId}).${serverCount} Trace uploads were awaited, so either the uploads failed (check for "Bitfab: Failed to create" errors above) or the replayed function is not wrapped with withSpan.`
-        );
-      }
+      throw new BitfabError(
+        `Replay completed but the server has no persisted trace for any of the ${completedCount} completed item(s) (testRunId ${testRunId}).${serverCount} Trace uploads were awaited, so either the uploads failed (check for "Bitfab: Failed to create" errors above) or the replayed function is not wrapped with withSpan.`
+      );
+    }
+    if (missing.length > 0) {
       try {
         console.error(
-          `Bitfab: server has no persisted trace for ${missing.length} of ${completedCount} completed replay item(s) (testRunId ${testRunId}).${serverCount} Their traceId is null and verdicts cannot be persisted for them. Missing: ${missing.join(", ")}`
+          `Bitfab: server has no persisted trace for ${missing.length} of ${completedCount} completed replay item(s) (testRunId ${testRunId}). Their replay token usage is unavailable and they cannot be labeled.`
         );
       } catch {
       }
     }
   }
-  const replayResult = {
+  const result = {
     items: resultItems,
     testRunId,
     testRunUrl: `${serviceUrl}${testRunUrl}`
   };
-  await writeReplayResultFile(replayResult);
-  return replayResult;
+  await writeReplayResultFile(result);
+  try {
+    options?.onProgress?.({
+      type: "complete",
+      testRunId,
+      completed: total,
+      total,
+      succeeded,
+      errored,
+      result
+    });
+  } catch {
+  }
+  return result;
 }
 async function writeReplayResultFile(result) {
   const resultPath = typeof process !== "undefined" ? process.env?.BITFAB_REPLAY_RESULT_PATH : void 0;
@@ -650,6 +712,7 @@ var init_replay = __esm({
   "src/replay.ts"() {
     "use strict";
     init_errors();
+    init_mockOverride();
     init_randomUuid();
     init_replayContext();
     init_serialize();
@@ -683,7 +746,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // src/version.generated.ts
-var __version__ = "0.28.11";
+var __version__ = "0.29.1";
 
 // src/constants.ts
 var DEFAULT_SERVICE_URL = "https://bitfab.ai";
@@ -1112,9 +1175,14 @@ var HttpClient = class {
   /**
    * Fetch the span tree for a root span.
    * Blocking GET request.
+   *
+   * Pass `includeOutputs: false` for a payload-free tree (structure +
+   * `externalSpanId` only), so recorded outputs are fetched lazily per mocked
+   * span instead of all up front. Omit it (default eager) for `mock: "all"`.
    */
-  async getSpanTree(externalSpanId) {
-    const url = `${this.serviceUrl}/api/sdk/replay/spanTree/${externalSpanId}`;
+  async getSpanTree(externalSpanId, options) {
+    const query = options?.includeOutputs === false ? "?includeOutputs=false" : "";
+    const url = `${this.serviceUrl}/api/sdk/replay/spanTree/${externalSpanId}${query}`;
     const controller = new AbortController();
     const timeoutId = setTimeout(() => controller.abort(), 3e4);
     try {
@@ -2700,6 +2768,9 @@ var BitfabLangGraphCallbackHandler = class {
   }
 };
 
+// src/client.ts
+init_mockOverride();
+
 // src/openaiAgentSdk.ts
 var BitfabOpenAIAgentHandler = class {
   constructor(config) {
@@ -3515,6 +3586,12 @@ var Bitfab = class {
   constructor(config) {
     /** Gate the empty-key warning to fire at most once. */
     this.apiKeyWarned = false;
+    /**
+     * Mock overrides registered via {@link Bitfab.registerMockOverride}, applied
+     * to every `replay` on this client (after any per-call `mockOverride`). In
+     * registration order; first matcher wins within this list.
+     */
+    this.mockOverrides = [];
     this.apiKeyConfig = config.apiKey;
     this.serviceUrl = config.serviceUrl ?? DEFAULT_SERVICE_URL;
     this.timeout = config.timeout ?? 12e4;
@@ -4099,7 +4176,7 @@ var Bitfab = class {
                   dbSnapshotUsage: {
                     neonBranchId: replayCtx.dbBranchLease.neonBranchId,
                     snapshotTimestamp: replayCtx.dbBranchLease.snapshotTimestamp,
-                    sourceTraceId: replayCtx.sourceBitfabTraceId,
+                    originalTraceId: replayCtx.sourceBitfabTraceId,
                     accessed: replayCtx.dbSnapshotAccessed === true
                   }
                 }
@@ -4127,24 +4204,77 @@ var Bitfab = class {
           const counterKey = `${traceFunctionKey}:${baseSpanParams.spanName}`;
           const callIndex = counters.get(counterKey) ?? 0;
           counters.set(counterKey, callIndex + 1);
-          const shouldMock = replayCtxForMock.mockStrategy === "all" || replayCtxForMock.mockStrategy === "marked" && options.mockOnReplay === true;
-          if (shouldMock) {
-            const mockKey = `${counterKey}:${callIndex}`;
-            const mockSpan = replayCtxForMock.mockTree.spans.get(mockKey);
-            if (mockSpan) {
-              let output = mockSpan.output;
-              if (mockSpan.outputMeta !== void 0 && mockSpan.outputMeta !== null) {
-                output = deserializeValue({
-                  json: mockSpan.output,
-                  meta: mockSpan.outputMeta
-                });
-              }
+          const mockKey = `${counterKey}:${callIndex}`;
+          const mockSpan = replayCtxForMock.mockTree.spans.get(mockKey);
+          const emitMock = (output) => {
+            void sendSpan({ result: output, mocked: true });
+            if (fnReturnsPromise) {
+              return Promise.resolve(output);
+            }
+            return output;
+          };
+          const emitMockAsync = (pending) => {
+            if (!fnReturnsPromise) {
+              throw new BitfabError(
+                `Cannot mock synchronous span "${traceFunctionKey}" with an asynchronously-resolved value (lazy recorded-output fetch or an async value function). Make the wrapped function async, or use mock: "all" so recorded outputs are fetched eagerly.`
+              );
+            }
+            return (async () => {
+              const output = await pending;
               void sendSpan({ result: output, mocked: true });
-              if (fnReturnsPromise) {
-                return Promise.resolve(output);
-              }
               return output;
+            })();
+          };
+          const resolveRecordedOutput = () => {
+            const hasInlineOutput = mockSpan?.output !== void 0 || mockSpan?.outputMeta !== void 0;
+            if (!hasInlineOutput && replayCtxForMock.fetchSpanOutput && mockSpan?.externalSpanId) {
+              return replayCtxForMock.fetchSpanOutput(mockSpan.externalSpanId);
+            }
+            if (!mockSpan) {
+              return Promise.reject(
+                new BitfabError(
+                  `No recorded span to source output for "${traceFunctionKey}".`
+                )
+              );
             }
+            let output = mockSpan.output;
+            if (mockSpan.outputMeta !== void 0 && mockSpan.outputMeta !== null) {
+              output = deserializeValue({
+                json: mockSpan.output,
+                meta: mockSpan.outputMeta
+              });
+            }
+            return output;
+          };
+          if (replayCtxForMock.mockOverrides?.length) {
+            const nodeMeta = {
+              traceFunctionKey,
+              spanName: baseSpanParams.spanName,
+              type: options.type ?? "custom",
+              originalSpanId: mockSpan?.sourceSpanId
+            };
+            const override = replayCtxForMock.mockOverrides.find(
+              (o) => o.match(nodeMeta)
+            );
+            if (override) {
+              const injected = resolveMockValue(override.value, {
+                node: nodeMeta,
+                inputs: args,
+                getOriginalOutput: () => Promise.resolve(resolveRecordedOutput())
+              });
+              if (injected instanceof Promise) {
+                return emitMockAsync(injected);
+              }
+              return emitMock(injected);
+            }
+          }
+          const shouldMock = replayCtxForMock.mockStrategy === "all" || replayCtxForMock.mockStrategy === "marked" && options.mockOnReplay === true;
+          if (shouldMock && mockSpan) {
+            const recorded = resolveRecordedOutput();
+            if (recorded instanceof Promise) {
+              return emitMockAsync(recorded);
+            }
+            return emitMock(recorded);
           }
         }
         const recordSpan = (result) => {
@@ -4350,8 +4480,11 @@ var Bitfab = class {
         ...params.dbSnapshotUsage.snapshotTimestamp && {
           snapshot_timestamp: params.dbSnapshotUsage.snapshotTimestamp
         },
-        ...params.dbSnapshotUsage.sourceTraceId && {
-          source_trace_id: params.dbSnapshotUsage.sourceTraceId
+        ...params.dbSnapshotUsage.originalTraceId && {
+          original_trace_id: params.dbSnapshotUsage.originalTraceId,
+          // Deprecated wire alias, kept so this SDK still reports usage
+          // against servers that predate the rename.
+          source_trace_id: params.dbSnapshotUsage.originalTraceId
         },
         accessed: params.dbSnapshotUsage.accessed
       };
@@ -4424,26 +4557,14 @@ var Bitfab = class {
       ...params.mocked && { mocked: true }
     });
   }
-  /**
-   * Replay historical traces through a function and create a test run.
-   *
-   * Fetches the last N traces for the given trace function key, re-runs each
-   * through the provided function, and returns comparison data.
-   *
-   * Accepts either a `withSpan`-wrapped function (under the same key) or any
-   * plain callable: plain callables are wrapped internally so each replayed
-   * invocation records a trace tied to the test run. The plain-callable form
-   * is how handler-instrumented workflows (LangGraph/LangChain, Claude Agent
-   * SDK) replay - those record traces under a key with no `withSpan`-wrapped
-   * root in the app.
-   *
-   * @param traceFunctionKey - The trace function key to replay
-   * @param fn - The function to run recorded inputs through
-   * @param options - Optional replay options. When `traceIds` is passed,
-   *   `limit` is ignored (with a warning): an explicit ID list already
-   *   determines how many traces replay.
-   * @returns ReplayResult with items, testRunId, and testRunUrl
-   */
+  registerMockOverride(overrideOrMatch, value) {
+    const override = typeof overrideOrMatch === "function" ? { match: overrideOrMatch, value } : overrideOrMatch;
+    this.mockOverrides.push(override);
+  }
+  /** Remove all overrides registered via {@link registerMockOverride}. */
+  clearMockOverrides() {
+    this.mockOverrides.length = 0;
+  }
   async replay(traceFunctionKey, fn, options) {
     const wrappedKey = fn._bitfabTraceFunctionKey;
     let replayFn = fn;
@@ -4464,7 +4585,8 @@ var Bitfab = class {
       this.serviceUrl,
       traceFunctionKey,
       replayFn,
-      options
+      options,
+      this.mockOverrides
     );
   }
 };