Skip to content

Commit 5dd279a

Browse files
committed
Use parent's tools in subagents. Leave note on what tools are actually available. Fixes prompt caching!
1 parent 731bc2e commit 5dd279a

11 files changed

+139
-14
lines changed

packages/agent-runtime/src/__tests__/cost-aggregation.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ describe('Cost Aggregation System', () => {
9090
repoUrl: undefined,
9191
signal: new AbortController().signal,
9292
system: 'Test system prompt',
93+
tools: {},
9394
userId: 'test-user',
9495
userInputId: 'test-input',
9596
writeToClient: () => {},

packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,83 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
407407
// allowing the LLM provider to cache and reuse the system prompt
408408
})
409409

410+
it('should pass parent tools and add subagent tools message when inheritParentSystemPrompt is true', async () => {
411+
const sessionState = getInitialSessionState(mockFileContext)
412+
413+
// Create a child that inherits system prompt and has specific tools
414+
const childWithTools: AgentTemplate = {
415+
id: 'child-with-tools',
416+
displayName: 'Child With Tools',
417+
outputMode: 'last_message',
418+
inputSchema: {},
419+
spawnerPrompt: '',
420+
model: 'anthropic/claude-sonnet-4',
421+
includeMessageHistory: false,
422+
inheritParentSystemPrompt: true,
423+
mcpServers: {},
424+
toolNames: ['read_files', 'code_search'],
425+
spawnableAgents: [],
426+
systemPrompt: '',
427+
instructionsPrompt: '',
428+
stepPrompt: '',
429+
}
430+
431+
mockLocalAgentTemplates['child-with-tools'] = childWithTools
432+
433+
// Run parent agent first
434+
await loopAgentSteps({
435+
...loopAgentStepsBaseParams,
436+
userInputId: 'test-parent',
437+
prompt: 'Parent task',
438+
agentType: 'parent',
439+
agentState: sessionState.mainAgentState,
440+
})
441+
442+
const parentMessages = capturedMessages
443+
const parentSystemPrompt = (parentMessages[0].content[0] as TextPart).text
444+
445+
// Mock parent tools
446+
const parentTools = { read_files: {}, write_file: {}, code_search: {} }
447+
448+
// Run child agent with inheritParentSystemPrompt=true and parentTools
449+
capturedMessages = []
450+
const childAgentState = {
451+
...sessionState.mainAgentState,
452+
agentId: 'child-agent',
453+
agentType: 'child-with-tools' as const,
454+
messageHistory: [],
455+
}
456+
457+
await loopAgentSteps({
458+
...loopAgentStepsBaseParams,
459+
userInputId: 'test-child',
460+
prompt: 'Child task',
461+
agentType: 'child-with-tools',
462+
agentState: childAgentState,
463+
parentSystemPrompt: parentSystemPrompt,
464+
parentTools: parentTools as any,
465+
})
466+
467+
const childMessages = capturedMessages
468+
469+
// Verify child uses parent's system prompt
470+
expect(childMessages[0].role).toBe('system')
471+
expect((childMessages[0].content[0] as TextPart).text).toBe(
472+
parentSystemPrompt,
473+
)
474+
475+
// Verify there's a message about subagent tools
476+
const subagentToolsMessage = childMessages.find(
477+
(msg) =>
478+
msg.role === 'user' &&
479+
msg.content[0].type === 'text' &&
480+
msg.content[0].text.includes('subagent') &&
481+
msg.content[0].text.includes('read_files') &&
482+
msg.content[0].text.includes('code_search'),
483+
)
484+
expect(subagentToolsMessage).toBeTruthy()
485+
})
486+
410487
it('should support both inheritParentSystemPrompt and includeMessageHistory together', async () => {
411488
const sessionState = getInitialSessionState(mockFileContext)
412489

packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ describe('Spawn Agents Message History', () => {
6868
sendSubagentChunk: mockSendSubagentChunk,
6969
signal: new AbortController().signal,
7070
system: 'Test system prompt',
71+
tools: {},
7172
userId: TEST_USER_ID,
7273
userInputId: 'test-input',
7374
writeToClient: () => {},

packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ describe('Spawn Agents Permissions', () => {
6767
sendSubagentChunk: mockSendSubagentChunk,
6868
signal: new AbortController().signal,
6969
system: 'Test system prompt',
70+
tools: {},
7071
userId: TEST_USER_ID,
7172
userInputId: 'test-input',
7273
writeToClient: () => {},

packages/agent-runtime/src/__tests__/subagent-streaming.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ describe('Subagent Streaming', () => {
7070
sendSubagentChunk: mockSendSubagentChunk,
7171
signal: new AbortController().signal,
7272
system: 'Test system prompt',
73+
tools: {},
7374
userId: TEST_USER_ID,
7475
userInputId: 'test-input',
7576
writeToClient: mockWriteToClient,

packages/agent-runtime/src/run-agent-step.ts

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ import type {
5858
CustomToolDefinitions,
5959
ProjectFileContext,
6060
} from '@codebuff/common/util/file'
61+
import type { ToolSet } from 'ai'
6162

6263
async function additionalToolDefinitions(
6364
params: {
@@ -481,6 +482,7 @@ export async function loopAgentSteps(
481482
localAgentTemplates: Record<string, AgentTemplate>
482483
logger: Logger
483484
parentSystemPrompt?: string
485+
parentTools?: ToolSet
484486
prompt: string | undefined
485487
signal: AbortSignal
486488
spawnParams: Record<string, any> | undefined
@@ -554,6 +556,7 @@ export async function loopAgentSteps(
554556
localAgentTemplates,
555557
logger,
556558
parentSystemPrompt,
559+
parentTools,
557560
prompt,
558561
signal,
559562
spawnParams,
@@ -631,18 +634,29 @@ export async function loopAgentSteps(
631634
},
632635
})) ?? ''
633636

634-
const tools = await getToolSet({
635-
toolNames: agentTemplate.toolNames,
636-
additionalToolDefinitions: async () => {
637-
if (!cachedAdditionalToolDefinitions) {
638-
cachedAdditionalToolDefinitions = await additionalToolDefinitions({
639-
...params,
640-
agentTemplate,
641-
})
642-
}
643-
return cachedAdditionalToolDefinitions
644-
},
645-
})
637+
// Use parent's tools for prompt caching when inheritParentSystemPrompt is true
638+
const useParentTools =
639+
agentTemplate.inheritParentSystemPrompt && parentTools !== undefined
640+
641+
const tools = useParentTools
642+
? parentTools
643+
: await getToolSet({
644+
toolNames: agentTemplate.toolNames,
645+
additionalToolDefinitions: async () => {
646+
if (!cachedAdditionalToolDefinitions) {
647+
cachedAdditionalToolDefinitions = await additionalToolDefinitions({
648+
...params,
649+
agentTemplate,
650+
})
651+
}
652+
return cachedAdditionalToolDefinitions
653+
},
654+
})
655+
656+
// Build a message explaining the subagent's tool access when using parent's tools
657+
const subagentToolsMessage = useParentTools
658+
? `You are a subagent that only has access to the following tools: ${agentTemplate.toolNames.join(', ')}. Do not attempt to use any other tools.`
659+
: undefined
646660

647661
const hasUserMessage = Boolean(
648662
prompt || (spawnParams && Object.keys(spawnParams).length > 0),
@@ -651,6 +665,14 @@ export async function loopAgentSteps(
651665
const initialMessages = buildArray<Message>(
652666
...agentState.messageHistory,
653667

668+
// Add subagent tools message before user prompt when using parent's tools for caching
669+
subagentToolsMessage &&
670+
userMessage({
671+
content: withSystemTags(subagentToolsMessage),
672+
tags: ['SUBAGENT_TOOLS'],
673+
keepDuringTruncation: true,
674+
}),
675+
654676
hasUserMessage && [
655677
{
656678
// Actual user message!
@@ -877,7 +899,10 @@ export async function loopAgentSteps(
877899
)
878900

879901
// Re-throw NetworkError and PaymentRequiredError to allow SDK retry wrapper to handle it
880-
if (error instanceof Error && (error.name === 'NetworkError' || error.name === 'PaymentRequiredError')) {
902+
if (
903+
error instanceof Error &&
904+
(error.name === 'NetworkError' || error.name === 'PaymentRequiredError')
905+
) {
881906
throw error
882907
}
883908

packages/agent-runtime/src/tools/handlers/handler-function-type.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import type { Logger } from '@codebuff/common/types/contracts/logger'
1818
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
1919
import type { AgentState, Subgoal } from '@codebuff/common/types/session-state'
2020
import type { ProjectFileContext } from '@codebuff/common/util/file'
21+
import type { ToolSet } from 'ai'
2122

2223
type PresentOrAbsent<K extends PropertyKey, V> =
2324
| { [P in K]: V }
@@ -49,6 +50,7 @@ export type CodebuffToolHandlerFunction<T extends ToolName = ToolName> = (
4950
sendSubagentChunk: SendSubagentChunkFn
5051
signal: AbortSignal
5152
system: string
53+
tools?: ToolSet
5254
trackEvent: TrackEventFn
5355
userId: string | undefined
5456
userInputId: string

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-inline.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import type { ParamsExcluding } from '@codebuff/common/types/function-params'
1717
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
1818
import type { AgentState } from '@codebuff/common/types/session-state'
1919
import type { ProjectFileContext } from '@codebuff/common/util/file'
20+
import type { ToolSet } from 'ai'
2021

2122
type ToolName = 'spawn_agent_inline'
2223
export const handleSpawnAgentInline = (async (
@@ -32,6 +33,7 @@ export const handleSpawnAgentInline = (async (
3233
localAgentTemplates: Record<string, AgentTemplate>
3334
logger: Logger
3435
system: string
36+
tools?: ToolSet
3537
userId: string | undefined
3638
userInputId: string
3739
writeToClient: (chunk: string | PrintModeEvent) => void
@@ -44,6 +46,7 @@ export const handleSpawnAgentInline = (async (
4446
| 'parentAgentState'
4547
| 'agentState'
4648
| 'parentSystemPrompt'
49+
| 'parentTools'
4750
| 'onResponseChunk'
4851
| 'clearUserPromptMessagesAfterResponse'
4952
| 'fingerprintId'
@@ -57,6 +60,7 @@ export const handleSpawnAgentInline = (async (
5760
agentTemplate: parentAgentTemplate,
5861
fingerprintId,
5962
system,
63+
tools: parentTools = {},
6064
userInputId,
6165
writeToClient,
6266
} = params
@@ -105,6 +109,9 @@ export const handleSpawnAgentInline = (async (
105109
agentState: childAgentState,
106110
fingerprintId,
107111
parentSystemPrompt: system,
112+
parentTools: agentTemplate.inheritParentSystemPrompt
113+
? parentTools
114+
: undefined,
108115
onResponseChunk: (chunk) => {
109116
// Inherits parent's onResponseChunk, except for context-pruner (TODO: add an option for it to be silent?)
110117
if (agentType !== 'context-pruner') {

packages/agent-runtime/src/tools/handlers/tool/spawn-agent-utils.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import type {
1212
ParamsExcluding,
1313
OptionalFields,
1414
} from '@codebuff/common/types/function-params'
15+
import type { ToolSet } from 'ai'
1516
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
1617
import type {
1718
AgentState,
@@ -231,6 +232,7 @@ export async function executeSubagent(
231232
{
232233
agentTemplate: AgentTemplate
233234
parentAgentState: AgentState
235+
parentTools?: ToolSet
234236
onResponseChunk: (chunk: string | PrintModeEvent) => void
235237
isOnlyChild?: boolean
236238
ancestorRunIds: string[]

packages/agent-runtime/src/tools/handlers/tool/spawn-agents.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import type { Logger } from '@codebuff/common/types/contracts/logger'
1818
import type { ParamsExcluding } from '@codebuff/common/types/function-params'
1919
import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
2020
import type { AgentState } from '@codebuff/common/types/session-state'
21+
import type { ToolSet } from 'ai'
2122

2223
export type SendSubagentChunk = (data: {
2324
userInputId: string
@@ -40,6 +41,7 @@ export const handleSpawnAgents = (async (
4041
localAgentTemplates: Record<string, AgentTemplate>
4142
logger: Logger
4243
system: string
44+
tools?: ToolSet
4345
userId: string | undefined
4446
userInputId: string
4547
sendSubagentChunk: SendSubagentChunk
@@ -59,6 +61,7 @@ export const handleSpawnAgents = (async (
5961
| 'fingerprintId'
6062
| 'isOnlyChild'
6163
| 'parentSystemPrompt'
64+
| 'parentTools'
6265
| 'onResponseChunk'
6366
>,
6467
): Promise<{ output: CodebuffToolOutput<ToolName> }> => {
@@ -70,6 +73,7 @@ export const handleSpawnAgents = (async (
7073
agentTemplate: parentAgentTemplate,
7174
fingerprintId,
7275
system: parentSystemPrompt,
76+
tools: parentTools = {},
7377
userInputId,
7478
sendSubagentChunk,
7579
writeToClient,
@@ -118,6 +122,9 @@ export const handleSpawnAgents = (async (
118122
fingerprintId,
119123
isOnlyChild: agents.length === 1,
120124
parentSystemPrompt,
125+
parentTools: agentTemplate.inheritParentSystemPrompt
126+
? parentTools
127+
: undefined,
121128
onResponseChunk: (chunk: string | PrintModeEvent) => {
122129
if (typeof chunk === 'string') {
123130
sendSubagentChunk({

0 commit comments

Comments
 (0)