From 9ef47f6a7f562c9e0bdfc61de7ab4a1eb26c6dd3 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Tue, 7 Apr 2026 15:30:39 +0200 Subject: [PATCH 1/4] fix: Fix claude agent SDK tool nesting --- .../claude-agent-sdk-v0.1.span-events.json | 6 +- .../claude-agent-sdk-v0.2.76.span-events.json | 21 +- .../claude-agent-sdk-v0.2.79.span-events.json | 21 +- .../claude-agent-sdk-v0.2.81.span-events.json | 21 +- .../assertions.ts | 116 +++++++-- .../scenario.impl.mjs | 43 ++-- .../scenario.test.ts | 2 + ...ude-agent-sdk-instrumentation-constants.ts | 2 + .../claude-agent-sdk-local-tool-spans.ts | 225 ++++++++++++++++++ .../plugins/claude-agent-sdk-plugin.ts | 217 +++++++++++++++-- .../claude-agent-sdk/claude-agent-sdk.ts | 121 +++++++++- 11 files changed, 724 insertions(+), 71 deletions(-) create mode 100644 js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts create mode 100644 js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json index 57289ec3f..f980b9320 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -194,7 +194,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -290,7 +290,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json index ce4c4a2e6..850acdbca 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -194,7 +194,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -278,6 +278,21 @@ ], "type": "task" }, - "tool": null + "tool": { + "has_input": true, + "has_output": true, + "metadata": { + "gen_ai.tool.name": "calculator", + "mcp.server": "calculator" + }, + "metric_keys": [], + "name": "tool: calculator/calculator", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "tool" + } } } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json index ce4c4a2e6..850acdbca 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -194,7 +194,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -278,6 +278,21 @@ ], "type": "task" }, - "tool": null + "tool": { + "has_input": true, + "has_output": true, + "metadata": { + "gen_ai.tool.name": "calculator", + "mcp.server": "calculator" + }, + "metric_keys": [], + "name": "tool: calculator/calculator", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "tool" + } } } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json index ce4c4a2e6..850acdbca 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -194,7 +194,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -278,6 +278,21 @@ ], "type": "task" }, - "tool": null + "tool": { + "has_input": true, + "has_output": true, + "metadata": { + "gen_ai.tool.name": "calculator", + "mcp.server": "calculator" + }, + "metric_keys": [], + "name": "tool: calculator/calculator", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "tool" + } } } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts b/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts index 2e2985821..79f789ac9 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/assertions.ts @@ -92,6 +92,35 @@ function summarizeSpan( return summary; } +function findToolSpanByOperation( + events: CapturedLogEvent[], + operation: "add" | "divide" | "multiply" | "subtract", +): CapturedLogEvent | undefined { + return findAllSpans(events, "tool: calculator/calculator").find((event) => { + const input = event.input as { operation?: string } | undefined; + return input?.operation === operation; + }); +} + +function findToolSpanByLocalHandler( + events: CapturedLogEvent[], + handlerSpanName: string, +): CapturedLogEvent | undefined { + const handlerSpan = findAllSpans(events, handlerSpanName).at(-1); + if (!handlerSpan) { + return undefined; + } + + const parentId = handlerSpan.span.parentIds[0]; + if (!parentId) { + return undefined; + } + + return findAllSpans(events, "tool: calculator/calculator").find( + (event) => event.span.id === parentId, + ); +} + function buildSpanSummary(events: CapturedLogEvent[]): Json { const root = findLatestSpan(events, ROOT_NAME); const basicOperation = findLatestSpan(events, "claude-agent-basic-operation"); @@ -151,21 +180,21 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json { failureTask?.span.id, ).at(-1); - const basicTool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(basicTask?.span.id ?? ""), - ); + const basicTool = + findToolSpanByLocalHandler(events, "calculator-local-handler-multiply") ?? + findToolSpanByOperation(events, "multiply"); const subAgentTask = events.find( (event) => event.span.type === "task" && event.span.parentIds.includes(subAgentTaskRoot?.span.id ?? "") && event.span.name?.startsWith("Agent:"), ); - const subAgentTool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(subAgentTask?.span.id ?? ""), - ); - const failureTool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(failureTask?.span.id ?? ""), - ); + const subAgentTool = + findToolSpanByLocalHandler(events, "calculator-local-handler-add") ?? + findToolSpanByOperation(events, "add"); + const failureTool = + findToolSpanByLocalHandler(events, "calculator-local-handler-divide") ?? + findToolSpanByOperation(events, "divide"); return normalizeForSnapshot({ async_prompt: { @@ -197,6 +226,7 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json { } export function defineClaudeAgentSDKInstrumentationAssertions(options: { + assertLocalToolHandlerParenting?: boolean; name: string; runScenario: RunClaudeAgentSDKScenario; snapshotName: string; @@ -243,18 +273,59 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: { "anthropic.messages.create", task?.span.id, ).at(-1); - const tool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(task?.span.id ?? ""), - ); + const tool = + findToolSpanByLocalHandler( + events, + "calculator-local-handler-multiply", + ) ?? findToolSpanByOperation(events, "multiply"); expect(operation).toBeDefined(); expect(task).toBeDefined(); expect(llm).toBeDefined(); expect(tool).toBeDefined(); expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); - expect(tool?.span.parentIds).toEqual([task?.span.id ?? ""]); }); + if (options.assertLocalToolHandlerParenting) { + test( + "nests local tool handler spans under tool spans", + testConfig, + () => { + const basicTool = + findToolSpanByLocalHandler( + events, + "calculator-local-handler-multiply", + ) ?? findToolSpanByOperation(events, "multiply"); + const basicHandler = findAllSpans( + events, + "calculator-local-handler-multiply", + ).at(-1); + + const failureTool = + findToolSpanByLocalHandler( + events, + "calculator-local-handler-divide", + ) ?? findToolSpanByOperation(events, "divide"); + const failureHandler = findAllSpans( + events, + "calculator-local-handler-divide", + ).at(-1); + + expect(basicTool).toBeDefined(); + expect(basicHandler).toBeDefined(); + expect(basicHandler?.span.parentIds).toEqual([ + basicTool?.span.id ?? "", + ]); + + expect(failureTool).toBeDefined(); + expect(failureHandler).toBeDefined(); + expect(failureHandler?.span.parentIds).toEqual([ + failureTool?.span.id ?? "", + ]); + }, + ); + } + test( "captures async prompt input on both task and llm spans", testConfig, @@ -309,8 +380,11 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: { event.span.parentIds.includes(taskRoot?.span.id ?? "") && event.span.name?.startsWith("Agent:"), ); - const tool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(nestedTask?.span.id ?? ""), + const tool = + findToolSpanByLocalHandler(events, "calculator-local-handler-add") ?? + findToolSpanByOperation(events, "add"); + const toolParent = events.find( + (event) => event.span.id === tool?.span.parentIds[0], ); expect(operation).toBeDefined(); @@ -318,8 +392,12 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: { expect(llm).toBeDefined(); expect(nestedTask).toBeDefined(); if (tool) { - expect(tool.span.parentIds).toContain(nestedTask?.span.id ?? ""); expect(tool.span.parentIds).not.toContain(taskRoot?.span.id ?? ""); + if (toolParent?.span.type === "llm") { + expect(toolParent.span.parentIds).not.toContain( + taskRoot?.span.id ?? "", + ); + } } }); @@ -338,9 +416,9 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: { "anthropic.messages.create", task?.span.id, ).at(-1); - const tool = findAllSpans(events, "tool: calculator/calculator").find( - (event) => event.span.parentIds.includes(task?.span.id ?? ""), - ); + const tool = + findToolSpanByLocalHandler(events, "calculator-local-handler-divide") ?? + findToolSpanByOperation(events, "divide"); expect(operation).toBeDefined(); expect(task).toBeDefined(); diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs index 879c14416..66a5c1bc2 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.impl.mjs @@ -1,4 +1,4 @@ -import { wrapClaudeAgentSDK } from "braintrust"; +import { traced, wrapClaudeAgentSDK } from "braintrust"; import { collectAsync, runOperation, @@ -33,27 +33,28 @@ async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) { b: z.number(), }, async (args) => { - let result; - - switch (args.operation) { - case "add": - result = args.a + args.b; - break; - case "subtract": - result = args.a - args.b; - break; - case "multiply": - result = args.a * args.b; - break; - case "divide": - if (args.b === 0) { - throw new Error("division by zero"); + const result = await traced( + async () => { + switch (args.operation) { + case "add": + return args.a + args.b; + case "subtract": + return args.a - args.b; + case "multiply": + return args.a * args.b; + case "divide": + if (args.b === 0) { + throw new Error("division by zero"); + } + return args.a / args.b; + default: + throw new Error(`unsupported operation: ${args.operation}`); } - result = args.a / args.b; - break; - default: - throw new Error(`unsupported operation: ${args.operation}`); - } + }, + { + name: `calculator-local-handler-${args.operation}`, + }, + ); return { content: [ diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts index b4e530c98..52f3703e5 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/scenario.test.ts @@ -48,6 +48,7 @@ const claudeAgentSDKScenarios = await Promise.all( describe("wrapped instrumentation", () => { for (const scenario of claudeAgentSDKScenarios) { defineClaudeAgentSDKInstrumentationAssertions({ + assertLocalToolHandlerParenting: true, name: `claude agent sdk ${scenario.version}`, runScenario: async ({ runScenarioDir }) => { await runScenarioDir({ @@ -67,6 +68,7 @@ describe("wrapped instrumentation", () => { describe("auto-hook instrumentation", () => { for (const scenario of claudeAgentSDKScenarios) { defineClaudeAgentSDKInstrumentationAssertions({ + assertLocalToolHandlerParenting: true, name: `claude agent sdk ${scenario.version}`, runScenario: async ({ runNodeScenarioDir }) => { await runNodeScenarioDir({ diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts b/js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts new file mode 100644 index 000000000..9a5125510 --- /dev/null +++ b/js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts @@ -0,0 +1,2 @@ +export const CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION = + "__braintrust_skip_local_tool_hooks"; diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts new file mode 100644 index 000000000..b84a0ddac --- /dev/null +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts @@ -0,0 +1,225 @@ +import { startSpan, withCurrent } from "../../logger"; +import { SpanTypeAttribute } from "../../../util/index"; + +export type LocalToolSpanMetadata = { + serverName?: string; + toolName: string; +}; + +type LocalToolHandler = (...args: unknown[]) => unknown; + +const LOCAL_TOOL_HANDLER_WRAPPED = Symbol.for( + "braintrust.claude_agent_sdk.local_tool_handler_wrapped", +); + +function toErrorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isPromiseLike(value: unknown): value is Promise { + return ( + value !== null && + typeof value === "object" && + "then" in value && + typeof value.then === "function" + ); +} + +export function wrapLocalClaudeToolHandler( + handler: LocalToolHandler, + getMetadata: () => LocalToolSpanMetadata, +): LocalToolHandler { + if ( + (handler as LocalToolHandler & { [LOCAL_TOOL_HANDLER_WRAPPED]?: boolean })[ + LOCAL_TOOL_HANDLER_WRAPPED + ] + ) { + return handler; + } + + const wrappedHandler: LocalToolHandler = function wrappedLocalToolHandler( + this: unknown, + ...handlerArgs: unknown[] + ) { + const metadata = getMetadata(); + const spanName = metadata.serverName + ? `tool: ${metadata.serverName}/${metadata.toolName}` + : `tool: ${metadata.toolName}`; + const span = startSpan({ + event: { + input: handlerArgs[0], + metadata: { + "gen_ai.tool.name": metadata.toolName, + ...(metadata.serverName && { + "mcp.server": metadata.serverName, + }), + }, + }, + name: spanName, + spanAttributes: { type: SpanTypeAttribute.TOOL }, + }); + + const runHandler = () => Reflect.apply(handler, this, handlerArgs); + const finalizeSuccess = (result: unknown) => { + span.log({ output: result }); + span.end(); + return result; + }; + const finalizeError = (error: unknown) => { + span.log({ error: toErrorMessage(error) }); + span.end(); + throw error; + }; + + return withCurrent(span, () => { + try { + const result = runHandler(); + if (isPromiseLike(result)) { + return result.then(finalizeSuccess, finalizeError); + } + return finalizeSuccess(result); + } catch (error) { + return finalizeError(error); + } + }); + }; + + Object.defineProperty(wrappedHandler, LOCAL_TOOL_HANDLER_WRAPPED, { + configurable: false, + enumerable: false, + value: true, + writable: false, + }); + + return wrappedHandler; +} + +type LocalToolRegistration = { + [key: string]: unknown; + handler?: unknown; +}; + +function getRegisteredTools( + instance: unknown, +): + | Map + | Record + | undefined { + if (!instance || typeof instance !== "object") { + return undefined; + } + + if (!("_registeredTools" in instance)) { + return undefined; + } + + const registeredTools = Reflect.get(instance, "_registeredTools"); + if (registeredTools instanceof Map) { + return registeredTools as Map; + } + + if (registeredTools && typeof registeredTools === "object") { + return registeredTools as Record; + } + + return undefined; +} + +export function wrapLocalMcpServerToolHandlers( + serverName: string, + serverConfig: unknown, +): boolean { + if (!serverConfig || typeof serverConfig !== "object") { + return false; + } + + if (!("instance" in serverConfig)) { + return false; + } + + const instance = Reflect.get(serverConfig, "instance"); + const registeredTools = getRegisteredTools(instance); + if (!registeredTools) { + return false; + } + + let wrappedAny = false; + const wrapHandler = (toolName: string, registration: unknown) => { + if (!registration || typeof registration !== "object") { + return; + } + + const handler = Reflect.get(registration, "handler"); + if (typeof handler !== "function") { + return; + } + + const wrappedHandler = wrapLocalClaudeToolHandler(handler, () => ({ + serverName, + toolName, + })); + if (wrappedHandler !== handler) { + Reflect.set(registration, "handler", wrappedHandler); + wrappedAny = true; + } + }; + + if (registeredTools instanceof Map) { + for (const [toolName, registration] of registeredTools.entries()) { + wrapHandler(toolName, registration); + } + return wrappedAny; + } + + for (const [toolName, registration] of Object.entries(registeredTools)) { + wrapHandler(toolName, registration); + } + + return wrappedAny; +} + +export function collectLocalMcpServerToolHookNames( + serverName: string, + serverConfig: unknown, +): Set { + const toolNames = new Set(); + + if (!serverConfig || typeof serverConfig !== "object") { + return toolNames; + } + + if ("instance" in serverConfig) { + const instance = Reflect.get(serverConfig, "instance"); + const registeredTools = getRegisteredTools(instance); + if (registeredTools instanceof Map) { + for (const toolName of registeredTools.keys()) { + toolNames.add(toolName); + toolNames.add(`mcp__${serverName}__${toolName}`); + } + } else if (registeredTools) { + for (const toolName of Object.keys(registeredTools)) { + toolNames.add(toolName); + toolNames.add(`mcp__${serverName}__${toolName}`); + } + } + } + + if ("tools" in serverConfig) { + const rawTools = Reflect.get(serverConfig, "tools"); + if (Array.isArray(rawTools)) { + for (const tool of rawTools) { + if (!tool || typeof tool !== "object") { + continue; + } + const toolName = Reflect.get(tool, "name"); + if (typeof toolName !== "string") { + continue; + } + toolNames.add(toolName); + toolNames.add(`mcp__${serverName}__${toolName}`); + } + } + } + + return toolNames; +} diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index e7c1fa8da..5b2b17860 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -11,6 +11,11 @@ import { finalizeAnthropicTokens, } from "../../wrappers/anthropic-tokens-util"; import { claudeAgentSDKChannels } from "./claude-agent-sdk-channels"; +import { CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION } from "./claude-agent-sdk-instrumentation-constants"; +import { + collectLocalMcpServerToolHookNames, + wrapLocalMcpServerToolHandlers, +} from "./claude-agent-sdk-local-tool-spans"; import type { ClaudeAgentSDKHookCallback, ClaudeAgentSDKHookCallbackMatcher, @@ -28,6 +33,15 @@ type ParsedToolName = { toolName: string; }; type ParentSpanResolver = (toolUseID: string) => Promise; +type LLMSpanResult = { + finalMessage: ClaudeConversationMessage | undefined; + spanExport: string; +}; +const ROOT_LLM_PARENT_KEY = "__root__"; + +function llmParentKey(parentToolUseId: string | null): string { + return parentToolUseId ?? ROOT_LLM_PARENT_KEY; +} function isSubAgentToolName(toolName: string): boolean { return toolName === "Agent" || toolName === "Task"; @@ -153,7 +167,8 @@ async function createLLMSpanForMessages( startTime: number, capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined, parentSpan: string, -): Promise { + existingSpan?: Span, +): Promise { if (messages.length === 0) { return undefined; } @@ -181,14 +196,16 @@ async function createLLMSpanForMessages( c !== undefined, ); - const span = startSpan({ - name: "anthropic.messages.create", - parent: parentSpan, - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - startTime, - }); + const span = + existingSpan ?? + startSpan({ + name: "anthropic.messages.create", + parent: parentSpan, + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + startTime, + }); span.log({ input, @@ -197,11 +214,18 @@ async function createLLMSpanForMessages( output: outputs, }); + const spanExport = await span.export(); await span.end(); - return lastMessage.message?.content && lastMessage.message?.role - ? { content: lastMessage.message.content, role: lastMessage.message.role } - : undefined; + const finalMessage = + lastMessage.message?.content && lastMessage.message?.role + ? { content: lastMessage.message.content, role: lastMessage.message.role } + : undefined; + + return { + finalMessage, + spanExport, + }; } function getMcpServerMetadata( @@ -259,10 +283,61 @@ function parseToolName(rawToolName: string): ParsedToolName { }; } +function isLocalToolUse( + rawToolName: string, + mcpServers: ClaudeAgentSDKMcpServersConfig | undefined, +): boolean { + const parsed = parseToolName(rawToolName); + if (!parsed.mcpServer || !mcpServers) { + return false; + } + + const serverConfig = mcpServers[parsed.mcpServer]; + if (!serverConfig || typeof serverConfig !== "object") { + return false; + } + + return serverConfig.type === "sdk" || "transport" in serverConfig; +} + +function prepareLocalToolHandlersInMcpServers( + mcpServers: ClaudeAgentSDKMcpServersConfig | undefined, +): { hasLocalToolHandlers: boolean; localToolHookNames: Set } { + const localToolHookNames = new Set(); + if (!mcpServers) { + return { + hasLocalToolHandlers: false, + localToolHookNames, + }; + } + + let hasLocalToolHandlers = false; + for (const [serverName, serverConfig] of Object.entries(mcpServers)) { + const toolNames = collectLocalMcpServerToolHookNames( + serverName, + serverConfig, + ); + for (const toolName of toolNames) { + localToolHookNames.add(toolName); + } + if (toolNames.size > 0) { + hasLocalToolHandlers = true; + } + + if (wrapLocalMcpServerToolHandlers(serverName, serverConfig)) { + hasLocalToolHandlers = true; + } + } + + return { hasLocalToolHandlers, localToolHookNames }; +} + function createToolTracingHooks( resolveParentSpan: ParentSpanResolver, activeToolSpans: Map, mcpServers: ClaudeAgentSDKMcpServersConfig | undefined, + localToolHookNames: Set, + skipLocalToolHooks: boolean, subAgentSpans: Map, endedSubAgentSpans: Set, ): { @@ -275,6 +350,14 @@ function createToolTracingHooks( return {}; } + if ( + skipLocalToolHooks && + (isLocalToolUse(input.tool_name, mcpServers) || + localToolHookNames.has(input.tool_name)) + ) { + return {}; + } + if (isSubAgentToolName(input.tool_name)) { return {}; } @@ -307,6 +390,14 @@ function createToolTracingHooks( return {}; } + if ( + skipLocalToolHooks && + (isLocalToolUse(input.tool_name, mcpServers) || + localToolHookNames.has(input.tool_name)) + ) { + return {}; + } + const subAgentSpan = subAgentSpans.get(toolUseID); if (subAgentSpan) { try { @@ -360,6 +451,14 @@ function createToolTracingHooks( return {}; } + if ( + skipLocalToolHooks && + (isLocalToolUse(input.tool_name, mcpServers) || + localToolHookNames.has(input.tool_name)) + ) { + return {}; + } + const subAgentSpan = subAgentSpans.get(toolUseID); if (subAgentSpan) { try { @@ -404,6 +503,8 @@ function injectTracingHooks( options: ClaudeAgentSDKQueryOptions, resolveParentSpan: ParentSpanResolver, activeToolSpans: Map, + localToolHookNames: Set, + skipLocalToolHooks: boolean, subAgentSpans: Map, endedSubAgentSpans: Set, ): ClaudeAgentSDKQueryOptions { @@ -412,6 +513,8 @@ function injectTracingHooks( resolveParentSpan, activeToolSpans, options.mcpServers, + localToolHookNames, + skipLocalToolHooks, subAgentSpans, endedSubAgentSpans, ); @@ -442,6 +545,7 @@ function injectTracingHooks( type QueryState = { accumulatedOutputTokens: number; + activeLlmSpansByParentToolUse: Map; activeToolSpans: Map; capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined; currentMessageId: string | undefined; @@ -457,6 +561,8 @@ type QueryState = { promptStarted: () => boolean; span: Span; subAgentSpans: Map; + latestLlmParentBySubAgentToolUse: Map; + latestRootLlmParentRef: { value: string | undefined }; toolUseToParent: Map; }; @@ -466,6 +572,7 @@ async function finalizeCurrentMessageGroup(state: QueryState): Promise { } const parentToolUseId = state.currentMessages[0]?.parent_tool_use_id ?? null; + const parentKey = llmParentKey(parentToolUseId); let parentSpan = await state.span.export(); if (parentToolUseId) { const subAgentSpan = state.subAgentSpans.get(parentToolUseId); @@ -473,8 +580,9 @@ async function finalizeCurrentMessageGroup(state: QueryState): Promise { parentSpan = await subAgentSpan.export(); } } + const existingLlmSpan = state.activeLlmSpansByParentToolUse.get(parentKey); - const finalMessage = await createLLMSpanForMessages( + const llmSpanResult = await createLLMSpanForMessages( state.currentMessages, state.originalPrompt, state.finalResults, @@ -482,10 +590,23 @@ async function finalizeCurrentMessageGroup(state: QueryState): Promise { state.currentMessageStartTime, state.capturedPromptMessages, parentSpan, + existingLlmSpan, ); + state.activeLlmSpansByParentToolUse.delete(parentKey); + + if (llmSpanResult) { + if (parentToolUseId) { + state.latestLlmParentBySubAgentToolUse.set( + parentToolUseId, + llmSpanResult.spanExport, + ); + } else { + state.latestRootLlmParentRef.value = llmSpanResult.spanExport; + } - if (finalMessage) { - state.finalResults.push(finalMessage); + if (llmSpanResult.finalMessage) { + state.finalResults.push(llmSpanResult.finalMessage); + } } const lastMessage = state.currentMessages[state.currentMessages.length - 1]; @@ -677,6 +798,11 @@ async function finalizeQuerySpan(state: QueryState): Promise { } } } finally { + for (const llmSpan of state.activeLlmSpansByParentToolUse.values()) { + llmSpan.end(); + } + state.activeLlmSpansByParentToolUse.clear(); + for (const [id, subAgentSpan] of state.subAgentSpans) { if (!state.endedSubAgentSpans.has(id)) { subAgentSpan.end(); @@ -761,26 +887,80 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { } const activeToolSpans = new Map(); + const activeLlmSpansByParentToolUse = new Map(); const subAgentSpans = new Map(); const endedSubAgentSpans = new Set(); const toolUseToParent = new Map(); + const latestLlmParentBySubAgentToolUse = new Map(); + const latestRootLlmParentRef = { + value: undefined as string | undefined, + }; const pendingSubAgentNames = new Map(); + const { hasLocalToolHandlers, localToolHookNames } = + prepareLocalToolHandlersInMcpServers(options.mcpServers); + const skipLocalToolHooks = + options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true || + hasLocalToolHandlers; const optionsWithHooks = injectTracingHooks( options, async (toolUseID) => { const parentToolUseId = toolUseToParent.get(toolUseID); + const parentKey = llmParentKey(parentToolUseId ?? null); + const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); + if (activeLlmSpan) { + return activeLlmSpan.export(); + } + if (parentToolUseId) { + const parentLlm = + latestLlmParentBySubAgentToolUse.get(parentToolUseId); + if (parentLlm) { + return parentLlm; + } + const subAgentSpan = await ensureSubAgentSpan( pendingSubAgentNames, span, subAgentSpans, parentToolUseId, ); - return subAgentSpan.export(); + const parentSpan = await subAgentSpan.export(); + const llmSpan = startSpan({ + name: "anthropic.messages.create", + parent: parentSpan, + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + startTime: getCurrentUnixTimestamp(), + }); + activeLlmSpansByParentToolUse.set(parentKey, llmSpan); + const llmSpanExport = await llmSpan.export(); + latestLlmParentBySubAgentToolUse.set( + parentToolUseId, + llmSpanExport, + ); + return llmSpanExport; + } + if (latestRootLlmParentRef.value) { + return latestRootLlmParentRef.value; } - return span.export(); + const parentSpan = await span.export(); + const llmSpan = startSpan({ + name: "anthropic.messages.create", + parent: parentSpan, + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + startTime: getCurrentUnixTimestamp(), + }); + activeLlmSpansByParentToolUse.set(parentKey, llmSpan); + const llmSpanExport = await llmSpan.export(); + latestRootLlmParentRef.value = llmSpanExport; + return llmSpanExport; }, activeToolSpans, + localToolHookNames, + skipLocalToolHooks, subAgentSpans, endedSubAgentSpans, ); @@ -790,6 +970,7 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { spans.set(event, { accumulatedOutputTokens: 0, + activeLlmSpansByParentToolUse, activeToolSpans, capturedPromptMessages, currentMessageId: undefined, @@ -805,6 +986,8 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { promptStarted: () => promptStarted, span, subAgentSpans, + latestLlmParentBySubAgentToolUse, + latestRootLlmParentRef, toolUseToParent, }); }, diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts index 28f2af611..8ba922048 100644 --- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts +++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts @@ -1,9 +1,16 @@ import { claudeAgentSDKChannels } from "../../instrumentation/plugins/claude-agent-sdk-channels"; +import { CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION } from "../../instrumentation/plugins/claude-agent-sdk-instrumentation-constants"; +import { wrapLocalClaudeToolHandler } from "../../instrumentation/plugins/claude-agent-sdk-local-tool-spans"; import type { ClaudeAgentSDKModule, ClaudeAgentSDKQueryParams, } from "../../vendor-sdk-types/claude-agent-sdk"; +type LocalToolMetadata = { + serverName?: string; + toolName: string; +}; + /** * Wraps the Claude Agent SDK with Braintrust tracing. Query calls only publish * tracing-channel events; the Claude Agent SDK plugin owns all span lifecycle @@ -36,25 +43,115 @@ function wrapClaudeAgentQuery( const proxy = new Proxy(queryFn, { apply(target, thisArg, argArray) { const params = (argArray[0] ?? {}) as ClaudeAgentSDKQueryParams; + const wrappedParams: ClaudeAgentSDKQueryParams = { + ...params, + options: { + ...(params.options ?? {}), + [CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION]: true, + }, + }; const invocationTarget: unknown = thisArg === proxy || thisArg === undefined ? (defaultThis ?? thisArg) : thisArg; return claudeAgentSDKChannels.query.traceSync( // Async iterator shenanigans are handled purely in the plugin which consumes this channel emission. - () => Reflect.apply(target, invocationTarget, [params]), + () => Reflect.apply(target, invocationTarget, [wrappedParams]), // The channel carries no extra context fields, but the generated // StartOf<> type for Record is overly strict here. - { arguments: [params] } as never, + { arguments: [wrappedParams] } as never, + ); + }, + }); + + return proxy; +} + +function wrapClaudeAgentTool( + toolFn: ClaudeAgentSDKModule["tool"], + localToolMetadataByTool: WeakMap, + defaultThis?: unknown, +): ClaudeAgentSDKModule["tool"] { + const proxy = new Proxy(toolFn, { + apply(target, thisArg, argArray) { + const invocationTarget: unknown = + thisArg === proxy || thisArg === undefined + ? (defaultThis ?? thisArg) + : thisArg; + const wrappedArgs = [...argArray]; + + const toolName = wrappedArgs[0]; + let handlerIndex = -1; + for (let i = wrappedArgs.length - 1; i >= 0; i -= 1) { + if (typeof wrappedArgs[i] === "function") { + handlerIndex = i; + break; + } + } + if (typeof toolName !== "string" || handlerIndex === -1) { + return Reflect.apply(target, invocationTarget, wrappedArgs); + } + + const localToolMetadata: LocalToolMetadata = { toolName }; + const originalHandler = wrappedArgs[handlerIndex] as ( + ...args: unknown[] + ) => unknown; + wrappedArgs[handlerIndex] = wrapLocalClaudeToolHandler( + originalHandler, + () => localToolMetadata, ); + + const wrappedTool = Reflect.apply(target, invocationTarget, wrappedArgs); + if (wrappedTool && typeof wrappedTool === "object") { + localToolMetadataByTool.set(wrappedTool, localToolMetadata); + } + + return wrappedTool; }, }); return proxy; } +function wrapCreateSdkMcpServer( + createSdkMcpServerFn: (...args: unknown[]) => unknown, + localToolMetadataByTool: WeakMap, + defaultThis?: unknown, +): (...args: unknown[]) => unknown { + const proxy = new Proxy(createSdkMcpServerFn, { + apply(target, thisArg, argArray) { + const invocationTarget: unknown = + thisArg === proxy || thisArg === undefined + ? (defaultThis ?? thisArg) + : thisArg; + const config = argArray[0] as + | { name?: unknown; tools?: unknown[] } + | undefined; + const serverName = config?.name; + + if (typeof serverName === "string" && Array.isArray(config?.tools)) { + for (const tool of config.tools) { + if (!tool || typeof tool !== "object") { + continue; + } + + const metadata = localToolMetadataByTool.get(tool); + if (metadata) { + metadata.serverName = serverName; + } + } + } + + return Reflect.apply(target, invocationTarget, argArray); + }, + }); + + return proxy as (...args: unknown[]) => unknown; +} + function claudeAgentSDKProxy(sdk: ClaudeAgentSDKModule): ClaudeAgentSDKModule { const cache = new Map(); + const localToolMetadataByTool = new WeakMap(); return new Proxy(sdk, { get(target, prop, receiver) { @@ -70,6 +167,26 @@ function claudeAgentSDKProxy(sdk: ClaudeAgentSDKModule): ClaudeAgentSDKModule { return wrappedQuery; } + if (prop === "tool" && typeof value === "function") { + const wrappedTool = wrapClaudeAgentTool( + target.tool, + localToolMetadataByTool, + target, + ); + cache.set(prop, wrappedTool); + return wrappedTool; + } + + if (prop === "createSdkMcpServer" && typeof value === "function") { + const wrappedCreateSdkMcpServer = wrapCreateSdkMcpServer( + value as (...args: unknown[]) => unknown, + localToolMetadataByTool, + target, + ); + cache.set(prop, wrappedCreateSdkMcpServer); + return wrappedCreateSdkMcpServer; + } + if (typeof value === "function") { const bound = value.bind(target); cache.set(prop, bound); From 7067d867e6f680fa8c59c74976ed6e857ee03f48 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Wed, 8 Apr 2026 10:53:15 +0200 Subject: [PATCH 2/4] fix nesting --- .../claude-agent-sdk-local-tool-context.ts | 176 ++++++++++++++++++ .../claude-agent-sdk-local-tool-spans.ts | 99 ++++++---- .../plugins/claude-agent-sdk-plugin.ts | 110 ++++++----- .../claude-agent-sdk/claude-agent-sdk.test.ts | 54 +++++- .../claude-agent-sdk/claude-agent-sdk.ts | 1 - 5 files changed, 353 insertions(+), 87 deletions(-) create mode 100644 js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts new file mode 100644 index 000000000..d5fc71b74 --- /dev/null +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts @@ -0,0 +1,176 @@ +import iso from "../../isomorph"; + +type LocalToolParentResolver = (toolUseId: string) => Promise; + +export type ClaudeAgentSDKLocalToolContext = { + resolveLocalToolParent?: LocalToolParentResolver; +}; + +const LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED = Symbol.for( + "braintrust.claude_agent_sdk.local_tool_context_async_iterator_patched", +); + +type AsyncLocalStorageLike = { + enterWith: (store: T) => void; + getStore: () => T | undefined; + run: (store: T, callback: () => R) => R; +}; + +function createLocalToolContextStore(): AsyncLocalStorageLike { + const maybeIsoWithAsyncLocalStorage = iso as { + newAsyncLocalStorage?: () => AsyncLocalStorageLike; + }; + + if ( + typeof maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage === "function" + ) { + return maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage(); + } + + let currentStore: ClaudeAgentSDKLocalToolContext | undefined; + return { + enterWith(store) { + currentStore = store; + }, + getStore() { + return currentStore; + }, + run(store, callback) { + const previousStore = currentStore; + currentStore = store; + try { + return callback(); + } finally { + currentStore = previousStore; + } + }, + }; +} + +const localToolContextStore = createLocalToolContextStore(); +let fallbackLocalToolParentResolver: LocalToolParentResolver | undefined; + +export function createClaudeLocalToolContext(): ClaudeAgentSDKLocalToolContext { + return {}; +} + +export function runWithClaudeLocalToolContext( + callback: () => R, + context?: ClaudeAgentSDKLocalToolContext, +): R { + return localToolContextStore.run( + context ?? createClaudeLocalToolContext(), + callback, + ); +} + +export function ensureClaudeLocalToolContext(): + | ClaudeAgentSDKLocalToolContext + | undefined { + const existing = localToolContextStore.getStore(); + if (existing) { + return existing; + } + + const created: ClaudeAgentSDKLocalToolContext = {}; + localToolContextStore.enterWith(created); + return created; +} + +export function setClaudeLocalToolParentResolver( + resolver: LocalToolParentResolver, +): void { + fallbackLocalToolParentResolver = resolver; + const context = ensureClaudeLocalToolContext(); + if (!context) { + return; + } + context.resolveLocalToolParent = resolver; +} + +export function getClaudeLocalToolParentResolver(): + | LocalToolParentResolver + | undefined { + return ( + localToolContextStore.getStore()?.resolveLocalToolParent ?? + fallbackLocalToolParentResolver + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + value !== null && + typeof value === "object" && + Symbol.asyncIterator in value && + typeof value[Symbol.asyncIterator] === "function" + ); +} + +export function bindClaudeLocalToolContextToAsyncIterable( + result: T, + localToolContext: ClaudeAgentSDKLocalToolContext, +): T { + if ( + !isAsyncIterable(result) || + Object.isFrozen(result) || + Object.isSealed(result) + ) { + return result; + } + + const stream = result as AsyncIterable & { + [Symbol.asyncIterator]: (() => AsyncIterator) & { + [LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]?: boolean; + }; + }; + const originalAsyncIterator = stream[Symbol.asyncIterator]; + if (originalAsyncIterator[LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]) { + return result; + } + + const patchedAsyncIterator = function (this: unknown) { + return runWithClaudeLocalToolContext(() => { + const iterator = Reflect.apply(originalAsyncIterator, this, []); + if (!iterator || typeof iterator !== "object") { + return iterator; + } + + const patchMethod = (methodName: "next" | "return" | "throw") => { + const originalMethod = Reflect.get(iterator, methodName); + if (typeof originalMethod !== "function") { + return; + } + + Reflect.set(iterator, methodName, (...args: unknown[]) => + runWithClaudeLocalToolContext( + () => + Reflect.apply( + originalMethod as (...methodArgs: unknown[]) => unknown, + iterator, + args, + ), + localToolContext, + ), + ); + }; + + patchMethod("next"); + patchMethod("return"); + patchMethod("throw"); + return iterator; + }, localToolContext); + }; + + Object.defineProperty( + patchedAsyncIterator, + LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED, + { + configurable: false, + enumerable: false, + value: true, + writable: false, + }, + ); + Reflect.set(stream, Symbol.asyncIterator, patchedAsyncIterator); + return result; +} diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts index b84a0ddac..24f7cd9d2 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts @@ -1,5 +1,6 @@ import { startSpan, withCurrent } from "../../logger"; import { SpanTypeAttribute } from "../../../util/index"; +import { getClaudeLocalToolParentResolver } from "./claude-agent-sdk-local-tool-context"; export type LocalToolSpanMetadata = { serverName?: string; @@ -25,6 +26,20 @@ function isPromiseLike(value: unknown): value is Promise { ); } +function getToolUseIdFromExtra(extra: unknown): string | undefined { + if (!extra || typeof extra !== "object" || !("_meta" in extra)) { + return undefined; + } + + const meta = Reflect.get(extra, "_meta"); + if (!meta || typeof meta !== "object") { + return undefined; + } + + const toolUseId = Reflect.get(meta, "claudecode/toolUseId"); + return typeof toolUseId === "string" ? toolUseId : undefined; +} + export function wrapLocalClaudeToolHandler( handler: LocalToolHandler, getMetadata: () => LocalToolSpanMetadata, @@ -42,46 +57,62 @@ export function wrapLocalClaudeToolHandler( ...handlerArgs: unknown[] ) { const metadata = getMetadata(); + const rawToolName = metadata.serverName + ? `mcp__${metadata.serverName}__${metadata.toolName}` + : metadata.toolName; + const toolUseId = getToolUseIdFromExtra(handlerArgs[1]); + const localToolParentResolver = getClaudeLocalToolParentResolver(); const spanName = metadata.serverName ? `tool: ${metadata.serverName}/${metadata.toolName}` : `tool: ${metadata.toolName}`; - const span = startSpan({ - event: { - input: handlerArgs[0], - metadata: { - "gen_ai.tool.name": metadata.toolName, - ...(metadata.serverName && { - "mcp.server": metadata.serverName, - }), + const runWithResolvedParent = async () => { + const parent = + toolUseId && localToolParentResolver + ? await localToolParentResolver(toolUseId).catch(() => undefined) + : undefined; + const span = startSpan({ + event: { + input: handlerArgs[0], + metadata: { + "claude_agent_sdk.raw_tool_name": rawToolName, + "gen_ai.tool.name": metadata.toolName, + ...(toolUseId && { "gen_ai.tool.call.id": toolUseId }), + ...(metadata.serverName && { + "mcp.server": metadata.serverName, + }), + }, }, - }, - name: spanName, - spanAttributes: { type: SpanTypeAttribute.TOOL }, - }); - - const runHandler = () => Reflect.apply(handler, this, handlerArgs); - const finalizeSuccess = (result: unknown) => { - span.log({ output: result }); - span.end(); - return result; - }; - const finalizeError = (error: unknown) => { - span.log({ error: toErrorMessage(error) }); - span.end(); - throw error; - }; + name: spanName, + ...(parent && { parent }), + spanAttributes: { type: SpanTypeAttribute.TOOL }, + }); + + const runHandler = () => Reflect.apply(handler, this, handlerArgs); + const finalizeSuccess = (result: unknown) => { + span.log({ output: result }); + span.end(); + return result; + }; + const finalizeError = (error: unknown) => { + span.log({ error: toErrorMessage(error) }); + span.end(); + throw error; + }; - return withCurrent(span, () => { - try { - const result = runHandler(); - if (isPromiseLike(result)) { - return result.then(finalizeSuccess, finalizeError); + return withCurrent(span, () => { + try { + const result = runHandler(); + if (isPromiseLike(result)) { + return result.then(finalizeSuccess, finalizeError); + } + return finalizeSuccess(result); + } catch (error) { + return finalizeError(error); } - return finalizeSuccess(result); - } catch (error) { - return finalizeError(error); - } - }); + }); + }; + + return runWithResolvedParent(); }; Object.defineProperty(wrappedHandler, LOCAL_TOOL_HANDLER_WRAPPED, { diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index 5b2b17860..094bbfbb5 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -16,6 +16,12 @@ import { collectLocalMcpServerToolHookNames, wrapLocalMcpServerToolHandlers, } from "./claude-agent-sdk-local-tool-spans"; +import { + bindClaudeLocalToolContextToAsyncIterable, + createClaudeLocalToolContext, + setClaudeLocalToolParentResolver, + type ClaudeAgentSDKLocalToolContext, +} from "./claude-agent-sdk-local-tool-context"; import type { ClaudeAgentSDKHookCallback, ClaudeAgentSDKHookCallbackMatcher, @@ -564,6 +570,7 @@ type QueryState = { latestLlmParentBySubAgentToolUse: Map; latestRootLlmParentRef: { value: string | undefined }; toolUseToParent: Map; + localToolContext: ClaudeAgentSDKLocalToolContext; }; async function finalizeCurrentMessageGroup(state: QueryState): Promise { @@ -896,55 +903,36 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { value: undefined as string | undefined, }; const pendingSubAgentNames = new Map(); + const localToolContext = createClaudeLocalToolContext(); const { hasLocalToolHandlers, localToolHookNames } = prepareLocalToolHandlersInMcpServers(options.mcpServers); const skipLocalToolHooks = options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true || hasLocalToolHandlers; - const optionsWithHooks = injectTracingHooks( - options, - async (toolUseID) => { - const parentToolUseId = toolUseToParent.get(toolUseID); - const parentKey = llmParentKey(parentToolUseId ?? null); - const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); - if (activeLlmSpan) { - return activeLlmSpan.export(); + const resolveToolUseParentSpan: ParentSpanResolver = async ( + toolUseID, + ) => { + const parentToolUseId = toolUseToParent.get(toolUseID); + const parentKey = llmParentKey(parentToolUseId ?? null); + const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); + if (activeLlmSpan) { + return activeLlmSpan.export(); + } + + if (parentToolUseId) { + const parentLlm = + latestLlmParentBySubAgentToolUse.get(parentToolUseId); + if (parentLlm) { + return parentLlm; } - if (parentToolUseId) { - const parentLlm = - latestLlmParentBySubAgentToolUse.get(parentToolUseId); - if (parentLlm) { - return parentLlm; - } - - const subAgentSpan = await ensureSubAgentSpan( - pendingSubAgentNames, - span, - subAgentSpans, - parentToolUseId, - ); - const parentSpan = await subAgentSpan.export(); - const llmSpan = startSpan({ - name: "anthropic.messages.create", - parent: parentSpan, - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - startTime: getCurrentUnixTimestamp(), - }); - activeLlmSpansByParentToolUse.set(parentKey, llmSpan); - const llmSpanExport = await llmSpan.export(); - latestLlmParentBySubAgentToolUse.set( - parentToolUseId, - llmSpanExport, - ); - return llmSpanExport; - } - if (latestRootLlmParentRef.value) { - return latestRootLlmParentRef.value; - } - const parentSpan = await span.export(); + const subAgentSpan = await ensureSubAgentSpan( + pendingSubAgentNames, + span, + subAgentSpans, + parentToolUseId, + ); + const parentSpan = await subAgentSpan.export(); const llmSpan = startSpan({ name: "anthropic.messages.create", parent: parentSpan, @@ -955,9 +943,37 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { }); activeLlmSpansByParentToolUse.set(parentKey, llmSpan); const llmSpanExport = await llmSpan.export(); - latestRootLlmParentRef.value = llmSpanExport; + latestLlmParentBySubAgentToolUse.set( + parentToolUseId, + llmSpanExport, + ); return llmSpanExport; - }, + } + + if (latestRootLlmParentRef.value) { + return latestRootLlmParentRef.value; + } + + const parentSpan = await span.export(); + const llmSpan = startSpan({ + name: "anthropic.messages.create", + parent: parentSpan, + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + startTime: getCurrentUnixTimestamp(), + }); + activeLlmSpansByParentToolUse.set(parentKey, llmSpan); + const llmSpanExport = await llmSpan.export(); + latestRootLlmParentRef.value = llmSpanExport; + return llmSpanExport; + }; + + localToolContext.resolveLocalToolParent = resolveToolUseParentSpan; + setClaudeLocalToolParentResolver(resolveToolUseParentSpan); + const optionsWithHooks = injectTracingHooks( + options, + resolveToolUseParentSpan, activeToolSpans, localToolHookNames, skipLocalToolHooks, @@ -989,6 +1005,7 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { latestLlmParentBySubAgentToolUse, latestRootLlmParentRef, toolUseToParent, + localToolContext, }); }, @@ -998,7 +1015,10 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { return; } - const eventResult = event.result; + const eventResult = bindClaudeLocalToolContextToAsyncIterable( + event.result, + state.localToolContext, + ); if (eventResult === undefined) { state.span.end(); spans.delete(event); diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.test.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.test.ts index 9ddeb7036..c28db748e 100644 --- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.test.ts +++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.test.ts @@ -12,8 +12,6 @@ import { initLogger, _exportsForTestingOnly } from "../../logger"; import { configureNode } from "../../node/config"; import { z } from "zod/v3"; -debugger; - const makePromptMessage = (content: string) => ({ type: "user", message: { role: "user", content }, @@ -761,10 +759,14 @@ describe.skipIf(!claudeSDK)("claude-agent-sdk integration tests", () => { ); }); - // Verify span hierarchy (all children should reference the root task span) + // Verify span hierarchy const rootSpanId = taskSpan!.span_id; spans - .filter((s) => s.span_id !== rootSpanId) + .filter( + (s) => + s.span_id !== rootSpanId && + (s["span_attributes"] as Record).type !== "tool", + ) .forEach((span) => { expect(span.root_span_id).toBe(rootSpanId); expect(span.span_parents).toContain(rootSpanId); @@ -933,15 +935,53 @@ describe.skipIf(!claudeSDK)("claude-agent-sdk integration tests", () => { ); expect(subAgentLlmSpans.length).toBeGreaterThanOrEqual(1); - // Tool spans within the sub-agent should be parented under the sub-agent, not root + const spanById = new Map( + spans.map((span) => [span.span_id, span] as const), + ); + const isDescendantOf = ( + span: (typeof spans)[number], + ancestorId: string, + ): boolean => { + const queue = [...((span.span_parents as string[] | undefined) ?? [])]; + const visited = new Set(); + + while (queue.length > 0) { + const parentId = queue.shift(); + if (!parentId || visited.has(parentId)) { + continue; + } + if (parentId === ancestorId) { + return true; + } + visited.add(parentId); + const parentSpan = spanById.get(parentId); + if (!parentSpan) { + continue; + } + const parentAncestors = + (parentSpan.span_parents as string[] | undefined) ?? []; + queue.push(...parentAncestors); + } + + return false; + }; + + // Local tool spans should be nested under the sub-agent. const subAgentToolSpans = spans.filter( (s) => (s["span_attributes"] as Record).type === "tool" && - (s.span_parents as string[])?.includes(subAgentSpan!.span_id as string), + (s["span_attributes"] as Record).name === + "tool: calculator/calculator" && + isDescendantOf(s, subAgentSpan!.span_id as string), ); expect(subAgentToolSpans.length).toBeGreaterThanOrEqual(1); subAgentToolSpans.forEach((toolSpan) => { - expect(toolSpan.span_parents).not.toContain(rootSpan!.span_id); + const metadata = toolSpan.metadata as Record; + expect(metadata["gen_ai.tool.name"]).toBe("calculator"); + expect(metadata["mcp.server"]).toBe("calculator"); + expect(metadata["claude_agent_sdk.raw_tool_name"]).toBe( + "mcp__calculator__calculator", + ); }); }, 60000); }); diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts index 8ba922048..db24f0bcf 100644 --- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts +++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts @@ -55,7 +55,6 @@ function wrapClaudeAgentQuery( ? (defaultThis ?? thisArg) : thisArg; return claudeAgentSDKChannels.query.traceSync( - // Async iterator shenanigans are handled purely in the plugin which consumes this channel emission. () => Reflect.apply(target, invocationTarget, [wrappedParams]), // The channel carries no extra context fields, but the generated // StartOf<> type for Record is overly strict here. From 645270f1e5b19a5d067abf97c7324cacc9ea46d0 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Wed, 8 Apr 2026 12:08:53 +0200 Subject: [PATCH 3/4] fix ci? --- .../claude-agent-sdk-v0.1.span-events.json | 34 +++++++++---------- .../claude-agent-sdk-v0.2.76.span-events.json | 34 +++++++++---------- .../claude-agent-sdk-v0.2.79.span-events.json | 34 +++++++++---------- .../claude-agent-sdk-v0.2.81.span-events.json | 34 +++++++++---------- .../claude-agent-sdk-local-tool-context.ts | 4 +-- 5 files changed, 70 insertions(+), 70 deletions(-) diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json index f980b9320..a7132668d 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.1.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -145,9 +145,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -160,7 +160,7 @@ "metric_keys": [], "name": "claude-agent-failure-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -175,9 +175,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -192,9 +192,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -229,9 +229,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -242,9 +242,9 @@ "metric_keys": [], "name": "Agent: math-expert", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -257,7 +257,7 @@ "metric_keys": [], "name": "claude-agent-subagent-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -272,9 +272,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -288,9 +288,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json index 850acdbca..7a702f0a9 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.76.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -145,9 +145,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -160,7 +160,7 @@ "metric_keys": [], "name": "claude-agent-failure-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -175,9 +175,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -192,9 +192,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -229,9 +229,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -242,9 +242,9 @@ "metric_keys": [], "name": "Agent: sub-agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -257,7 +257,7 @@ "metric_keys": [], "name": "claude-agent-subagent-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -272,9 +272,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -288,9 +288,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json index 850acdbca..7a702f0a9 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.79.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -145,9 +145,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -160,7 +160,7 @@ "metric_keys": [], "name": "claude-agent-failure-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -175,9 +175,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -192,9 +192,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -229,9 +229,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -242,9 +242,9 @@ "metric_keys": [], "name": "Agent: sub-agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -257,7 +257,7 @@ "metric_keys": [], "name": "claude-agent-subagent-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -272,9 +272,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -288,9 +288,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } diff --git a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json index 850acdbca..7a702f0a9 100644 --- a/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json +++ b/e2e/scenarios/claude-agent-sdk-instrumentation/__snapshots__/claude-agent-sdk-v0.2.81.span-events.json @@ -123,7 +123,7 @@ "root_span_id": "", "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -145,9 +145,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -160,7 +160,7 @@ "metric_keys": [], "name": "claude-agent-failure-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -175,9 +175,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -192,9 +192,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } @@ -229,9 +229,9 @@ ], "name": "anthropic.messages.create", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "llm" }, @@ -242,9 +242,9 @@ "metric_keys": [], "name": "Agent: sub-agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -257,7 +257,7 @@ "metric_keys": [], "name": "claude-agent-subagent-operation", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ "" ], @@ -272,9 +272,9 @@ "metric_keys": [], "name": "Claude Agent", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "task" }, @@ -288,9 +288,9 @@ "metric_keys": [], "name": "tool: calculator/calculator", "root_span_id": "", - "span_id": "", + "span_id": "", "span_parents": [ - "" + "" ], "type": "tool" } diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts index d5fc71b74..ef96d0efe 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts @@ -54,7 +54,7 @@ export function createClaudeLocalToolContext(): ClaudeAgentSDKLocalToolContext { return {}; } -export function runWithClaudeLocalToolContext( +function runWithClaudeLocalToolContext( callback: () => R, context?: ClaudeAgentSDKLocalToolContext, ): R { @@ -64,7 +64,7 @@ export function runWithClaudeLocalToolContext( ); } -export function ensureClaudeLocalToolContext(): +function ensureClaudeLocalToolContext(): | ClaudeAgentSDKLocalToolContext | undefined { const existing = localToolContextStore.getStore(); From ffd5bb518ce5bb0630f99aa0dcb8d70c65be5266 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Wed, 8 Apr 2026 15:32:22 +0200 Subject: [PATCH 4/4] fix span time --- .../plugins/claude-agent-sdk-plugin.ts | 60 +++++++++---------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts index 094bbfbb5..bf92a77c4 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts @@ -729,6 +729,31 @@ async function handleStreamMessage( } if (message.type === "assistant" && message.message?.usage) { + const parentToolUseId = message.parent_tool_use_id ?? null; + const parentKey = llmParentKey(parentToolUseId); + if (!state.activeLlmSpansByParentToolUse.has(parentKey)) { + let llmParentSpan = await state.span.export(); + if (parentToolUseId) { + const subAgentSpan = await ensureSubAgentSpan( + state.pendingSubAgentNames, + state.span, + state.subAgentSpans, + parentToolUseId, + ); + llmParentSpan = await subAgentSpan.export(); + } + + const llmSpan = startSpan({ + name: "anthropic.messages.create", + parent: llmParentSpan, + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + startTime: state.currentMessageStartTime, + }); + state.activeLlmSpansByParentToolUse.set(parentKey, llmSpan); + } + state.currentMessages.push(message); } @@ -912,8 +937,8 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { const resolveToolUseParentSpan: ParentSpanResolver = async ( toolUseID, ) => { - const parentToolUseId = toolUseToParent.get(toolUseID); - const parentKey = llmParentKey(parentToolUseId ?? null); + const parentToolUseId = toolUseToParent.get(toolUseID) ?? null; + const parentKey = llmParentKey(parentToolUseId); const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey); if (activeLlmSpan) { return activeLlmSpan.export(); @@ -932,41 +957,14 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { subAgentSpans, parentToolUseId, ); - const parentSpan = await subAgentSpan.export(); - const llmSpan = startSpan({ - name: "anthropic.messages.create", - parent: parentSpan, - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - startTime: getCurrentUnixTimestamp(), - }); - activeLlmSpansByParentToolUse.set(parentKey, llmSpan); - const llmSpanExport = await llmSpan.export(); - latestLlmParentBySubAgentToolUse.set( - parentToolUseId, - llmSpanExport, - ); - return llmSpanExport; + return subAgentSpan.export(); } if (latestRootLlmParentRef.value) { return latestRootLlmParentRef.value; } - const parentSpan = await span.export(); - const llmSpan = startSpan({ - name: "anthropic.messages.create", - parent: parentSpan, - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - startTime: getCurrentUnixTimestamp(), - }); - activeLlmSpansByParentToolUse.set(parentKey, llmSpan); - const llmSpanExport = await llmSpan.export(); - latestRootLlmParentRef.value = llmSpanExport; - return llmSpanExport; + return span.export(); }; localToolContext.resolveLocalToolParent = resolveToolUseParentSpan;