Skip to content

Commit 045a3e5

Browse files
lforstlforst
committed
fix(claude-agent-sdk): Nest built-in tools under sub-agents
1 parent 3121158 commit 045a3e5

5 files changed

Lines changed: 172 additions & 12 deletions

File tree

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

Lines changed: 115 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,38 @@ function findSpanById(
169169
return events.find((event) => event.span.id === spanId);
170170
}
171171

172+
function isDescendantOf(
173+
events: CapturedLogEvent[],
174+
event: CapturedLogEvent | undefined,
175+
ancestorId: string | undefined,
176+
): boolean {
177+
if (!event || !ancestorId) {
178+
return false;
179+
}
180+
181+
const spanById = new Map(events.map((span) => [span.span.id, span] as const));
182+
const queue = [...event.span.parentIds];
183+
const visited = new Set<string>();
184+
185+
while (queue.length > 0) {
186+
const parentId = queue.shift();
187+
if (!parentId || visited.has(parentId)) {
188+
continue;
189+
}
190+
if (parentId === ancestorId) {
191+
return true;
192+
}
193+
194+
visited.add(parentId);
195+
const parentSpan = spanById.get(parentId);
196+
if (parentSpan) {
197+
queue.push(...parentSpan.span.parentIds);
198+
}
199+
}
200+
201+
return false;
202+
}
203+
172204
function hasSubAgentHandoffToolName(
173205
event: CapturedLogEvent | undefined,
174206
): boolean {
@@ -185,10 +217,13 @@ function hasSubAgentHandoffToolName(
185217

186218
function findSubAgentTaskSpan(
187219
events: CapturedLogEvent[],
220+
ancestorId?: string,
188221
): CapturedLogEvent | undefined {
189222
return events.find(
190223
(event) =>
191-
event.span.type === "task" && event.span.name?.startsWith("Agent:"),
224+
event.span.type === "task" &&
225+
event.span.name?.startsWith("Agent:") &&
226+
(!ancestorId || isDescendantOf(events, event, ancestorId)),
192227
);
193228
}
194229

@@ -202,6 +237,14 @@ function findSubAgentHandoffTool(
202237
return hasSubAgentHandoffToolName(parentSpan) ? parentSpan : undefined;
203238
}
204239

240+
function findOperationTaskRoot(
241+
events: CapturedLogEvent[],
242+
operationName: string,
243+
): CapturedLogEvent | undefined {
244+
const operation = findLatestSpan(events, operationName);
245+
return findChildSpans(events, "Claude Agent", operation?.span.id).at(-1);
246+
}
247+
205248
function buildSpanSummary(events: CapturedLogEvent[]): Json {
206249
const root = findLatestSpan(events, ROOT_NAME);
207250
const basicOperation = findLatestSpan(events, "claude-agent-basic-operation");
@@ -252,7 +295,7 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json {
252295
const input = event.input as Array<{ content?: string }> | undefined;
253296
return Array.isArray(input) && input.some((item) => item.content);
254297
});
255-
const subAgentTask = findSubAgentTaskSpan(events);
298+
const subAgentTask = findSubAgentTaskSpan(events, subAgentTaskRoot?.span.id);
256299
const subAgentLlm = findChildSpans(
257300
events,
258301
"anthropic.messages.create",
@@ -268,9 +311,16 @@ function buildSpanSummary(events: CapturedLogEvent[]): Json {
268311
const basicTool =
269312
findToolSpanByLocalHandler(events, "calculator-local-handler-multiply") ??
270313
findToolSpanByOperation(events, "multiply");
271-
const subAgentTool =
314+
const subAgentToolCandidate =
272315
findToolSpanByLocalHandler(events, "calculator-local-handler-add") ??
273316
findToolSpanByOperation(events, "add");
317+
const subAgentTool = isDescendantOf(
318+
events,
319+
subAgentToolCandidate,
320+
subAgentTask?.span.id,
321+
)
322+
? subAgentToolCandidate
323+
: undefined;
274324
const failureTool =
275325
findToolSpanByLocalHandler(events, "calculator-local-handler-divide") ??
276326
findToolSpanByOperation(events, "divide");
@@ -447,15 +497,14 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: {
447497
events,
448498
"claude-agent-subagent-operation",
449499
);
450-
const taskRoot = findChildSpans(
500+
const taskRoot = findOperationTaskRoot(
451501
events,
452-
"Claude Agent",
453-
operation?.span.id,
454-
).at(-1);
502+
"claude-agent-subagent-operation",
503+
);
455504
const llm = findAllSpans(events, "anthropic.messages.create").find(
456505
(event) => event.span.parentIds.includes(taskRoot?.span.id ?? ""),
457506
);
458-
const nestedTask = findSubAgentTaskSpan(events);
507+
const nestedTask = findSubAgentTaskSpan(events, taskRoot?.span.id);
459508
const handoffTool = findSubAgentHandoffTool(events, nestedTask);
460509
const nestedTaskLlm = findChildSpans(
461510
events,
@@ -502,6 +551,64 @@ export function defineClaudeAgentSDKInstrumentationAssertions(options: {
502551
}
503552
});
504553

554+
if (options.expectTaskLifecycleDetails) {
555+
test(
556+
"nests built-in subagent Bash under the subagent llm",
557+
testConfig,
558+
() => {
559+
const operation = findLatestSpan(
560+
events,
561+
"claude-agent-subagent-built-in-tool-operation",
562+
);
563+
const taskRoot = findOperationTaskRoot(
564+
events,
565+
"claude-agent-subagent-built-in-tool-operation",
566+
);
567+
const taskRootLlm = findAllSpans(
568+
events,
569+
"anthropic.messages.create",
570+
).find((event) =>
571+
event.span.parentIds.includes(taskRoot?.span.id ?? ""),
572+
);
573+
const nestedTask = findSubAgentTaskSpan(events, taskRoot?.span.id);
574+
const handoffTool = findSubAgentHandoffTool(events, nestedTask);
575+
const nestedTaskLlm = findChildSpans(
576+
events,
577+
"anthropic.messages.create",
578+
nestedTask?.span.id,
579+
).at(-1);
580+
const bashTool = findAllSpans(events, "tool: Bash").find((event) =>
581+
isDescendantOf(events, event, taskRoot?.span.id),
582+
);
583+
const bashToolParent = findSpanById(
584+
events,
585+
bashTool?.span.parentIds[0],
586+
);
587+
588+
expect(operation).toBeDefined();
589+
expect(taskRoot).toBeDefined();
590+
expect(nestedTask).toBeDefined();
591+
expect(handoffTool).toBeDefined();
592+
expect(nestedTaskLlm).toBeDefined();
593+
expect(bashTool).toBeDefined();
594+
expect(isDescendantOf(events, bashTool, nestedTask?.span.id)).toBe(
595+
true,
596+
);
597+
expect(bashTool?.span.parentIds).not.toContain(
598+
taskRootLlm?.span.id ?? "",
599+
);
600+
if (bashToolParent?.span.type === "llm") {
601+
expect(bashToolParent.span.parentIds).toContain(
602+
nestedTask?.span.id ?? "",
603+
);
604+
}
605+
if (bashToolParent?.span.type === "task") {
606+
expect(bashToolParent.span.id).toBe(nestedTask?.span.id);
607+
}
608+
},
609+
);
610+
}
611+
505612
test("captures tool failure details", testConfig, () => {
506613
const operation = findLatestSpan(
507614
events,

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
import { z } from "zod";
88

99
const CLAUDE_AGENT_MODEL = "claude-haiku-4-5-20251001";
10+
const CLAUDE_AGENT_TOP_LEVEL_MODEL = "claude-sonnet-4-5-20250929";
1011

1112
export const ROOT_NAME = "claude-agent-sdk-root";
1213
export const SCENARIO_NAME = "claude-agent-sdk-traces";
@@ -139,6 +140,33 @@ async function runClaudeAgentSDKScenario({ decorateSDK, sdk }) {
139140
},
140141
);
141142

143+
await runOperation(
144+
"claude-agent-subagent-built-in-tool-operation",
145+
"subagent-built-in-tool",
146+
async () => {
147+
await collectAsync(
148+
query({
149+
prompt:
150+
'You MUST call the Agent tool now with subagent_type="echo" and description "echo greeting". Do not call Bash yourself. Do not answer with text yourself; delegate to the echo sub-agent.',
151+
options: {
152+
agents: {
153+
echo: {
154+
description: "Runs one bash echo and reports back.",
155+
model: CLAUDE_AGENT_MODEL,
156+
prompt:
157+
"Run `echo hello` via Bash exactly once, then reply with only the word done.",
158+
tools: ["Bash"],
159+
},
160+
},
161+
allowedTools: ["Agent", "Bash"],
162+
model: CLAUDE_AGENT_TOP_LEVEL_MODEL,
163+
permissionMode: "bypassPermissions",
164+
},
165+
}),
166+
);
167+
},
168+
);
169+
142170
await runOperation(
143171
"claude-agent-failure-operation",
144172
"failure",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import { defineClaudeAgentSDKInstrumentationAssertions } from "./assertions";
99
const scenarioDir = await prepareScenarioDir({
1010
scenarioDir: resolveScenarioDir(import.meta.url),
1111
});
12-
const TIMEOUT_MS = 120_000;
12+
const TIMEOUT_MS = 180_000;
1313
const claudeAgentSDKScenarios = await Promise.all(
1414
[
1515
{

‎js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,10 @@ type ParsedToolName = {
3838
rawToolName: string;
3939
toolName: string;
4040
};
41-
type ParentSpanResolver = (toolUseID: string) => Promise<string>;
41+
type ParentSpanResolver = (
42+
toolUseID: string,
43+
context?: { agentId?: string },
44+
) => Promise<string>;
4245
type LLMSpanResult = {
4346
finalMessage: ClaudeConversationMessage | undefined;
4447
spanExport: string;
@@ -188,6 +191,18 @@ function resolveTaskToolUseId(
188191
return undefined;
189192
}
190193

194+
function seedTaskToolUseIdMapping(
195+
taskIdToToolUseId: Map<string, string>,
196+
message: ClaudeAgentSDKMessage,
197+
): void {
198+
if (
199+
typeof message.task_id === "string" &&
200+
typeof message.tool_use_id === "string"
201+
) {
202+
taskIdToToolUseId.set(message.task_id, message.tool_use_id);
203+
}
204+
}
205+
191206
function extractUsageFromMessage(
192207
message: ClaudeAgentSDKMessage,
193208
): Record<string, number> {
@@ -479,7 +494,9 @@ function createToolTracingHooks(
479494
},
480495
},
481496
name: parsed.displayName,
482-
parent: await resolveParentSpan(toolUseID),
497+
parent: await resolveParentSpan(toolUseID, {
498+
agentId: input.agent_id,
499+
}),
483500
spanAttributes: { type: SpanTypeAttribute.TOOL },
484501
});
485502

@@ -869,6 +886,8 @@ function maybeTrackToolUseContext(
869886
state: QueryState,
870887
message: ClaudeAgentSDKMessage,
871888
): void {
889+
seedTaskToolUseIdMapping(state.taskIdToToolUseId, message);
890+
872891
if (
873892
message.type !== "assistant" ||
874893
!Array.isArray(message.message?.content)
@@ -1311,8 +1330,13 @@ export class ClaudeAgentSDKPlugin extends BasePlugin {
13111330
hasLocalToolHandlers;
13121331
const resolveToolUseParentSpan: ParentSpanResolver = async (
13131332
toolUseID,
1333+
context,
13141334
) => {
1315-
const parentToolUseId = toolUseToParent.get(toolUseID) ?? null;
1335+
const parentToolUseId =
1336+
toolUseToParent.get(toolUseID) ??
1337+
(context?.agentId
1338+
? (taskIdToToolUseId.get(context.agentId) ?? null)
1339+
: null);
13161340
const parentKey = llmParentKey(parentToolUseId);
13171341
const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey);
13181342
if (activeLlmSpan) {

‎js/src/vendor-sdk-types/claude-agent-sdk.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ interface BaseHookInput {
2323
session_id: string;
2424
transcript_path: string;
2525
cwd: string;
26+
agent_id?: string;
2627
permission_mode?: string;
2728
}
2829

0 commit comments

Comments
 (0)