From cc7e520e686518b29725d7b80018444a4f110e44 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Thu, 13 Aug 2026 15:27:26 +0000 Subject: [PATCH 1/2] fix(claude-agent-sdk): Correctly parent subagent tool spans --- .../assertions.ts | 65 ++++- .../scenario.impl.mjs | 83 ++++-- .../scenario.test.ts | 14 +- .../claude-agent-sdk-plugin.streaming.test.ts | 257 ++++++++++++++++++ .../plugins/claude-agent-sdk-plugin.ts | 18 ++ 5 files changed, 397 insertions(+), 40 deletions(-) create mode 100644 js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts b/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts index c7983324d..da76ed62c 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts +++ b/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; + nodeArgs?: string[]; runContext?: ScenarioRunContext; scenarioDir: string; - timeoutMs: number; - }) => Promise; + timeoutMs?: number; + }) => Promise<{ stdout: string }>; runScenarioDir: (options: { - entry: string; + entry?: string; + env?: Record; runContext?: ScenarioRunContext; scenarioDir: string; - timeoutMs: number; - }) => Promise; -}) => Promise; + 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; 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", diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs index 4d565d2a0..680fa4fa7 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs +++ b/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) { diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts index 44f1044ef..ee79c0b1e 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts +++ b/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, diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts new file mode 100644 index 000000000..42101e066 --- /dev/null +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts @@ -0,0 +1,257 @@ +import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import type { BackgroundLogEvent } from "../../../util/index"; +import { + _exportsForTestingOnly, + initLogger, + type TestBackgroundLogger, +} from "../../logger"; +import { configureNode } from "../../node/config"; +import type { + ClaudeAgentSDKMessage, + ClaudeAgentSDKQueryOptions, + ClaudeAgentSDKQueryParams, +} from "../../vendor-sdk-types/claude-agent-sdk"; +import { wrapClaudeAgentSDK } from "../../wrappers/claude-agent-sdk/claude-agent-sdk"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +type ControlledStream = { + finish: () => void; + push: (value: T) => void; + stream: AsyncIterableIterator; +}; + +type CapturedSpanEvent = BackgroundLogEvent & { + metadata?: Record; + span_attributes?: Record; + span_id?: string; + span_parents?: string[]; +}; + +function isCapturedSpanEvent( + event: BackgroundLogEvent, +): event is CapturedSpanEvent { + return "span_attributes" in event; +} + +function makeControlledStream(): ControlledStream { + const queue: T[] = []; + const waiters: Array<(result: IteratorResult) => void> = []; + let done = false; + + return { + finish() { + done = true; + for (const waiter of waiters.splice(0)) { + waiter({ done: true, value: undefined }); + } + }, + push(value) { + const waiter = waiters.shift(); + if (waiter) { + waiter({ done: false, value }); + } else { + queue.push(value); + } + }, + stream: { + [Symbol.asyncIterator]() { + return this; + }, + next() { + const value = queue.shift(); + if (value !== undefined) { + return Promise.resolve({ done: false as const, value }); + } + if (done) { + return Promise.resolve({ done: true as const, value: undefined }); + } + return new Promise((resolve) => waiters.push(resolve)); + }, + }, + }; +} + +function assistantToolUseMessage(options: { + messageId: string; + parentToolUseId: string | null; + toolName: string; + toolUseId: string; +}): ClaudeAgentSDKMessage { + return { + type: "assistant", + parent_tool_use_id: options.parentToolUseId, + message: { + id: options.messageId, + role: "assistant", + model: "claude-scripted", + usage: { input_tokens: 10, output_tokens: 5 }, + content: [ + { + type: "tool_use", + id: options.toolUseId, + name: options.toolName, + input: + options.toolName === "Task" + ? { + subagent_type: "metadata-checker", + description: "Metadata check", + prompt: "Check the metadata fields.", + } + : { field: "company" }, + }, + ], + }, + }; +} + +describe("Claude Agent SDK streaming instrumentation", () => { + let backgroundLogger: TestBackgroundLogger; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectId: "test-project-id", + projectName: "claude-agent-sdk-plugin.streaming.test.ts", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + it("parents a local tool to its subagent when execution races stream consumption", async () => { + const controlled = makeControlledStream(); + let injectedOptions: ClaudeAgentSDKQueryOptions | undefined; + const sdk = wrapClaudeAgentSDK({ + query: (params: ClaudeAgentSDKQueryParams) => { + injectedOptions = params.options; + return controlled.stream; + }, + tool: ( + _name: string, + _description: string, + _schema: unknown, + handler: (args: TArgs, ...extra: unknown[]) => unknown, + ) => ({ handler }), + createSdkMcpServer: (config: { name: string; tools: unknown[] }) => ({ + type: "sdk" as const, + name: config.name, + instance: {}, + }), + }); + + const getMetadata = sdk.tool( + "get_metadata", + "Returns a synthetic metadata field.", + { field: "string" }, + async ({ field }: { field: string }) => ({ + content: [{ type: "text", text: JSON.stringify({ field }) }], + }), + ); + const server = sdk.createSdkMcpServer({ + name: "skillrepro", + tools: [getMetadata], + }); + const result = sdk.query({ + prompt: "Delegate metadata checking to a subagent.", + options: { + model: "claude-scripted", + mcpServers: { skillrepro: server }, + }, + }); + const iterator = result[Symbol.asyncIterator](); + + const taskToolUseId = "toolu_task_1"; + const localToolUseId = "toolu_mcp_before_consumption"; + const agentId = "agent-metadata-checker"; + + controlled.push( + assistantToolUseMessage({ + messageId: "msg_orchestrator", + parentToolUseId: null, + toolName: "Task", + toolUseId: taskToolUseId, + }), + ); + await iterator.next(); + controlled.push({ + type: "system", + subtype: "task_started", + task_id: agentId, + tool_use_id: taskToolUseId, + description: "Metadata check", + prompt: "Check the metadata fields.", + task_type: "local_agent", + }); + await iterator.next(); + + for (const matcher of injectedOptions?.hooks?.PreToolUse ?? []) { + for (const hook of matcher.hooks) { + await hook( + { + hook_event_name: "PreToolUse", + agent_id: agentId, + cwd: "/test", + session_id: "scripted-session", + tool_input: { field: "company" }, + tool_name: "mcp__skillrepro__get_metadata", + transcript_path: "/test/transcript.jsonl", + }, + localToolUseId, + { signal: new AbortController().signal }, + ); + } + } + + // Reproduce SDK-222: the local handler starts before the application pulls + // the subagent assistant message containing this tool-use ID. + const localToolResult = getMetadata.handler( + { field: "company" }, + { _meta: { "claudecode/toolUseId": localToolUseId } }, + ); + await localToolResult; + + controlled.push( + assistantToolUseMessage({ + messageId: "msg_subagent", + parentToolUseId: taskToolUseId, + toolName: "mcp__skillrepro__get_metadata", + toolUseId: localToolUseId, + }), + ); + await iterator.next(); + + controlled.push({ + type: "result", + num_turns: 2, + session_id: "scripted-session", + usage: { input_tokens: 20, output_tokens: 10 }, + }); + await iterator.next(); + controlled.finish(); + await expect(iterator.next()).resolves.toMatchObject({ done: true }); + + const events = (await backgroundLogger.drain()).filter(isCapturedSpanEvent); + const subagentSpan = events.find((event) => + String(event.span_attributes?.name).startsWith("Agent: "), + ); + const toolSpan = events.find( + (event) => + event.span_attributes?.name === "tool: skillrepro/get_metadata" && + event.metadata?.["gen_ai.tool.call.id"] === localToolUseId, + ); + + expect(subagentSpan).toBeDefined(); + expect(toolSpan).toBeDefined(); + expect(toolSpan?.span_parents).toEqual([subagentSpan?.span_id]); + }); +}); diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index 8a88c5daf..86237a6df 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -493,6 +493,8 @@ function prepareLocalToolHandlersInMcpServers( function createToolTracingHooks( resolveParentSpan: ParentSpanResolver, + taskIdToToolUseId: Map, + toolUseToParent: Map, activeToolSpans: Map, mcpServers: ClaudeAgentSDKMcpServersConfig | undefined, localToolHookNames: Set, @@ -512,6 +514,16 @@ function createToolTracingHooks( return {}; } + if (!toolUseToParent.has(toolUseID) && input.agent_id) { + // Subagent hook inputs identify the same task that lifecycle messages + // expose as task_id, so this correlation is available before the SDK + // yields the assistant tool_use block to the application. + const parentToolUseId = taskIdToToolUseId.get(input.agent_id); + if (parentToolUseId) { + toolUseToParent.set(toolUseID, parentToolUseId); + } + } + if ( skipLocalToolHooks && (isLocalToolUse(input.tool_name, mcpServers) || @@ -787,6 +799,8 @@ function createToolTracingHooks( function injectTracingHooks( options: ClaudeAgentSDKQueryOptions, resolveParentSpan: ParentSpanResolver, + taskIdToToolUseId: Map, + toolUseToParent: Map, activeToolSpans: Map, localToolHookNames: Set, skipLocalToolHooks: boolean, @@ -802,6 +816,8 @@ function injectTracingHooks( subagentStop, } = createToolTracingHooks( resolveParentSpan, + taskIdToToolUseId, + toolUseToParent, activeToolSpans, options.mcpServers, localToolHookNames, @@ -1564,6 +1580,8 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { const optionsWithHooks = injectTracingHooks( options, resolveToolUseParentSpan, + taskIdToToolUseId, + toolUseToParent, activeToolSpans, localToolHookNames, skipLocalToolHooks, From 74b61ecbefb40691c9921f436f188f625a67581f Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Thu, 13 Aug 2026 15:31:59 +0000 Subject: [PATCH 2/2] Update PR #2362 --- .changeset/claude-agent-sdk-subagent-tool-parenting.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/claude-agent-sdk-subagent-tool-parenting.md diff --git a/.changeset/claude-agent-sdk-subagent-tool-parenting.md b/.changeset/claude-agent-sdk-subagent-tool-parenting.md new file mode 100644 index 000000000..503deeb4d --- /dev/null +++ b/.changeset/claude-agent-sdk-subagent-tool-parenting.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix(claude-agent-sdk): Correctly parent subagent tool spans