@bitfab/sdk 0.17.0 → 0.18.1

package/dist/index.d.cts

@@ -225,7 +225,7 @@ interface CodeChangeFile {
  * LangGraph/LangChain callback handler for Bitfab tracing.
  *
  * Hooks into LangGraph's callback system to capture graph node execution,
- * LLM calls, and tool invocations as Bitfab spans — without requiring users
+ * LLM calls, and tool invocations as Bitfab spans, without requiring users
  * to wrap their functions with withSpan (which fails on non-serializable args).
  *
  * Duck-typed to match LangChain.js's BaseCallbackHandler interface.
@@ -238,8 +238,8 @@ interface ActiveSpanContext$1 {
 /**
  * LangChain/LangGraph callback handler that sends traces to Bitfab.
  *
- * Duck-typed to match LangChain.js's BaseCallbackHandler — no
- * `@langchain/core` dependency required. Pass as a callback:
+ * Duck-typed to match LangChain.js's BaseCallbackHandler, so no
+ * `@langchain/core` dependency is required. Pass as a callback:
  *
  * ```typescript
  * const handler = bitfab.getLangGraphCallbackHandler("my-agent");
@@ -978,11 +978,15 @@ declare class Bitfab {
      * Fetches the last N traces for the given trace function key, re-runs each
      * through the provided function, and returns comparison data.
      *
-     * The function must have been wrapped with `withSpan` — replay injects
-     * `testRunId` via async context so new spans are linked to the test run.
+     * 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 replay (must be the return value of `withSpan`)
+     * @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.
@@ -1051,7 +1055,7 @@ declare class BitfabFunction {
 /**
  * SDK version from package.json (injected at build time)
  */
-declare const __version__ = "0.17.0";
+declare const __version__ = "0.18.1";
 
 /**
  * Constants for the Bitfab SDK.

package/dist/index.d.ts

@@ -225,7 +225,7 @@ interface CodeChangeFile {
  * LangGraph/LangChain callback handler for Bitfab tracing.
  *
  * Hooks into LangGraph's callback system to capture graph node execution,
- * LLM calls, and tool invocations as Bitfab spans — without requiring users
+ * LLM calls, and tool invocations as Bitfab spans, without requiring users
  * to wrap their functions with withSpan (which fails on non-serializable args).
  *
  * Duck-typed to match LangChain.js's BaseCallbackHandler interface.
@@ -238,8 +238,8 @@ interface ActiveSpanContext$1 {
 /**
  * LangChain/LangGraph callback handler that sends traces to Bitfab.
  *
- * Duck-typed to match LangChain.js's BaseCallbackHandler — no
- * `@langchain/core` dependency required. Pass as a callback:
+ * Duck-typed to match LangChain.js's BaseCallbackHandler, so no
+ * `@langchain/core` dependency is required. Pass as a callback:
  *
  * ```typescript
  * const handler = bitfab.getLangGraphCallbackHandler("my-agent");
@@ -978,11 +978,15 @@ declare class Bitfab {
      * Fetches the last N traces for the given trace function key, re-runs each
      * through the provided function, and returns comparison data.
      *
-     * The function must have been wrapped with `withSpan` — replay injects
-     * `testRunId` via async context so new spans are linked to the test run.
+     * 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 replay (must be the return value of `withSpan`)
+     * @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.
@@ -1051,7 +1055,7 @@ declare class BitfabFunction {
 /**
  * SDK version from package.json (injected at build time)
  */
-declare const __version__ = "0.17.0";
+declare const __version__ = "0.18.1";
 
 /**
  * Constants for the Bitfab SDK.

package/dist/index.js

@@ -11,7 +11,7 @@ import {
   flushTraces,
   getCurrentSpan,
   getCurrentTrace
-} from "./chunk-M6N633CX.js";
+} from "./chunk-ILIUTS5D.js";
 import {
   BitfabError
 } from "./chunk-QT7HWOKU.js";

package/dist/node.cjs

@@ -447,7 +447,7 @@ registerAsyncLocalStorageClass(
 );
 
 // src/version.generated.ts
-var __version__ = "0.17.0";
+var __version__ = "0.18.1";
 
 // src/constants.ts
 var DEFAULT_SERVICE_URL = "https://bitfab.ai";
@@ -1677,21 +1677,131 @@ function extractModelName(serialized, metadata) {
   }
   return void 0;
 }
+function asTokenCount(value) {
+  return typeof value === "number" && Number.isFinite(value) ? value : null;
+}
+function normalizeTokenUsage(raw) {
+  if (typeof raw !== "object" || raw === null || Array.isArray(raw)) {
+    return null;
+  }
+  const u = raw;
+  if ("cache_read_input_tokens" in u || "cache_creation_input_tokens" in u) {
+    const cacheRead = asTokenCount(u.cache_read_input_tokens);
+    const cacheCreation = asTokenCount(u.cache_creation_input_tokens);
+    const baseInput = asTokenCount(u.input_tokens);
+    const outputTokens = asTokenCount(u.output_tokens);
+    if (cacheRead === null && cacheCreation === null && baseInput === null && outputTokens === null) {
+      return null;
+    }
+    const inputTokens = (baseInput ?? 0) + (cacheRead ?? 0) + (cacheCreation ?? 0);
+    return {
+      inputTokens,
+      outputTokens,
+      totalTokens: inputTokens + (outputTokens ?? 0),
+      cachedInputTokens: cacheRead
+    };
+  }
+  if ("prompt_tokens" in u || "completion_tokens" in u || "promptTokens" in u || "completionTokens" in u) {
+    const promptDetails = u.prompt_tokens_details ?? {};
+    return withAnyTokenCount({
+      inputTokens: asTokenCount(u.prompt_tokens) ?? asTokenCount(u.promptTokens),
+      outputTokens: asTokenCount(u.completion_tokens) ?? asTokenCount(u.completionTokens),
+      totalTokens: asTokenCount(u.total_tokens) ?? asTokenCount(u.totalTokens),
+      cachedInputTokens: asTokenCount(promptDetails.cached_tokens)
+    });
+  }
+  if ("prompt_token_count" in u || "candidates_token_count" in u) {
+    return withAnyTokenCount({
+      inputTokens: asTokenCount(u.prompt_token_count),
+      outputTokens: asTokenCount(u.candidates_token_count),
+      totalTokens: asTokenCount(u.total_token_count),
+      cachedInputTokens: asTokenCount(u.cached_content_token_count)
+    });
+  }
+  if ("input_tokens" in u || "output_tokens" in u) {
+    const inputDetails = u.input_token_details ?? {};
+    const inputTokens = asTokenCount(u.input_tokens);
+    const outputTokens = asTokenCount(u.output_tokens);
+    let totalTokens = asTokenCount(u.total_tokens);
+    if (totalTokens === null && inputTokens !== null && outputTokens !== null) {
+      totalTokens = inputTokens + outputTokens;
+    }
+    return withAnyTokenCount({
+      inputTokens,
+      outputTokens,
+      totalTokens,
+      cachedInputTokens: asTokenCount(inputDetails.cache_read)
+    });
+  }
+  return null;
+}
+function withAnyTokenCount(usage) {
+  const hasCount = usage.inputTokens !== null || usage.outputTokens !== null || usage.totalTokens !== null || usage.cachedInputTokens !== null;
+  return hasCount ? usage : null;
+}
+function addUsage(totals, usage) {
+  for (const key of [
+    "inputTokens",
+    "outputTokens",
+    "totalTokens",
+    "cachedInputTokens"
+  ]) {
+    const value = usage[key];
+    if (value !== null) {
+      totals[key] = (totals[key] ?? 0) + value;
+    }
+  }
+}
+function usageFromGenerations(generations) {
+  if (!generations?.length) {
+    return null;
+  }
+  const totals = {
+    inputTokens: null,
+    outputTokens: null,
+    totalTokens: null,
+    cachedInputTokens: null
+  };
+  let found = false;
+  for (const batch of generations) {
+    if (!Array.isArray(batch)) {
+      continue;
+    }
+    for (const gen of batch) {
+      const msg = gen?.message;
+      if (!msg || typeof msg !== "object") {
+        continue;
+      }
+      const responseMetadata = msg.response_metadata;
+      const usage = normalizeTokenUsage(msg.usage_metadata) ?? normalizeTokenUsage(responseMetadata?.token_usage) ?? normalizeTokenUsage(responseMetadata?.usage) ?? normalizeTokenUsage(responseMetadata?.tokenUsage);
+      if (!usage) {
+        continue;
+      }
+      found = true;
+      addUsage(totals, usage);
+    }
+  }
+  return found ? totals : null;
+}
 function extractUsage2(output) {
+  const generations = output.generations;
+  const llmOutput = output.llmOutput ?? output.llm_output;
+  const normalized = usageFromGenerations(generations) ?? normalizeTokenUsage(llmOutput?.tokenUsage) ?? normalizeTokenUsage(llmOutput?.token_usage) ?? normalizeTokenUsage(llmOutput?.usage);
   const usage = {};
-  const llmOutput = output.llmOutput;
-  const tokenUsage = llmOutput?.tokenUsage ?? llmOutput?.token_usage ?? llmOutput?.usage ?? {};
-  const inputTokens = tokenUsage.promptTokens ?? tokenUsage.prompt_tokens ?? tokenUsage.input_tokens;
-  const outputTokens = tokenUsage.completionTokens ?? tokenUsage.completion_tokens ?? tokenUsage.output_tokens;
-  const totalTokens = tokenUsage.totalTokens ?? tokenUsage.total_tokens;
-  if (inputTokens !== void 0 && inputTokens !== null) {
-    usage.inputTokens = inputTokens;
+  if (!normalized) {
+    return usage;
+  }
+  if (normalized.inputTokens !== null) {
+    usage.inputTokens = normalized.inputTokens;
+  }
+  if (normalized.outputTokens !== null) {
+    usage.outputTokens = normalized.outputTokens;
   }
-  if (outputTokens !== void 0 && outputTokens !== null) {
-    usage.outputTokens = outputTokens;
+  if (normalized.totalTokens !== null) {
+    usage.totalTokens = normalized.totalTokens;
   }
-  if (totalTokens !== void 0 && totalTokens !== null) {
-    usage.totalTokens = totalTokens;
+  if (normalized.cachedInputTokens !== null) {
+    usage.cachedInputTokens = normalized.cachedInputTokens;
   }
   return usage;
 }
@@ -3079,6 +3189,9 @@ var Bitfab = class {
       };
       return runWithSpanStack(newStack, executeWithContext);
     };
+    Object.defineProperty(wrappedFn, "_bitfabTraceFunctionKey", {
+      value: traceFunctionKey
+    });
     return wrappedFn;
   }
   /**
@@ -3250,23 +3363,40 @@ var Bitfab = class {
    * Fetches the last N traces for the given trace function key, re-runs each
    * through the provided function, and returns comparison data.
    *
-   * The function must have been wrapped with `withSpan` — replay injects
-   * `testRunId` via async context so new spans are linked to the test run.
+   * 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 replay (must be the return value of `withSpan`)
+   * @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
    */
   async replay(traceFunctionKey, fn, options) {
+    const wrappedKey = fn._bitfabTraceFunctionKey;
+    let replayFn = fn;
+    if (wrappedKey === void 0) {
+      replayFn = this.withSpan(
+        traceFunctionKey,
+        { name: fn.name || "Replay", type: "agent" },
+        fn
+      );
+    } else if (wrappedKey !== traceFunctionKey) {
+      throw new BitfabError(
+        `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 Promise.resolve().then(() => (init_replay(), replay_exports));
     return doReplay(
       this.httpClient,
       this.serviceUrl,
       traceFunctionKey,
-      fn,
+      replayFn,
       options
     );
   }