bitfab 0.18.0 → 0.18.2

package/dist/index.cjs

@@ -442,7 +442,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // src/version.generated.ts
-var __version__ = "0.18.0";
+var __version__ = "0.18.2";
 
 // src/constants.ts
 var DEFAULT_SERVICE_URL = "https://bitfab.ai";
@@ -1672,21 +1672,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 (outputTokens !== void 0 && outputTokens !== null) {
-    usage.outputTokens = outputTokens;
+  if (normalized.outputTokens !== null) {
+    usage.outputTokens = normalized.outputTokens;
   }
-  if (totalTokens !== void 0 && totalTokens !== null) {
-    usage.totalTokens = totalTokens;
+  if (normalized.totalTokens !== null) {
+    usage.totalTokens = normalized.totalTokens;
+  }
+  if (normalized.cachedInputTokens !== null) {
+    usage.cachedInputTokens = normalized.cachedInputTokens;
   }
   return usage;
 }
@@ -2047,7 +2157,9 @@ var ReplayEnvironment = class {
    * Throws if read outside a replay item.
    */
   get databaseUrl() {
-    return this.require().databaseUrl;
+    const snapshot = this.require();
+    this.markAccessed();
+    return snapshot.databaseUrl;
   }
   /** When the per-trace branch URL stops being valid. ISO-8601. */
   get expiresAt() {
@@ -2075,7 +2187,24 @@ var ReplayEnvironment = class {
   }
   /** Non-throwing variant for callers that handle the inactive case. */
   snapshot() {
-    return this.read();
+    const snapshot = this.read();
+    if (snapshot) {
+      this.markAccessed();
+    }
+    return snapshot;
+  }
+  /**
+   * Record on the replay context that customer code obtained the branch
+   * URL. Only `databaseUrl` and `snapshot()` count — `active`, `readOnly`
+   * and friends inspect the lease without exposing the connection string,
+   * so they don't prove the replayed code could have connected to the
+   * branch.
+   */
+  markAccessed() {
+    const ctx = getReplayContext();
+    if (ctx?.dbBranchLease) {
+      ctx.dbSnapshotAccessed = true;
+    }
   }
   read() {
     const ctx = getReplayContext();
@@ -3007,7 +3136,19 @@ var Bitfab = class {
               contexts: traceState?.contexts ?? [],
               testRunId: traceState?.testRunId,
               inputSourceTraceId: traceState?.inputSourceTraceId,
-              dbSnapshotRef: traceState?.dbSnapshotRef
+              dbSnapshotRef: traceState?.dbSnapshotRef,
+              // Built AFTER the wrapped fn finished, so `accessed` reflects
+              // whether customer code obtained the branch URL during this
+              // item. Omitted entirely when no lease was attached, so the
+              // server can distinguish "no branch" from "branch ignored".
+              ...replayCtx?.dbBranchLease && {
+                dbSnapshotUsage: {
+                  neonBranchId: replayCtx.dbBranchLease.neonBranchId,
+                  snapshotTimestamp: replayCtx.dbBranchLease.snapshotTimestamp,
+                  sourceTraceId: replayCtx.sourceBitfabTraceId,
+                  accessed: replayCtx.dbSnapshotAccessed === true
+                }
+              }
             });
             activeTraceStates.delete(traceId);
             if (persistenceCollector) {
@@ -3179,6 +3320,18 @@ var Bitfab = class {
     if (params.dbSnapshotRef) {
       rawTrace.db_snapshot_ref = params.dbSnapshotRef;
     }
+    if (params.dbSnapshotUsage) {
+      rawTrace.db_snapshot_usage = {
+        neon_branch_id: params.dbSnapshotUsage.neonBranchId,
+        ...params.dbSnapshotUsage.snapshotTimestamp && {
+          snapshot_timestamp: params.dbSnapshotUsage.snapshotTimestamp
+        },
+        ...params.dbSnapshotUsage.sourceTraceId && {
+          source_trace_id: params.dbSnapshotUsage.sourceTraceId
+        },
+        accessed: params.dbSnapshotUsage.accessed
+      };
+    }
     return this.httpClient.sendExternalTrace({
       type: "sdk-function",
       source: "typescript-sdk-function",