fix(claude-agent-sdk): Correctly parent subagent tool spans

.changeset/claude-agent-sdk-subagent-tool-parenting.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+fix(claude-agent-sdk): Correctly parent subagent tool spans

e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts

@@ -23,19 +23,45 @@ import { ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs";
 
 type RunClaudeAgentSDKScenario = (harness: {
   runNodeScenarioDir: (options: {
-    entry: string;
-    nodeArgs: string[];
+    entry?: string;
+    env?: Record<string, string>;
+    nodeArgs?: string[];
     runContext?: ScenarioRunContext;
     scenarioDir: string;
-    timeoutMs: number;
-  }) => Promise<unknown>;
+    timeoutMs?: number;
+  }) => Promise<{ stdout: string }>;
   runScenarioDir: (options: {
-    entry: string;
+    entry?: string;
+    env?: Record<string, string>;
     runContext?: ScenarioRunContext;
     scenarioDir: string;
-    timeoutMs: number;
-  }) => Promise<unknown>;
-}) => Promise<void>;
+    timeoutMs?: number;
+  }) => Promise<{ stdout: string }>;
+}) => Promise<{ stdout: string }>;
+
+function parseScenarioResult(stdout: string): {
+  subagentRacedToolUseId: string;
+} {
+  const prefix = "CLAUDE_AGENT_E2E_RESULT=";
+  const resultLine = stdout.split("\n").find((line) => line.startsWith(prefix));
+  if (!resultLine) {
+    throw new Error("Claude Agent SDK scenario did not report its e2e result");
+  }
+
+  const result: unknown = JSON.parse(resultLine.slice(prefix.length));
+  if (
+    !result ||
+    typeof result !== "object" ||
+    !("subagentRacedToolUseId" in result) ||
+    typeof result.subagentRacedToolUseId !== "string"
+  ) {
+    throw new Error(
+      "Claude Agent SDK scenario did not race local tool execution with stream consumption",
+    );
+  }
+
+  return { subagentRacedToolUseId: result.subagentRacedToolUseId };
+}
 
 const SNAPSHOT_METADATA_KEYS = [
   "provider",
@@ -437,10 +463,12 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: {
 
   describe(options.name, () => {
     let events: CapturedLogEvent[] = [];
+    let scenarioResult: ReturnType<typeof parseScenarioResult>;
 
     beforeAll(async () => {
       await withScenarioHarness(async (harness) => {
-        await options.runScenario(harness);
+        const result = await options.runScenario(harness);
+        scenarioResult = parseScenarioResult(result.stdout);
         events = harness.events();
       });
     }, timeoutMs);
@@ -651,6 +679,25 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: {
       expect(tool?.span.parentIds).not.toContain(nestedTaskLlm?.span.id ?? "");
     });
 
+    test(
+      "parents a local subagent tool when execution races stream consumption",
+      testConfig,
+      () => {
+        const taskRoot = findOperationTaskRoot(
+          events,
+          "claude-agent-subagent-operation",
+        );
+        const nestedTask = findSubAgentTaskSpan(events, taskRoot?.span.id);
+        const tool = findAllSpans(events, "tool: calculator/calculator").find(
+          (event) =>
+            event.row.metadata?.["gen_ai.tool.call.id"] ===
+            scenarioResult.subagentRacedToolUseId,
+        );
+        expect(tool).toBeDefined();
+        expect(tool?.span.parentIds).toEqual([nestedTask?.span.id ?? ""]);
+      },
+    );
+
     if (options.expectTaskLifecycleDetails) {
       test(
         "orders built-in Agent and Bash after their llm siblings",

e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs

@@ -25,6 +25,7 @@ function makePromptMessage(content) {
 async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) {
   const instrumentedSDK = decorateSDK ? decorateSDK(sdk) : sdk;
   const { createSdkMcpServer, query, tool } = instrumentedSDK;
+  let subagentRacedToolUseId = null;
   const calculator = tool(
     "calculator",
     "Performs basic arithmetic operations",
@@ -52,9 +53,7 @@ async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) {
               throw new Error(`unsupported operation: ${args.operation}`);
           }
         },
-        {
-          name: `calculator-local-handler-${args.operation}`,
-        },
+        { name: `calculator-local-handler-${args.operation}` },
       );
 
       return {
@@ -115,28 +114,62 @@ async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) {
         "claude-agent-subagent-operation",
         "subagent",
         async () => {
-          await collectAsync(
-            query({
-              prompt:
-                "Spawn a math-expert subagent to add 15 and 27 using the calculator tool. Report the result. Do not solve it yourself.",
-              options: {
-                agents: {
-                  "math-expert": {
-                    description: "Math specialist",
-                    model: CLAUDE_AGENT_MODEL,
-                    prompt:
-                      "You are a math expert. Use the calculator tool for calculations. Be concise.",
-                  },
-                },
-                allowedTools: ["Task"],
-                mcpServers: {
-                  calculator: calculatorServer,
+          const consumedToolUseIds = new Set();
+          const result = query({
+            prompt:
+              "Spawn a math-expert subagent to add 15 and 27 using the calculator tool. Report the result. Do not solve it yourself.",
+            options: {
+              agents: {
+                "math-expert": {
+                  description: "Math specialist",
+                  model: CLAUDE_AGENT_MODEL,
+                  prompt:
+                    "You are a math expert. Use the calculator tool for calculations. Be concise.",
                 },
-                model: CLAUDE_AGENT_MODEL,
-                permissionMode: "bypassPermissions",
               },
-            }),
-          );
+              allowedTools: ["Task"],
+              hooks: {
+                PreToolUse: [
+                  {
+                    hooks: [
+                      async (input, toolUseId) => {
+                        if (
+                          input.tool_name === "mcp__calculator__calculator" &&
+                          input.tool_input?.operation === "add" &&
+                          typeof toolUseId === "string" &&
+                          !consumedToolUseIds.has(toolUseId)
+                        ) {
+                          subagentRacedToolUseId = toolUseId;
+                        }
+                        return {};
+                      },
+                    ],
+                  },
+                ],
+              },
+              mcpServers: {
+                calculator: calculatorServer,
+              },
+              model: CLAUDE_AGENT_MODEL,
+              permissionMode: "bypassPermissions",
+            },
+          });
+
+          for await (const record of result) {
+            if (
+              record.type === "assistant" &&
+              Array.isArray(record.message?.content)
+            ) {
+              for (const block of record.message.content) {
+                if (
+                  block?.type === "tool_use" &&
+                  typeof block.id === "string"
+                ) {
+                  consumedToolUseIds.add(block.id);
+                }
+              }
+            }
+          }
         },
       );
 
@@ -193,6 +226,10 @@ async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) {
     projectNameBase: "e2e-claude-agent-sdk-instrumentation",
     rootName: ROOT_NAME,
   });
+
+  process.stdout.write(
+    `CLAUDE_AGENT_E2E_RESULT=${JSON.stringify({ subagentRacedToolUseId })}\n`,
+  );
 }
 
 export async function runWrappedClaudeAgentSDKInstrumentation(sdk) {

e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts

@@ -46,8 +46,8 @@ describe.concurrent("wrapped instrumentation", () => {
         assertLocalToolHandlerParenting: true,
         expectTaskLifecycleDetails: scenario.expectTaskLifecycleDetails,
         name: "scenario",
-        runScenario: async ({ runScenarioDir }) => {
-          await runScenarioDir({
+        runScenario: ({ runScenarioDir }) =>
+          runScenarioDir({
             entry: scenario.wrapperEntry,
             env: { CLAUDE_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName },
             runContext: {
@@ -56,8 +56,7 @@ describe.concurrent("wrapped instrumentation", () => {
             },
             scenarioDir,
             timeoutMs: TIMEOUT_MS,
-          });
-        },
+          }),
         snapshotName: `${scenario.snapshotName}-wrapped`,
         testFileUrl: import.meta.url,
         timeoutMs: TIMEOUT_MS,
@@ -73,8 +72,8 @@ describe.concurrent("auto-hook instrumentation", () => {
         assertLocalToolHandlerParenting: true,
         expectTaskLifecycleDetails: scenario.expectTaskLifecycleDetails,
         name: "scenario",
-        runScenario: async ({ runNodeScenarioDir }) => {
-          await runNodeScenarioDir({
+        runScenario: ({ runNodeScenarioDir }) =>
+          runNodeScenarioDir({
             entry: scenario.autoEntry,
             env: { CLAUDE_AGENT_SDK_PACKAGE_NAME: scenario.dependencyName },
             nodeArgs: ["--import", "braintrust/hook.mjs"],
@@ -84,8 +83,7 @@ describe.concurrent("auto-hook instrumentation", () => {
             },
             scenarioDir,
             timeoutMs: TIMEOUT_MS,
-          });
-        },
+          }),
         snapshotName: `${scenario.snapshotName}-auto-hook`,
         testFileUrl: import.meta.url,
         timeoutMs: TIMEOUT_MS,