fix(pi-sdk): Fix nesting of child spans in tool spans

.changeset/pi-tool-current-span.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+Keep Pi Coding Agent tool spans current during execution so nested spans attach to the tool span.

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-auto-hook.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-auto-hook.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-latest-wrapped.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v0-wrapped.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-auto-hook.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-auto-hook.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-latest-wrapped.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/__snapshots__/pi-coding-agent-v081-wrapped.span-tree.txt

@@ -119,3 +119,4 @@ span_tree:
                   "gen_ai.tool.name": "bash",
                   "pi_coding_agent.tool.name": "bash"
                 }
+                └── pi-tool-child

e2e/scenarios/pi-coding-agent-instrumentation/assertions.ts

@@ -13,7 +13,11 @@ import {
   type SpanTreeEntry,
   type SpanTreeFields,
 } from "../../helpers/span-tree";
-import { ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs";
+import {
+  ROOT_NAME,
+  SCENARIO_NAME,
+  TOOL_CHILD_SPAN_NAME,
+} from "./scenario.impl.mjs";
 
 type RunPiCodingAgentScenario = (harness: {
   runNodeScenarioDir: (options: {
@@ -104,13 +108,19 @@ function summarize(events: CapturedLogEvent[]): SpanTreeEntry[] {
     task?.span.id,
   ).at(-1);
   const tool = findChildSpans(events, "bash", task?.span.id).at(-1);
+  const toolChild = findChildSpans(
+    events,
+    TOOL_CHILD_SPAN_NAME,
+    tool?.span.id,
+  ).at(-1);
 
   return [
     findLatestSpan(events, ROOT_NAME),
     operation,
     task,
     llm,
     tool,
+    toolChild,
   ].flatMap((event) =>
     event
       ? [
@@ -181,6 +191,11 @@ export function definePiCodingAgentInstrumentationAssertions(options: {
           task?.span.id,
         ).at(-1);
         const tool = findChildSpans(events, "bash", task?.span.id).at(-1);
+        const toolChild = findChildSpans(
+          events,
+          TOOL_CHILD_SPAN_NAME,
+          tool?.span.id,
+        ).at(-1);
 
         expect(operation).toBeDefined();
         expect(task).toBeDefined();
@@ -221,6 +236,9 @@ export function definePiCodingAgentInstrumentationAssertions(options: {
           expect.any(String),
         );
         expect(JSON.stringify(tool?.output)).toContain("pi_tool_ok");
+
+        expect(toolChild).toBeDefined();
+        expect(toolChild?.span.parentIds).toEqual([tool?.span.id ?? ""]);
       },
     );
 

e2e/scenarios/pi-coding-agent-instrumentation/scenario.impl.mjs

@@ -1,11 +1,12 @@
-import { wrapPiCodingAgentSDK } from "braintrust";
+import { startSpan, wrapPiCodingAgentSDK } from "braintrust";
 import {
   runOperation,
   runTracedScenario,
 } from "../../helpers/provider-runtime.mjs";
 
 export const ROOT_NAME = "pi-coding-agent-root";
 export const SCENARIO_NAME = "pi-coding-agent-instrumentation";
+export const TOOL_CHILD_SPAN_NAME = "pi-tool-child";
 
 async function runPiCodingAgentScenario({ decorateSDK, sdk }) {
   const instrumentedSDK = decorateSDK ? decorateSDK(sdk) : sdk;
@@ -75,6 +76,22 @@ async function runPiCodingAgentScenario({ decorateSDK, sdk }) {
             tools: ["bash"],
           });
           session = result.session;
+          const bashTool = session.agent.state.tools.find(
+            (tool) => tool.name === "bash",
+          );
+          if (!bashTool?.execute) {
+            throw new Error("Expected Pi Coding Agent bash tool executor");
+          }
+          const executeBashTool = bashTool.execute;
+          bashTool.execute = async function (...args) {
+            await Promise.resolve();
+            const childSpan = startSpan({ name: TOOL_CHILD_SPAN_NAME });
+            try {
+              return await Reflect.apply(executeBashTool, this, args);
+            } finally {
+              childSpan.end();
+            }
+          };
           if (useFakeStream) {
             let streamCallCount = 0;
             session.agent.streamFunction = async (streamModel) => {