braintrust 3.6.0 → 3.7.0

package/dev/dist/index.mjs

@@ -24,25 +24,51 @@ function patchTracingChannel(tracingChannelFn) {
   if (TracingChannel.prototype.tracePromise) {
     TracingChannel.prototype.tracePromise = function(fn, context = {}, thisArg, ...args) {
       const { start, end, asyncStart, asyncEnd, error } = this;
-      function reject2(err) {
+      function publishRejected(err) {
         context.error = err;
         error?.publish(context);
         asyncStart?.publish(context);
         asyncEnd?.publish(context);
-        return Promise.reject(err);
       }
-      function resolve(result) {
+      function publishResolved(result) {
         context.result = result;
         asyncStart?.publish(context);
         asyncEnd?.publish(context);
-        return result;
       }
       return start.runStores(context, () => {
         try {
           const result = Reflect.apply(fn, thisArg, args);
           end?.publish(context);
           if (result && (typeof result === "object" || typeof result === "function") && typeof result.then === "function") {
-            return result.then(resolve, reject2);
+            if (result.constructor === Promise) {
+              return result.then(
+                (res) => {
+                  publishResolved(res);
+                  return res;
+                },
+                (err) => {
+                  publishRejected(err);
+                  return Promise.reject(err);
+                }
+              );
+            }
+            void result.then(
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              (resolved) => {
+                try {
+                  publishResolved(resolved);
+                } catch {
+                }
+              },
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
+              (err) => {
+                try {
+                  publishRejected(err);
+                } catch {
+                }
+              }
+            );
+            return result;
           }
           context.result = result;
           asyncStart?.publish(context);
@@ -10468,7 +10494,7 @@ var AnthropicPlugin = class extends BasePlugin {
     this.unsubscribers.push(
       traceStreamingChannel(anthropicChannels.betaMessagesCreate, {
         ...anthropicConfig,
-        name: "anthropic.beta.messages.create"
+        name: "anthropic.messages.create"
       })
     );
   }
@@ -10491,9 +10517,12 @@ function parseMetricsFromUsage2(usage) {
   return metrics;
 }
 function aggregateAnthropicStreamChunks(chunks) {
-  const deltas = [];
+  const fallbackTextDeltas = [];
+  const contentBlocks = {};
+  const contentBlockDeltas = {};
   let metrics = {};
   let metadata = {};
+  let role;
   for (const event of chunks) {
     switch (event?.type) {
       case "message_start":
@@ -10501,15 +10530,43 @@ function aggregateAnthropicStreamChunks(chunks) {
           const initialMetrics = parseMetricsFromUsage2(event.message.usage);
           metrics = { ...metrics, ...initialMetrics };
         }
+        if (typeof event.message?.role === "string") {
+          role = event.message.role;
+        }
+        break;
+      case "content_block_start":
+        if (event.content_block) {
+          contentBlocks[event.index] = event.content_block;
+          contentBlockDeltas[event.index] = [];
+        }
         break;
       case "content_block_delta":
         if (event.delta?.type === "text_delta") {
           const text = event.delta.text;
           if (text) {
-            deltas.push(text);
+            if (contentBlocks[event.index] !== void 0 || contentBlockDeltas[event.index] !== void 0) {
+              contentBlockDeltas[event.index] ??= [];
+              contentBlockDeltas[event.index].push(text);
+            } else {
+              fallbackTextDeltas.push(text);
+            }
+          }
+        } else if (event.delta?.type === "input_json_delta") {
+          const partialJson = event.delta.partial_json;
+          if (partialJson) {
+            contentBlockDeltas[event.index] ??= [];
+            contentBlockDeltas[event.index].push(partialJson);
           }
         }
         break;
+      case "content_block_stop":
+        finalizeContentBlock(
+          event.index,
+          contentBlocks,
+          contentBlockDeltas,
+          fallbackTextDeltas
+        );
+        break;
       case "message_delta":
         if (event.usage) {
           const finalMetrics = parseMetricsFromUsage2(event.usage);
@@ -10521,7 +10578,21 @@ function aggregateAnthropicStreamChunks(chunks) {
         break;
     }
   }
-  const output = deltas.join("");
+  const orderedContent = Object.entries(contentBlocks).map(([index, block]) => ({
+    block,
+    index: Number(index)
+  })).filter(({ block }) => block !== void 0).sort((left, right) => left.index - right.index).map(({ block }) => block);
+  let output = fallbackTextDeltas.join("");
+  if (orderedContent.length > 0) {
+    if (orderedContent.every(isTextContentBlock)) {
+      output = orderedContent.map((block) => block.text).join("");
+    } else {
+      output = {
+        ...role ? { role } : {},
+        content: orderedContent
+      };
+    }
+  }
   const finalized = finalizeAnthropicTokens(metrics);
   const filteredMetrics = Object.fromEntries(
     Object.entries(finalized).filter(
@@ -10534,6 +10605,49 @@ function aggregateAnthropicStreamChunks(chunks) {
     metadata
   };
 }
+function finalizeContentBlock(index, contentBlocks, contentBlockDeltas, fallbackTextDeltas) {
+  const contentBlock = contentBlocks[index];
+  if (!contentBlock) {
+    return;
+  }
+  const text = contentBlockDeltas[index]?.join("") ?? "";
+  if (isToolUseContentBlock(contentBlock)) {
+    if (!text) {
+      return;
+    }
+    try {
+      contentBlocks[index] = {
+        ...contentBlock,
+        input: JSON.parse(text)
+      };
+    } catch {
+      fallbackTextDeltas.push(text);
+      delete contentBlocks[index];
+    }
+    return;
+  }
+  if (isTextContentBlock(contentBlock)) {
+    if (!text) {
+      delete contentBlocks[index];
+      return;
+    }
+    contentBlocks[index] = {
+      ...contentBlock,
+      text
+    };
+    return;
+  }
+  if (text) {
+    fallbackTextDeltas.push(text);
+  }
+  delete contentBlocks[index];
+}
+function isTextContentBlock(contentBlock) {
+  return contentBlock.type === "text";
+}
+function isToolUseContentBlock(contentBlock) {
+  return contentBlock.type === "tool_use";
+}
 function isAnthropicBase64ContentBlock(input) {
   return (input.type === "image" || input.type === "document") && isObject(input.source) && input.source.type === "base64";
 }
@@ -11656,12 +11770,15 @@ var claudeAgentSDKChannels = defineChannels(
   {
     query: channel({
       channelName: "query",
-      kind: "async"
+      kind: "sync-stream"
     })
   }
 );
 
 // src/instrumentation/plugins/claude-agent-sdk-plugin.ts
+function isSubAgentToolName(toolName) {
+  return toolName === "Agent" || toolName === "Task";
+}
 function filterSerializableOptions(options) {
   const allowedKeys = [
     "model",
@@ -11715,34 +11832,50 @@ function extractUsageFromMessage(message) {
   const cacheReadTokens = getNumberProperty(usage, "cache_read_input_tokens") || 0;
   const cacheCreationTokens = getNumberProperty(usage, "cache_creation_input_tokens") || 0;
   if (cacheReadTokens > 0 || cacheCreationTokens > 0) {
-    const cacheTokens = extractAnthropicCacheTokens(
-      cacheReadTokens,
-      cacheCreationTokens
+    Object.assign(
+      metrics,
+      extractAnthropicCacheTokens(cacheReadTokens, cacheCreationTokens)
     );
-    Object.assign(metrics, cacheTokens);
   }
   if (Object.keys(metrics).length > 0) {
     Object.assign(metrics, finalizeAnthropicTokens(metrics));
   }
   return metrics;
 }
-function buildLLMInput(prompt, conversationHistory) {
-  const promptMessage = typeof prompt === "string" ? { content: prompt, role: "user" } : void 0;
-  const inputParts = [
-    ...promptMessage ? [promptMessage] : [],
-    ...conversationHistory
-  ];
+function buildLLMInput(prompt, conversationHistory, capturedPromptMessages) {
+  const promptMessages = [];
+  if (typeof prompt === "string") {
+    promptMessages.push({ content: prompt, role: "user" });
+  } else if (capturedPromptMessages && capturedPromptMessages.length > 0) {
+    for (const msg of capturedPromptMessages) {
+      const role = msg.message?.role;
+      const content = msg.message?.content;
+      if (role && content !== void 0) {
+        promptMessages.push({ content, role });
+      }
+    }
+  }
+  const inputParts = [...promptMessages, ...conversationHistory];
   return inputParts.length > 0 ? inputParts : void 0;
 }
-async function createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, parentSpan) {
-  if (messages.length === 0) return void 0;
+function formatCapturedMessages(messages) {
+  return messages.length > 0 ? messages : [];
+}
+async function createLLMSpanForMessages(messages, prompt, conversationHistory, options, startTime, capturedPromptMessages, parentSpan) {
+  if (messages.length === 0) {
+    return void 0;
+  }
   const lastMessage = messages[messages.length - 1];
   if (lastMessage.type !== "assistant" || !lastMessage.message?.usage) {
     return void 0;
   }
   const model = lastMessage.message.model || options.model;
   const usage = extractUsageFromMessage(lastMessage);
-  const input = buildLLMInput(prompt, conversationHistory);
+  const input = buildLLMInput(
+    prompt,
+    conversationHistory,
+    capturedPromptMessages
+  );
   const outputs = messages.map(
     (m) => m.message?.content && m.message?.role ? { content: m.message.content, role: m.message.role } : void 0
   ).filter(
@@ -11750,21 +11883,359 @@ async function createLLMSpanForMessages(messages, prompt, conversationHistory, o
   );
   const span = startSpan({
     name: "anthropic.messages.create",
+    parent: parentSpan,
     spanAttributes: {
       type: "llm" /* LLM */
     },
-    startTime,
-    parent: parentSpan
+    startTime
   });
   span.log({
     input,
-    output: outputs,
     metadata: model ? { model } : void 0,
-    metrics: usage
+    metrics: usage,
+    output: outputs
   });
   await span.end();
   return lastMessage.message?.content && lastMessage.message?.role ? { content: lastMessage.message.content, role: lastMessage.message.role } : void 0;
 }
+function getMcpServerMetadata(serverName, mcpServers) {
+  if (!serverName || !mcpServers) {
+    return {};
+  }
+  const serverConfig = mcpServers[serverName];
+  if (!serverConfig) {
+    return {};
+  }
+  const metadata = {};
+  if (serverConfig.type) {
+    metadata["mcp.type"] = serverConfig.type;
+  } else if (typeof serverConfig === "object" && "transport" in serverConfig) {
+    metadata["mcp.type"] = "sdk";
+  }
+  if (serverConfig.url) {
+    metadata["mcp.url"] = serverConfig.url;
+  }
+  if (serverConfig.command) {
+    metadata["mcp.command"] = serverConfig.command;
+    if (serverConfig.args) {
+      metadata["mcp.args"] = serverConfig.args.join(" ");
+    }
+  }
+  return metadata;
+}
+function parseToolName(rawToolName) {
+  const mcpMatch = rawToolName.match(/^mcp__([^_]+)__(.+)$/);
+  if (mcpMatch) {
+    const [, mcpServer, toolName] = mcpMatch;
+    return {
+      displayName: `tool: ${mcpServer}/${toolName}`,
+      mcpServer,
+      rawToolName,
+      toolName
+    };
+  }
+  return {
+    displayName: `tool: ${rawToolName}`,
+    rawToolName,
+    toolName: rawToolName
+  };
+}
+function createToolTracingHooks(resolveParentSpan, activeToolSpans, mcpServers, subAgentSpans, endedSubAgentSpans) {
+  const preToolUse = async (input, toolUseID) => {
+    if (input.hook_event_name !== "PreToolUse" || !toolUseID) {
+      return {};
+    }
+    if (isSubAgentToolName(input.tool_name)) {
+      return {};
+    }
+    const parsed = parseToolName(input.tool_name);
+    const toolSpan = startSpan({
+      event: {
+        input: input.tool_input,
+        metadata: {
+          "claude_agent_sdk.cwd": input.cwd,
+          "claude_agent_sdk.raw_tool_name": parsed.rawToolName,
+          "claude_agent_sdk.session_id": input.session_id,
+          "gen_ai.tool.call.id": toolUseID,
+          "gen_ai.tool.name": parsed.toolName,
+          ...parsed.mcpServer && { "mcp.server": parsed.mcpServer },
+          ...getMcpServerMetadata(parsed.mcpServer, mcpServers)
+        }
+      },
+      name: parsed.displayName,
+      parent: await resolveParentSpan(toolUseID),
+      spanAttributes: { type: "tool" /* TOOL */ }
+    });
+    activeToolSpans.set(toolUseID, toolSpan);
+    return {};
+  };
+  const postToolUse = async (input, toolUseID) => {
+    if (input.hook_event_name !== "PostToolUse" || !toolUseID) {
+      return {};
+    }
+    const subAgentSpan = subAgentSpans.get(toolUseID);
+    if (subAgentSpan) {
+      try {
+        const response = input.tool_response;
+        const metadata = {};
+        if (response?.status) {
+          metadata["claude_agent_sdk.status"] = response.status;
+        }
+        if (response?.totalDurationMs) {
+          metadata["claude_agent_sdk.duration_ms"] = response.totalDurationMs;
+        }
+        if (response?.totalToolUseCount !== void 0) {
+          metadata["claude_agent_sdk.tool_use_count"] = response.totalToolUseCount;
+        }
+        subAgentSpan.log({
+          metadata,
+          output: response?.content
+        });
+      } finally {
+        subAgentSpan.end();
+        endedSubAgentSpans.add(toolUseID);
+      }
+      return {};
+    }
+    const toolSpan = activeToolSpans.get(toolUseID);
+    if (!toolSpan) {
+      return {};
+    }
+    try {
+      toolSpan.log({ output: input.tool_response });
+    } finally {
+      toolSpan.end();
+      activeToolSpans.delete(toolUseID);
+    }
+    return {};
+  };
+  const postToolUseFailure = async (input, toolUseID) => {
+    if (input.hook_event_name !== "PostToolUseFailure" || !toolUseID) {
+      return {};
+    }
+    const subAgentSpan = subAgentSpans.get(toolUseID);
+    if (subAgentSpan) {
+      try {
+        subAgentSpan.log({ error: input.error });
+      } finally {
+        subAgentSpan.end();
+        endedSubAgentSpans.add(toolUseID);
+      }
+      return {};
+    }
+    const toolSpan = activeToolSpans.get(toolUseID);
+    if (!toolSpan) {
+      return {};
+    }
+    const parsed = parseToolName(input.tool_name);
+    try {
+      toolSpan.log({
+        error: input.error,
+        metadata: {
+          "claude_agent_sdk.is_interrupt": input.is_interrupt,
+          "claude_agent_sdk.session_id": input.session_id,
+          "gen_ai.tool.call.id": toolUseID,
+          "gen_ai.tool.name": parsed.toolName,
+          ...parsed.mcpServer && { "mcp.server": parsed.mcpServer }
+        }
+      });
+    } finally {
+      toolSpan.end();
+      activeToolSpans.delete(toolUseID);
+    }
+    return {};
+  };
+  return { postToolUse, postToolUseFailure, preToolUse };
+}
+function injectTracingHooks(options, resolveParentSpan, activeToolSpans, subAgentSpans, endedSubAgentSpans) {
+  const { preToolUse, postToolUse, postToolUseFailure } = createToolTracingHooks(
+    resolveParentSpan,
+    activeToolSpans,
+    options.mcpServers,
+    subAgentSpans,
+    endedSubAgentSpans
+  );
+  const existingHooks = options.hooks ?? {};
+  return {
+    ...options,
+    hooks: {
+      ...existingHooks,
+      PostToolUse: [
+        ...existingHooks.PostToolUse ?? [],
+        { hooks: [postToolUse] }
+      ],
+      PostToolUseFailure: [
+        ...existingHooks.PostToolUseFailure ?? [],
+        {
+          hooks: [postToolUseFailure]
+        }
+      ],
+      PreToolUse: [
+        ...existingHooks.PreToolUse ?? [],
+        { hooks: [preToolUse] }
+      ]
+    }
+  };
+}
+async function finalizeCurrentMessageGroup(state) {
+  if (state.currentMessages.length === 0) {
+    return;
+  }
+  const parentToolUseId = state.currentMessages[0]?.parent_tool_use_id ?? null;
+  let parentSpan = await state.span.export();
+  if (parentToolUseId) {
+    const subAgentSpan = state.subAgentSpans.get(parentToolUseId);
+    if (subAgentSpan) {
+      parentSpan = await subAgentSpan.export();
+    }
+  }
+  const finalMessage = await createLLMSpanForMessages(
+    state.currentMessages,
+    state.originalPrompt,
+    state.finalResults,
+    state.options,
+    state.currentMessageStartTime,
+    state.capturedPromptMessages,
+    parentSpan
+  );
+  if (finalMessage) {
+    state.finalResults.push(finalMessage);
+  }
+  const lastMessage = state.currentMessages[state.currentMessages.length - 1];
+  if (lastMessage?.message?.usage) {
+    state.accumulatedOutputTokens += getNumberProperty(lastMessage.message.usage, "output_tokens") || 0;
+  }
+  state.currentMessages.length = 0;
+}
+function maybeTrackToolUseContext(state, message) {
+  if (message.type !== "assistant" || !Array.isArray(message.message?.content)) {
+    return;
+  }
+  const parentToolUseId = message.parent_tool_use_id ?? null;
+  for (const block of message.message.content) {
+    if (typeof block !== "object" || block === null || !("type" in block) || block.type !== "tool_use" || !("id" in block) || typeof block.id !== "string") {
+      continue;
+    }
+    state.toolUseToParent.set(block.id, parentToolUseId);
+    if (block.name === "Task" && typeof block.input === "object" && block.input !== null && "subagent_type" in block.input && typeof block.input.subagent_type === "string") {
+      state.pendingSubAgentNames.set(block.id, block.input.subagent_type);
+    }
+  }
+}
+async function maybeStartSubAgentSpan(state, message) {
+  if (!("parent_tool_use_id" in message)) {
+    return;
+  }
+  const parentToolUseId = message.parent_tool_use_id;
+  if (!parentToolUseId) {
+    return;
+  }
+  await ensureSubAgentSpan(
+    state.pendingSubAgentNames,
+    state.span,
+    state.subAgentSpans,
+    parentToolUseId
+  );
+}
+async function ensureSubAgentSpan(pendingSubAgentNames, rootSpan, subAgentSpans, parentToolUseId) {
+  const existingSpan = subAgentSpans.get(parentToolUseId);
+  if (existingSpan) {
+    return existingSpan;
+  }
+  const agentName = pendingSubAgentNames.get(parentToolUseId);
+  const spanName = agentName ? `Agent: ${agentName}` : "Agent: sub-agent";
+  const subAgentSpan = startSpan({
+    event: {
+      metadata: {
+        ...agentName && { "claude_agent_sdk.agent_type": agentName }
+      }
+    },
+    name: spanName,
+    parent: await rootSpan.export(),
+    spanAttributes: { type: "task" /* TASK */ }
+  });
+  subAgentSpans.set(parentToolUseId, subAgentSpan);
+  return subAgentSpan;
+}
+async function handleStreamMessage(state, message) {
+  maybeTrackToolUseContext(state, message);
+  await maybeStartSubAgentSpan(state, message);
+  const messageId = message.message?.id;
+  if (messageId && messageId !== state.currentMessageId) {
+    await finalizeCurrentMessageGroup(state);
+    state.currentMessageId = messageId;
+    state.currentMessageStartTime = getCurrentUnixTimestamp();
+  }
+  if (message.type === "assistant" && message.message?.usage) {
+    state.currentMessages.push(message);
+  }
+  if (message.type !== "result" || !message.usage) {
+    return;
+  }
+  const finalUsageMetrics = extractUsageFromMessage(message);
+  if (state.currentMessages.length > 0 && finalUsageMetrics.completion_tokens !== void 0) {
+    const lastMessage = state.currentMessages[state.currentMessages.length - 1];
+    if (lastMessage?.message?.usage) {
+      const adjustedTokens = finalUsageMetrics.completion_tokens - state.accumulatedOutputTokens;
+      if (adjustedTokens >= 0) {
+        lastMessage.message.usage.output_tokens = adjustedTokens;
+      }
+      const resultUsage = message.usage;
+      if (resultUsage && typeof resultUsage === "object") {
+        const cacheReadTokens = getNumberProperty(
+          resultUsage,
+          "cache_read_input_tokens"
+        );
+        if (cacheReadTokens !== void 0) {
+          lastMessage.message.usage.cache_read_input_tokens = cacheReadTokens;
+        }
+        const cacheCreationTokens = getNumberProperty(
+          resultUsage,
+          "cache_creation_input_tokens"
+        );
+        if (cacheCreationTokens !== void 0) {
+          lastMessage.message.usage.cache_creation_input_tokens = cacheCreationTokens;
+        }
+      }
+    }
+  }
+  const metadata = {};
+  if (message.num_turns !== void 0) {
+    metadata.num_turns = message.num_turns;
+  }
+  if (message.session_id !== void 0) {
+    metadata.session_id = message.session_id;
+  }
+  if (Object.keys(metadata).length > 0) {
+    state.span.log({ metadata });
+  }
+}
+async function finalizeQuerySpan(state) {
+  try {
+    await finalizeCurrentMessageGroup(state);
+    state.span.log({
+      output: state.finalResults.length > 0 ? state.finalResults[state.finalResults.length - 1] : void 0
+    });
+    if (state.capturedPromptMessages) {
+      if (state.promptStarted()) {
+        await state.promptDone;
+      }
+      if (state.capturedPromptMessages.length > 0) {
+        state.span.log({
+          input: formatCapturedMessages(state.capturedPromptMessages)
+        });
+      }
+    }
+  } finally {
+    for (const [id, subAgentSpan] of state.subAgentSpans) {
+      if (!state.endedSubAgentSpans.has(id)) {
+        subAgentSpan.end();
+      }
+    }
+    state.subAgentSpans.clear();
+    state.span.end();
+  }
+}
 var ClaudeAgentSDKPlugin = class extends BasePlugin {
   onEnable() {
     this.subscribeToQuery();
@@ -11775,19 +12246,36 @@ var ClaudeAgentSDKPlugin = class extends BasePlugin {
     }
     this.unsubscribers = [];
   }
-  /**
-   * Subscribe to the query channel for agent interactions.
-   * Handles streaming responses and traces both the top-level agent task
-   * and individual LLM calls.
-   */
   subscribeToQuery() {
     const channel2 = claudeAgentSDKChannels.query.tracingChannel();
     const spans = /* @__PURE__ */ new WeakMap();
     const handlers = {
       start: (event) => {
-        const params = event.arguments[0];
-        const prompt = params?.prompt;
-        const options = params?.options ?? {};
+        const params = event.arguments[0] ?? {};
+        const originalPrompt = params.prompt;
+        const options = params.options ?? {};
+        const promptIsAsyncIterable = isAsyncIterable(originalPrompt);
+        let promptStarted = false;
+        let capturedPromptMessages;
+        let resolvePromptDone;
+        const promptDone = new Promise((resolve) => {
+          resolvePromptDone = resolve;
+        });
+        if (promptIsAsyncIterable) {
+          capturedPromptMessages = [];
+          const promptStream = originalPrompt;
+          params.prompt = (async function* () {
+            promptStarted = true;
+            try {
+              for await (const message of promptStream) {
+                capturedPromptMessages.push(message);
+                yield message;
+              }
+            } finally {
+              resolvePromptDone?.();
+            }
+          })();
+        }
         const span = startSpan({
           name: "Claude Agent",
           spanAttributes: {
@@ -11797,163 +12285,115 @@ var ClaudeAgentSDKPlugin = class extends BasePlugin {
         const startTime = getCurrentUnixTimestamp();
         try {
           span.log({
-            input: typeof prompt === "string" ? prompt : {
-              type: "streaming",
-              description: "AsyncIterable<ClaudeAgentSDKMessage>"
-            },
+            input: typeof originalPrompt === "string" ? originalPrompt : promptIsAsyncIterable ? void 0 : originalPrompt !== void 0 ? String(originalPrompt) : void 0,
             metadata: filterSerializableOptions(options)
           });
         } catch (error) {
           console.error("Error extracting input for Claude Agent SDK:", error);
         }
+        const activeToolSpans = /* @__PURE__ */ new Map();
+        const subAgentSpans = /* @__PURE__ */ new Map();
+        const endedSubAgentSpans = /* @__PURE__ */ new Set();
+        const toolUseToParent = /* @__PURE__ */ new Map();
+        const pendingSubAgentNames = /* @__PURE__ */ new Map();
+        const optionsWithHooks = injectTracingHooks(
+          options,
+          async (toolUseID) => {
+            const parentToolUseId = toolUseToParent.get(toolUseID);
+            if (parentToolUseId) {
+              const subAgentSpan = await ensureSubAgentSpan(
+                pendingSubAgentNames,
+                span,
+                subAgentSpans,
+                parentToolUseId
+              );
+              return subAgentSpan.export();
+            }
+            return span.export();
+          },
+          activeToolSpans,
+          subAgentSpans,
+          endedSubAgentSpans
+        );
+        params.options = optionsWithHooks;
+        event.arguments[0] = params;
         spans.set(event, {
-          span,
-          startTime,
-          conversationHistory: [],
-          currentMessages: [],
+          accumulatedOutputTokens: 0,
+          activeToolSpans,
+          capturedPromptMessages,
           currentMessageId: void 0,
           currentMessageStartTime: startTime,
-          accumulatedOutputTokens: 0
+          currentMessages: [],
+          endedSubAgentSpans,
+          finalResults: [],
+          options: optionsWithHooks,
+          originalPrompt,
+          pendingSubAgentNames,
+          processing: Promise.resolve(),
+          promptDone,
+          promptStarted: () => promptStarted,
+          span,
+          subAgentSpans,
+          toolUseToParent
         });
       },
-      asyncEnd: (event) => {
-        const spanData = spans.get(event);
-        if (!spanData) {
+      end: (event) => {
+        const state = spans.get(event);
+        if (!state) {
           return;
         }
         const eventResult = event.result;
         if (eventResult === void 0) {
-          spanData.span.end();
+          state.span.end();
           spans.delete(event);
           return;
         }
         if (isAsyncIterable(eventResult)) {
           patchStreamIfNeeded(eventResult, {
-            onChunk: async (message) => {
-              const currentTime = getCurrentUnixTimestamp();
-              const params = event.arguments[0];
-              const prompt = params?.prompt;
-              const options = params?.options ?? {};
-              const messageId = message.message?.id;
-              if (messageId && messageId !== spanData.currentMessageId) {
-                if (spanData.currentMessages.length > 0) {
-                  const finalMessage = await createLLMSpanForMessages(
-                    spanData.currentMessages,
-                    prompt,
-                    spanData.conversationHistory,
-                    options,
-                    spanData.currentMessageStartTime,
-                    await spanData.span.export()
-                  );
-                  if (finalMessage) {
-                    spanData.conversationHistory.push(finalMessage);
-                  }
-                  const lastMessage = spanData.currentMessages[spanData.currentMessages.length - 1];
-                  if (lastMessage?.message?.usage) {
-                    const outputTokens = getNumberProperty(
-                      lastMessage.message.usage,
-                      "output_tokens"
-                    ) || 0;
-                    spanData.accumulatedOutputTokens += outputTokens;
-                  }
-                  spanData.currentMessages = [];
-                }
-                spanData.currentMessageId = messageId;
-                spanData.currentMessageStartTime = currentTime;
-              }
-              if (message.type === "assistant" && message.message?.usage) {
-                spanData.currentMessages.push(message);
-              }
-              if (message.type === "result" && message.usage) {
-                const finalUsageMetrics = extractUsageFromMessage(message);
-                if (spanData.currentMessages.length > 0 && finalUsageMetrics.completion_tokens !== void 0) {
-                  const lastMessage = spanData.currentMessages[spanData.currentMessages.length - 1];
-                  if (lastMessage?.message?.usage) {
-                    const adjustedTokens = finalUsageMetrics.completion_tokens - spanData.accumulatedOutputTokens;
-                    if (adjustedTokens >= 0) {
-                      lastMessage.message.usage.output_tokens = adjustedTokens;
-                    }
-                  }
-                }
-                const result_metadata = {};
-                if (message.num_turns !== void 0) {
-                  result_metadata.num_turns = message.num_turns;
-                }
-                if (message.session_id !== void 0) {
-                  result_metadata.session_id = message.session_id;
-                }
-                if (Object.keys(result_metadata).length > 0) {
-                  spanData.span.log({
-                    metadata: result_metadata
-                  });
-                }
-              }
-            },
-            onComplete: async () => {
-              try {
-                const params = event.arguments[0];
-                const prompt = params?.prompt;
-                const options = params?.options ?? {};
-                if (spanData.currentMessages.length > 0) {
-                  const finalMessage = await createLLMSpanForMessages(
-                    spanData.currentMessages,
-                    prompt,
-                    spanData.conversationHistory,
-                    options,
-                    spanData.currentMessageStartTime,
-                    await spanData.span.export()
-                  );
-                  if (finalMessage) {
-                    spanData.conversationHistory.push(finalMessage);
-                  }
-                }
-                spanData.span.log({
-                  output: spanData.conversationHistory.length > 0 ? spanData.conversationHistory[spanData.conversationHistory.length - 1] : void 0
-                });
-              } catch (error) {
+            onChunk: (message) => {
+              maybeTrackToolUseContext(state, message);
+              state.processing = state.processing.then(() => handleStreamMessage(state, message)).catch((error) => {
                 console.error(
-                  "Error extracting output for Claude Agent SDK:",
+                  "Error processing Claude Agent SDK stream chunk:",
                   error
                 );
-              } finally {
-                spanData.span.end();
+              });
+            },
+            onComplete: () => {
+              void state.processing.then(() => finalizeQuerySpan(state)).finally(() => {
                 spans.delete(event);
-              }
+              });
             },
             onError: (error) => {
-              spanData.span.log({
-                error: error.message
+              void state.processing.then(() => {
+                state.span.log({
+                  error: error.message
+                });
+              }).then(() => finalizeQuerySpan(state)).finally(() => {
+                spans.delete(event);
               });
-              spanData.span.end();
-              spans.delete(event);
             }
           });
-        } else {
-          try {
-            spanData.span.log({
-              output: eventResult
-            });
-          } catch (error) {
-            console.error(
-              "Error extracting output for Claude Agent SDK:",
-              error
-            );
-          } finally {
-            spanData.span.end();
-            spans.delete(event);
-          }
+          return;
+        }
+        try {
+          state.span.log({ output: eventResult });
+        } catch (error) {
+          console.error("Error extracting output for Claude Agent SDK:", error);
+        } finally {
+          state.span.end();
+          spans.delete(event);
         }
       },
       error: (event) => {
-        const spanData = spans.get(event);
-        if (!spanData || !event.error) {
+        const state = spans.get(event);
+        if (!state || !event.error) {
           return;
         }
-        const { span } = spanData;
-        span.log({
+        state.span.log({
           error: event.error.message
         });
-        span.end();
+        state.span.end();
         spans.delete(event);
       }
     };
@@ -11977,6 +12417,18 @@ var googleGenAIChannels = defineChannels("@google/genai", {
 });
 
 // src/instrumentation/plugins/google-genai-plugin.ts
+var GOOGLE_GENAI_INTERNAL_CONTEXT = {
+  caller_filename: "<node-internal>",
+  caller_functionname: "<node-internal>",
+  caller_lineno: 0
+};
+function createWrapperParityEvent(args) {
+  return {
+    context: GOOGLE_GENAI_INTERNAL_CONTEXT,
+    input: args.input,
+    metadata: args.metadata
+  };
+}
 var GoogleGenAIPlugin = class extends BasePlugin {
   onEnable() {
     this.subscribeToGoogleGenAIChannels();
@@ -11985,51 +12437,282 @@ var GoogleGenAIPlugin = class extends BasePlugin {
     this.unsubscribers = unsubscribeAll(this.unsubscribers);
   }
   subscribeToGoogleGenAIChannels() {
-    this.unsubscribers.push(
-      traceAsyncChannel(googleGenAIChannels.generateContent, {
-        name: "google-genai.generateContent",
-        type: "llm" /* LLM */,
-        extractInput: ([params]) => {
-          const input = serializeInput(params);
-          const metadata = extractMetadata(params);
-          return {
-            input,
-            metadata: { ...metadata, provider: "google-genai" }
-          };
-        },
-        extractOutput: (result) => {
-          return result;
-        },
-        extractMetrics: (result, startTime) => {
-          return extractGenerateContentMetrics(result, startTime);
-        }
-      })
+    this.subscribeToGenerateContentChannel();
+    this.subscribeToGenerateContentStreamChannel();
+  }
+  subscribeToGenerateContentChannel() {
+    const tracingChannel2 = googleGenAIChannels.generateContent.tracingChannel();
+    const states = /* @__PURE__ */ new WeakMap();
+    const unbindCurrentSpanStore = bindCurrentSpanStoreToStart2(
+      tracingChannel2,
+      states,
+      (event) => {
+        const params = event.arguments[0];
+        const input = serializeInput(params);
+        const metadata = extractMetadata(params);
+        const span = startSpan({
+          name: "generate_content",
+          spanAttributes: {
+            type: "llm" /* LLM */
+          },
+          event: createWrapperParityEvent({ input, metadata })
+        });
+        return {
+          span,
+          startTime: getCurrentUnixTimestamp()
+        };
+      }
     );
-    this.unsubscribers.push(
-      traceStreamingChannel(googleGenAIChannels.generateContentStream, {
-        name: "google-genai.generateContentStream",
-        type: "llm" /* LLM */,
-        extractInput: ([params]) => {
+    const handlers = {
+      start: (event) => {
+        ensureSpanState(states, event, () => {
+          const params = event.arguments[0];
           const input = serializeInput(params);
           const metadata = extractMetadata(params);
+          const span = startSpan({
+            name: "generate_content",
+            spanAttributes: {
+              type: "llm" /* LLM */
+            },
+            event: createWrapperParityEvent({ input, metadata })
+          });
           return {
-            input,
-            metadata: { ...metadata, provider: "google-genai" }
+            span,
+            startTime: getCurrentUnixTimestamp()
           };
-        },
-        extractOutput: (result) => {
-          return result;
-        },
-        extractMetrics: () => {
-          return {};
-        },
-        aggregateChunks: (chunks, _result, _endEvent, startTime) => {
-          return aggregateGenerateContentChunks(chunks, startTime);
+        });
+      },
+      asyncEnd: (event) => {
+        const spanState = states.get(event);
+        if (!spanState) {
+          return;
         }
-      })
-    );
+        try {
+          spanState.span.log({
+            metrics: cleanMetrics(
+              extractGenerateContentMetrics(
+                event.result,
+                spanState.startTime
+              )
+            ),
+            output: event.result
+          });
+        } finally {
+          spanState.span.end();
+          states.delete(event);
+        }
+      },
+      error: (event) => {
+        logErrorAndEndSpan(states, event);
+      }
+    };
+    tracingChannel2.subscribe(handlers);
+    this.unsubscribers.push(() => {
+      unbindCurrentSpanStore?.();
+      tracingChannel2.unsubscribe(handlers);
+    });
+  }
+  subscribeToGenerateContentStreamChannel() {
+    const tracingChannel2 = googleGenAIChannels.generateContentStream.tracingChannel();
+    const handlers = {
+      start: (event) => {
+        const streamEvent = event;
+        const params = event.arguments[0];
+        streamEvent.googleGenAIInput = serializeInput(params);
+        streamEvent.googleGenAIMetadata = extractMetadata(params);
+      },
+      asyncEnd: (event) => {
+        const streamEvent = event;
+        patchGoogleGenAIStreamingResult({
+          input: streamEvent.googleGenAIInput,
+          metadata: streamEvent.googleGenAIMetadata,
+          result: streamEvent.result
+        });
+      },
+      error: () => {
+      }
+    };
+    tracingChannel2.subscribe(handlers);
+    this.unsubscribers.push(() => {
+      tracingChannel2.unsubscribe(handlers);
+    });
   }
 };
+function ensureSpanState(states, event, create) {
+  const existing = states.get(event);
+  if (existing) {
+    return existing;
+  }
+  const created = create();
+  states.set(event, created);
+  return created;
+}
+function bindCurrentSpanStoreToStart2(tracingChannel2, states, create) {
+  const state = _internalGetGlobalState();
+  const startChannel = tracingChannel2.start;
+  const currentSpanStore = state?.contextManager ? state.contextManager[BRAINTRUST_CURRENT_SPAN_STORE] : void 0;
+  if (!startChannel?.bindStore || !currentSpanStore) {
+    return void 0;
+  }
+  startChannel.bindStore(
+    currentSpanStore,
+    (event) => ensureSpanState(
+      states,
+      event,
+      () => create(event)
+    ).span
+  );
+  return () => {
+    startChannel.unbindStore?.(currentSpanStore);
+  };
+}
+function logErrorAndEndSpan(states, event) {
+  const spanState = states.get(event);
+  if (!spanState) {
+    return;
+  }
+  spanState.span.log({
+    error: event.error.message
+  });
+  spanState.span.end();
+  states.delete(event);
+}
+function patchGoogleGenAIStreamingResult(args) {
+  const { input, metadata, result } = args;
+  if (!input || !metadata || !result || typeof result !== "object" || typeof result.next !== "function") {
+    return false;
+  }
+  const chunks = [];
+  let firstTokenTime = null;
+  let finalized = false;
+  let span = null;
+  let startTime = null;
+  const ensureSpan = () => {
+    if (!span) {
+      span = startSpan({
+        name: "generate_content_stream",
+        spanAttributes: {
+          type: "llm" /* LLM */
+        },
+        event: {
+          input,
+          metadata
+        }
+      });
+      startTime = getCurrentUnixTimestamp();
+    }
+    return span;
+  };
+  const finalize = (options) => {
+    if (finalized || !span) {
+      return;
+    }
+    finalized = true;
+    if (options.result) {
+      const { end, ...metricsWithoutEnd } = options.result.metrics;
+      span.log({
+        metrics: cleanMetrics(metricsWithoutEnd),
+        output: options.result.aggregated
+      });
+      span.end(typeof end === "number" ? { endTime: end } : void 0);
+      return;
+    }
+    if (options.error !== void 0) {
+      span.log({
+        error: options.error instanceof Error ? options.error.message : String(options.error)
+      });
+    }
+    span.end();
+  };
+  const patchIterator = (iterator) => {
+    if (typeof iterator !== "object" || iterator === null || "__braintrustGoogleGenAIPatched" in iterator) {
+      return iterator;
+    }
+    const iteratorRecord = iterator;
+    const originalNext = typeof iteratorRecord.next === "function" ? iteratorRecord.next.bind(iterator) : void 0;
+    const originalReturn = typeof iteratorRecord.return === "function" ? iteratorRecord.return.bind(iterator) : void 0;
+    const originalThrow = typeof iteratorRecord.throw === "function" ? iteratorRecord.throw.bind(iterator) : void 0;
+    const asyncIteratorMethod = iteratorRecord[Symbol.asyncIterator];
+    const originalAsyncIterator = typeof asyncIteratorMethod === "function" ? asyncIteratorMethod.bind(iterator) : void 0;
+    Object.defineProperty(iteratorRecord, "__braintrustGoogleGenAIPatched", {
+      configurable: true,
+      enumerable: false,
+      value: true,
+      writable: false
+    });
+    if (originalNext) {
+      iteratorRecord.next = async (...nextArgs) => {
+        ensureSpan();
+        try {
+          const nextResult = await originalNext(
+            ...nextArgs
+          );
+          if (!nextResult.done && nextResult.value) {
+            if (firstTokenTime === null) {
+              firstTokenTime = getCurrentUnixTimestamp();
+            }
+            chunks.push(nextResult.value);
+          }
+          if (nextResult.done && startTime !== null) {
+            finalize({
+              result: aggregateGenerateContentChunks(
+                chunks,
+                startTime,
+                firstTokenTime
+              )
+            });
+          }
+          return nextResult;
+        } catch (error) {
+          finalize({ error });
+          throw error;
+        }
+      };
+    }
+    if (originalReturn) {
+      iteratorRecord.return = async (...returnArgs) => {
+        ensureSpan();
+        try {
+          return await originalReturn(
+            ...returnArgs
+          );
+        } finally {
+          if (startTime !== null) {
+            finalize({
+              result: chunks.length > 0 ? aggregateGenerateContentChunks(
+                chunks,
+                startTime,
+                firstTokenTime
+              ) : void 0
+            });
+          } else {
+            finalize({});
+          }
+        }
+      };
+    }
+    if (originalThrow) {
+      iteratorRecord.throw = async (...throwArgs) => {
+        ensureSpan();
+        try {
+          return await originalThrow(
+            ...throwArgs
+          );
+        } catch (error) {
+          finalize({ error });
+          throw error;
+        }
+      };
+    }
+    iteratorRecord[Symbol.asyncIterator] = () => {
+      const asyncIterator = originalAsyncIterator ? originalAsyncIterator() : iterator;
+      return patchIterator(asyncIterator);
+    };
+    return iterator;
+  };
+  patchIterator(result);
+  return true;
+}
 function serializeInput(params) {
   const input = {
     model: params.model,
@@ -12038,11 +12721,13 @@ function serializeInput(params) {
   if (params.config) {
     const config = tryToDict(params.config);
     if (config) {
-      const tools = serializeTools(params);
-      if (tools) {
-        config.tools = tools;
-      }
-      input.config = config;
+      const filteredConfig = {};
+      Object.keys(config).forEach((key) => {
+        if (key !== "tools") {
+          filteredConfig[key] = config[key];
+        }
+      });
+      input.config = filteredConfig;
     }
   }
   return input;
@@ -12129,12 +12814,18 @@ function extractMetadata(params) {
       });
     }
   }
+  const tools = serializeTools(params);
+  if (tools) {
+    metadata.tools = tools;
+  }
   return metadata;
 }
 function extractGenerateContentMetrics(response, startTime) {
   const metrics = {};
-  if (startTime) {
+  if (startTime !== void 0) {
     const end = getCurrentUnixTimestamp();
+    metrics.start = startTime;
+    metrics.end = end;
     metrics.duration = end - startTime;
   }
   if (response?.usageMetadata) {
@@ -12159,19 +12850,18 @@ function populateUsageMetrics(metrics, usage) {
     metrics.completion_reasoning_tokens = usage.thoughtsTokenCount;
   }
 }
-function aggregateGenerateContentChunks(chunks, startTime) {
-  const metrics = {};
-  if (startTime !== void 0) {
-    const end = getCurrentUnixTimestamp();
-    metrics.duration = end - startTime;
-  }
-  let firstTokenTime = null;
-  if (chunks.length > 0 && firstTokenTime === null && startTime !== void 0) {
-    firstTokenTime = getCurrentUnixTimestamp();
+function aggregateGenerateContentChunks(chunks, startTime, firstTokenTime) {
+  const end = getCurrentUnixTimestamp();
+  const metrics = {
+    start: startTime,
+    end,
+    duration: end - startTime
+  };
+  if (firstTokenTime !== null) {
     metrics.time_to_first_token = firstTokenTime - startTime;
   }
   if (chunks.length === 0) {
-    return { output: {}, metrics };
+    return { aggregated: {}, metrics };
   }
   let text = "";
   let thoughtText = "";
@@ -12207,7 +12897,7 @@ function aggregateGenerateContentChunks(chunks, startTime) {
       }
     }
   }
-  const output = {};
+  const aggregated = {};
   const parts = [];
   if (thoughtText) {
     parts.push({ text: thoughtText, thought: true });
@@ -12233,16 +12923,25 @@ function aggregateGenerateContentChunks(chunks, startTime) {
       }
       candidates.push(candidateDict);
     }
-    output.candidates = candidates;
+    aggregated.candidates = candidates;
   }
   if (usageMetadata) {
-    output.usageMetadata = usageMetadata;
+    aggregated.usageMetadata = usageMetadata;
     populateUsageMetrics(metrics, usageMetadata);
   }
   if (text) {
-    output.text = text;
+    aggregated.text = text;
+  }
+  return { aggregated, metrics };
+}
+function cleanMetrics(metrics) {
+  const cleaned = {};
+  for (const [key, value] of Object.entries(metrics)) {
+    if (value !== null && value !== void 0) {
+      cleaned[key] = value;
+    }
   }
-  return { output, metrics };
+  return cleaned;
 }
 function tryToDict(obj) {
   if (obj === null || obj === void 0) {