bitfab 0.21.1 → 0.21.2

package/dist/index.cjs

@@ -499,7 +499,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // src/version.generated.ts
-var __version__ = "0.21.1";
+var __version__ = "0.21.2";
 
 // src/constants.ts
 var DEFAULT_SERVICE_URL = "https://bitfab.ai";
@@ -970,6 +970,13 @@ var BitfabClaudeAgentHandler = class {
     this.currentLlmHistorySnapshot = [];
     // Subagent tracking
     this.activeSubagentSpans = /* @__PURE__ */ new Map();
+    // Synthetic root span (handler-only replay). When an `input` is supplied to
+    // wrapQuery/wrapResponse, the handler emits a root `agent` span carrying that
+    // input, so a handler-instrumented run is replayable WITHOUT an enclosing
+    // withSpan — matching the LangGraph handler, which records the graph input as
+    // its root. The prompt is not present anywhere in the message stream, so it
+    // must be handed in explicitly.
+    this.hasRootInput = false;
     this.httpClient = new HttpClient({
       apiKey: config.apiKey,
       serviceUrl: config.serviceUrl ?? DEFAULT_SERVICE_URL,
@@ -1004,7 +1011,30 @@ var BitfabClaudeAgentHandler = class {
         return subagentSpanId;
       }
     }
-    return this.activeContext?.spanId ?? this.rootSpanId ?? null;
+    return this.rootSpanId ?? this.activeContext?.spanId ?? null;
+  }
+  // Emit the synthetic root `agent` span once, before any child spans. No-op
+  // unless an `input` was supplied AND there is no enclosing withSpan (in which
+  // case that outer span is already the replayable root).
+  maybeStartRootSpan() {
+    if (!this.hasRootInput || this.rootSpanId !== null) {
+      return;
+    }
+    this.ensureTrace();
+    if (this.activeContext !== null) {
+      return;
+    }
+    const spanId = crypto.randomUUID();
+    this.startSpan(spanId, this.traceFunctionKey, "agent", this.rootInput, null);
+    this.rootSpanId = spanId;
+  }
+  completeRootSpan() {
+    if (this.rootSpanId === null) {
+      return;
+    }
+    const spanId = this.rootSpanId;
+    this.rootSpanId = null;
+    this.completeSpan(spanId, this.rootOutput);
   }
   // ── span helpers ─────────────────────────────────────────────
   startSpan(spanId, name, spanType, inputData, parentId) {
@@ -1200,26 +1230,53 @@ var BitfabClaudeAgentHandler = class {
     return options;
   }
   /**
-   * Wrap a `ClaudeSDKClient.receiveResponse()` stream to capture LLM turns.
+   * Wrap any Claude Agent SDK message stream to capture LLM turns.
    *
-   * Yields every message unchanged while capturing AssistantMessage
-   * content as LLM turn spans.
+   * Yields every message unchanged while capturing assistant message
+   * content as LLM turn spans. Kept for naming symmetry with the Python
+   * SDK's `wrapResponse` (which wraps `ClaudeSDKClient.receiveResponse()`);
+   * in TypeScript, prefer `wrapQuery` around `query()`.
+   *
+   * Pass `{ input }` (the prompt) to record a replayable root span — see
+   * `wrapQuery`.
    */
-  async *wrapResponse(stream) {
+  async *wrapResponse(stream, opts) {
+    this.setRootInput(opts);
     yield* this.processStream(stream);
   }
   /**
    * Wrap a `query()` async iterator to capture LLM turns.
    *
-   * Same as `wrapResponse` but for the simpler `query()` API
-   * which does not support hooks (no tool/subagent spans).
+   * Tool and subagent spans are captured separately via the hooks injected
+   * by `instrumentOptions` into the `options` passed to `query()`.
+   *
+   * Pass `{ input }` — the prompt (or the serializable args that produced it)
+   * — to make a handler-only run replayable: the handler records a root `agent`
+   * span with that input, so `replay(key, fn)` can re-feed it. Omit it only
+   * when an enclosing `withSpan` already supplies the replayable root.
+   *
+   * ```typescript
+   * handler.wrapQuery(query({ prompt, options }), { input: prompt })
+   * ```
    */
-  async *wrapQuery(stream) {
+  async *wrapQuery(stream, opts) {
+    this.setRootInput(opts);
     yield* this.processStream(stream);
   }
+  setRootInput(opts) {
+    if (opts && opts.input !== void 0) {
+      this.hasRootInput = true;
+      this.rootInput = opts.input;
+    } else {
+      this.hasRootInput = false;
+      this.rootInput = void 0;
+    }
+    this.rootOutput = void 0;
+  }
   // ── stream processing ────────────────────────────────────────
   async *processStream(stream) {
     try {
+      this.maybeStartRootSpan();
       for await (const message of stream) {
         try {
           this.processMessage(message);
@@ -1230,6 +1287,7 @@ var BitfabClaudeAgentHandler = class {
     } finally {
       try {
         this.flushLlmTurn();
+        this.completeRootSpan();
         this.sendTraceCompletion();
       } catch {
       }
@@ -1237,18 +1295,19 @@ var BitfabClaudeAgentHandler = class {
     }
   }
   processMessage(message) {
-    const typeName = message.constructor?.name ?? "";
-    if (typeName === "AssistantMessage") {
+    const typeName = message.type;
+    if (typeName === "assistant") {
       this.handleAssistantMessage(message);
-    } else if (typeName === "UserMessage") {
+    } else if (typeName === "user") {
       this.handleUserMessage(message);
-    } else if (typeName === "ResultMessage") {
+    } else if (typeName === "result") {
       this.handleResultMessage(message);
     }
   }
   handleAssistantMessage(message) {
     this.ensureTrace();
-    const messageId = message.message_id;
+    const inner = message.message ?? {};
+    const messageId = inner.id ?? message.uuid;
     if (messageId !== this.currentLlmMessageId) {
       this.flushLlmTurn();
       this.conversationHistory.push(...this.pendingMessages);
@@ -1256,26 +1315,27 @@ var BitfabClaudeAgentHandler = class {
       this.currentLlmSpanId = crypto.randomUUID();
       this.currentLlmMessageId = messageId ?? null;
       this.currentLlmContent = [];
-      this.currentLlmModel = message.model ?? null;
+      this.currentLlmModel = inner.model ?? null;
       this.currentLlmUsage = {};
       this.currentLlmStartedAt = nowIso();
       this.currentLlmHistorySnapshot = [...this.conversationHistory];
     }
-    const content = message.content;
+    const content = inner.content;
     if (Array.isArray(content)) {
       this.currentLlmContent.push(...extractContentBlocks(content));
     }
-    const usage = extractUsage(message);
+    const usage = extractUsage(inner);
     if (Object.keys(usage).length > 0) {
       Object.assign(this.currentLlmUsage, usage);
     }
-    const model = message.model;
+    const model = inner.model;
     if (model) {
       this.currentLlmModel = model;
     }
   }
   handleUserMessage(message) {
-    const content = message.content;
+    const inner = message.message ?? {};
+    const content = inner.content;
     const toolUseResult = message.tool_use_result;
     if (toolUseResult !== void 0) {
       this.pendingMessages.push({
@@ -1292,6 +1352,10 @@ var BitfabClaudeAgentHandler = class {
   }
   handleResultMessage(message) {
     this.flushLlmTurn();
+    if (message.result !== void 0) {
+      this.rootOutput = message.result;
+    }
+    this.completeRootSpan();
     const metadata = {};
     for (const attr of [
       "num_turns",
@@ -1355,6 +1419,9 @@ var BitfabClaudeAgentHandler = class {
     this.runToSpan.clear();
     this.traceId = null;
     this.rootSpanId = null;
+    this.hasRootInput = false;
+    this.rootInput = void 0;
+    this.rootOutput = void 0;
     this.activeContext = null;
     this.traceStartedAt = null;
     this.conversationHistory = [];
@@ -2960,14 +3027,15 @@ var Bitfab = class {
    * execution as Bitfab spans with proper parent-child hierarchy.
    *
    * ```typescript
+   * import { query } from "@anthropic-ai/claude-agent-sdk";
+   *
    * const handler = client.getClaudeAgentHandler("my-agent");
    * const options = handler.instrumentOptions({
    *   model: "claude-sonnet-4-5-...",
    * });
-   * const sdkClient = new ClaudeSDKClient(options);
-   * await sdkClient.connect();
-   * await sdkClient.query("Do something");
-   * for await (const msg of handler.wrapResponse(sdkClient.receiveResponse())) {
+   * for await (const msg of handler.wrapQuery(
+   *   query({ prompt: "Do something", options })
+   * )) {
    *   // process messages
    * }
    * ```