@bitfab/sdk 0.21.1 → 0.22.0

package/dist/index.cjs

@@ -477,6 +477,7 @@ __export(index_exports, {
   BitfabFunction: () => BitfabFunction,
   BitfabLangChainCallbackHandler: () => BitfabLangGraphCallbackHandler,
   BitfabLangGraphCallbackHandler: () => BitfabLangGraphCallbackHandler,
+  BitfabOpenAIAgentHandler: () => BitfabOpenAIAgentHandler,
   BitfabOpenAITracingProcessor: () => BitfabOpenAITracingProcessor,
   DEFAULT_SERVICE_URL: () => DEFAULT_SERVICE_URL,
   ReplayEnvironment: () => ReplayEnvironment,
@@ -490,7 +491,7 @@ __export(index_exports, {
 module.exports = __toCommonJS(index_exports);
 
 // src/version.generated.ts
-var __version__ = "0.21.1";
+var __version__ = "0.22.0";
 
 // src/constants.ts
 var DEFAULT_SERVICE_URL = "https://bitfab.ai";
@@ -961,6 +962,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,
@@ -995,7 +1003,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) {
@@ -1191,26 +1222,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 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()`.
    *
-   * Yields every message unchanged while capturing AssistantMessage
-   * content as LLM turn spans.
+   * 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);
@@ -1221,6 +1279,7 @@ var BitfabClaudeAgentHandler = class {
     } finally {
       try {
         this.flushLlmTurn();
+        this.completeRootSpan();
         this.sendTraceCompletion();
       } catch {
       }
@@ -1228,18 +1287,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);
@@ -1247,26 +1307,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({
@@ -1283,6 +1344,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",
@@ -1346,6 +1411,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 = [];
@@ -2205,6 +2273,39 @@ var BitfabLangGraphCallbackHandler = class {
   }
 };
 
+// src/openaiAgentSdk.ts
+var BitfabOpenAIAgentHandler = class {
+  constructor(config) {
+    this.traceFunctionKey = config.traceFunctionKey;
+    this.withSpanFn = config.withSpan;
+  }
+  async wrapRun(agent, input, options) {
+    const { run } = await import("@openai/agents");
+    const isStreaming = options?.stream === true;
+    const finalize = async (result) => {
+      const res = result;
+      if (isStreaming && res?.completed) {
+        try {
+          await res.completed;
+        } catch {
+        }
+      }
+      return res?.finalOutput;
+    };
+    const options_ = { type: "agent", finalize };
+    const traced = this.withSpanFn(
+      this.traceFunctionKey,
+      options_,
+      (agentInput) => run(
+        agent,
+        agentInput,
+        options
+      )
+    );
+    return traced(input);
+  }
+};
+
 // src/client.ts
 init_replayContext();
 
@@ -2898,6 +2999,34 @@ var Bitfab = class {
       }
     });
   }
+  /**
+   * Get an OpenAI Agents SDK handler that records a replayable root span.
+   *
+   * The processor from {@link getOpenAiTracingProcessor} captures everything
+   * inside a run (LLM calls, tools, handoffs) but never sees the caller's
+   * input, so a processor-only run records an empty-input root and is not
+   * replayable. This handler's `wrapRun` is a drop-in for `run()` that opens a
+   * `withSpan` root carrying the input and final output; the processor's spans
+   * nest beneath it. Register the processor once at startup, then call
+   * `handler.wrapRun(agent, input)` in place of `run(agent, input)`.
+   *
+   * ```typescript
+   * import { addTraceProcessor, Agent, run } from "@openai/agents";
+   *
+   * addTraceProcessor(client.getOpenAiTracingProcessor());
+   * const handler = client.getOpenAiAgentHandler("research-topic");
+   * const result = await handler.wrapRun(agent, "Find X");
+   * ```
+   *
+   * @param traceFunctionKey - Groups traces under this key in Bitfab
+   * @returns A BitfabOpenAIAgentHandler configured for this client
+   */
+  getOpenAiAgentHandler(traceFunctionKey) {
+    return new BitfabOpenAIAgentHandler({
+      traceFunctionKey,
+      withSpan: this.withSpan.bind(this)
+    });
+  }
   /**
    * Get a LangGraph/LangChain callback handler for tracing.
    *
@@ -2951,14 +3080,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
    * }
    * ```
@@ -3647,6 +3777,7 @@ var finalizers = {
   BitfabFunction,
   BitfabLangChainCallbackHandler,
   BitfabLangGraphCallbackHandler,
+  BitfabOpenAIAgentHandler,
   BitfabOpenAITracingProcessor,
   DEFAULT_SERVICE_URL,
   ReplayEnvironment,