Skip to content

Commit 2d35017

Browse files
abhipatel12abhipatel12
authored
refactor: implement DelegateToAgentTool with discriminated union (#14769)
1 parent d5e469c commit 2d35017

12 files changed

Lines changed: 545 additions & 46 deletions

File tree

‎docs/core/policy-engine.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,8 @@ The Gemini CLI ships with a set of default policies to provide a safe
239239
out-of-the-box experience.
240240

241241
- **Read-only tools** (like `read_file`, `glob`) are generally **allowed**.
242+
- **Agent delegation** (like `delegate_to_agent`) is **allowed** (sub-agent
243+
actions are checked individually).
242244
- **Write tools** (like `write_file`, `run_shell_command`) default to
243245
**`ask_user`**.
244246
- In **`yolo`** mode, a high-priority rule allows all tools.

‎packages/core/src/agents/delegate-to-agent-tool.test.ts‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { describe, it, expect, vi, beforeEach } from 'vitest';
8+
import { DelegateToAgentTool } from './delegate-to-agent-tool.js';
9+
import { AgentRegistry } from './registry.js';
10+
import type { Config } from '../config/config.js';
11+
import type { AgentDefinition } from './types.js';
12+
import { SubagentInvocation } from './invocation.js';
13+
import type { MessageBus } from '../confirmation-bus/message-bus.js';
14+
import { MessageBusType } from '../confirmation-bus/types.js';
15+
import { DELEGATE_TO_AGENT_TOOL_NAME } from '../tools/tool-names.js';
16+
17+
vi.mock('./invocation.js', () => ({
18+
SubagentInvocation: vi.fn().mockImplementation(() => ({
19+
execute: vi
20+
.fn()
21+
.mockResolvedValue({ content: [{ type: 'text', text: 'Success' }] }),
22+
})),
23+
}));
24+
25+
describe('DelegateToAgentTool', () => {
26+
let registry: AgentRegistry;
27+
let config: Config;
28+
let tool: DelegateToAgentTool;
29+
let messageBus: MessageBus;
30+
31+
const mockAgentDef: AgentDefinition = {
32+
name: 'test_agent',
33+
description: 'A test agent',
34+
promptConfig: {},
35+
modelConfig: { model: 'test-model', temp: 0, top_p: 0 },
36+
inputConfig: {
37+
inputs: {
38+
arg1: { type: 'string', description: 'Argument 1', required: true },
39+
arg2: { type: 'number', description: 'Argument 2', required: false },
40+
},
41+
},
42+
runConfig: { max_turns: 1, max_time_minutes: 1 },
43+
toolConfig: { tools: [] },
44+
};
45+
46+
beforeEach(() => {
47+
config = {
48+
getDebugMode: () => false,
49+
modelConfigService: {
50+
registerRuntimeModelConfig: vi.fn(),
51+
},
52+
} as unknown as Config;
53+
54+
registry = new AgentRegistry(config);
55+
// Manually register the mock agent (bypassing protected method for testing)
56+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
57+
(registry as any).agents.set(mockAgentDef.name, mockAgentDef);
58+
59+
messageBus = {
60+
publish: vi.fn(),
61+
subscribe: vi.fn(),
62+
unsubscribe: vi.fn(),
63+
} as unknown as MessageBus;
64+
65+
tool = new DelegateToAgentTool(registry, config, messageBus);
66+
});
67+
68+
it('should use dynamic description from registry', () => {
69+
// registry has mockAgentDef registered in beforeEach
70+
expect(tool.description).toContain(
71+
'Delegates a task to a specialized sub-agent',
72+
);
73+
expect(tool.description).toContain(
74+
`- **${mockAgentDef.name}**: ${mockAgentDef.description}`,
75+
);
76+
});
77+
78+
it('should validate agent_name exists in registry', async () => {
79+
// Zod validation happens at build time now (or rather, build validates the schema)
80+
// Since we use discriminated union, an invalid agent_name won't match any option.
81+
expect(() =>
82+
tool.build({
83+
agent_name: 'non_existent_agent',
84+
}),
85+
).toThrow();
86+
});
87+
88+
it('should validate correct arguments', async () => {
89+
const invocation = tool.build({
90+
agent_name: 'test_agent',
91+
arg1: 'valid',
92+
});
93+
94+
const result = await invocation.execute(new AbortController().signal);
95+
expect(result).toEqual({ content: [{ type: 'text', text: 'Success' }] });
96+
expect(SubagentInvocation).toHaveBeenCalledWith(
97+
{ arg1: 'valid' },
98+
mockAgentDef,
99+
config,
100+
messageBus,
101+
);
102+
});
103+
104+
it('should throw error for missing required argument', async () => {
105+
// Missing arg1 should fail Zod validation
106+
expect(() =>
107+
tool.build({
108+
agent_name: 'test_agent',
109+
arg2: 123,
110+
}),
111+
).toThrow();
112+
});
113+
114+
it('should throw error for invalid argument type', async () => {
115+
// arg1 should be string, passing number
116+
expect(() =>
117+
tool.build({
118+
agent_name: 'test_agent',
119+
arg1: 123,
120+
}),
121+
).toThrow();
122+
});
123+
124+
it('should allow optional arguments to be omitted', async () => {
125+
const invocation = tool.build({
126+
agent_name: 'test_agent',
127+
arg1: 'valid',
128+
// arg2 is optional
129+
});
130+
131+
await expect(
132+
invocation.execute(new AbortController().signal),
133+
).resolves.toBeDefined();
134+
});
135+
136+
it('should throw error if an agent has an input named "agent_name"', () => {
137+
const invalidAgentDef: AgentDefinition = {
138+
...mockAgentDef,
139+
name: 'invalid_agent',
140+
inputConfig: {
141+
inputs: {
142+
agent_name: {
143+
type: 'string',
144+
description: 'Conflict',
145+
required: true,
146+
},
147+
},
148+
},
149+
};
150+
151+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
152+
(registry as any).agents.set(invalidAgentDef.name, invalidAgentDef);
153+
154+
expect(() => new DelegateToAgentTool(registry, config)).toThrow(
155+
"Agent 'invalid_agent' cannot have an input parameter named 'agent_name' as it is a reserved parameter for delegation.",
156+
);
157+
});
158+
159+
it('should use correct tool name "delegate_to_agent" when requesting confirmation', async () => {
160+
const invocation = tool.build({
161+
agent_name: 'test_agent',
162+
arg1: 'valid',
163+
});
164+
165+
// Trigger confirmation check
166+
const p = invocation.shouldConfirmExecute(new AbortController().signal);
167+
void p;
168+
169+
expect(messageBus.publish).toHaveBeenCalledWith(
170+
expect.objectContaining({
171+
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
172+
toolCall: expect.objectContaining({
173+
name: DELEGATE_TO_AGENT_TOOL_NAME,
174+
}),
175+
}),
176+
);
177+
});
178+
});

‎packages/core/src/agents/delegate-to-agent-tool.ts‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import { z } from 'zod';
8+
import { zodToJsonSchema } from 'zod-to-json-schema';
9+
import {
10+
BaseDeclarativeTool,
11+
Kind,
12+
type ToolInvocation,
13+
type ToolResult,
14+
BaseToolInvocation,
15+
} from '../tools/tools.js';
16+
import type { AnsiOutput } from '../utils/terminalSerializer.js';
17+
import { DELEGATE_TO_AGENT_TOOL_NAME } from '../tools/tool-names.js';
18+
import type { AgentRegistry } from './registry.js';
19+
import type { Config } from '../config/config.js';
20+
import type { MessageBus } from '../confirmation-bus/message-bus.js';
21+
import { SubagentInvocation } from './invocation.js';
22+
import type { AgentInputs } from './types.js';
23+
24+
type DelegateParams = { agent_name: string } & Record<string, unknown>;
25+
26+
export class DelegateToAgentTool extends BaseDeclarativeTool<
27+
DelegateParams,
28+
ToolResult
29+
> {
30+
constructor(
31+
private readonly registry: AgentRegistry,
32+
private readonly config: Config,
33+
messageBus?: MessageBus,
34+
) {
35+
const definitions = registry.getAllDefinitions();
36+
37+
let schema: z.ZodTypeAny;
38+
39+
if (definitions.length === 0) {
40+
// Fallback if no agents are registered (mostly for testing/safety)
41+
schema = z.object({
42+
agent_name: z.string().describe('No agents are currently available.'),
43+
});
44+
} else {
45+
const agentSchemas = definitions.map((def) => {
46+
const inputShape: Record<string, z.ZodTypeAny> = {
47+
agent_name: z.literal(def.name).describe(def.description),
48+
};
49+
50+
for (const [key, inputDef] of Object.entries(def.inputConfig.inputs)) {
51+
if (key === 'agent_name') {
52+
throw new Error(
53+
`Agent '${def.name}' cannot have an input parameter named 'agent_name' as it is a reserved parameter for delegation.`,
54+
);
55+
}
56+
57+
let validator: z.ZodTypeAny;
58+
59+
// Map input types to Zod
60+
switch (inputDef.type) {
61+
case 'string':
62+
validator = z.string();
63+
break;
64+
case 'number':
65+
validator = z.number();
66+
break;
67+
case 'boolean':
68+
validator = z.boolean();
69+
break;
70+
case 'integer':
71+
validator = z.number().int();
72+
break;
73+
case 'string[]':
74+
validator = z.array(z.string());
75+
break;
76+
case 'number[]':
77+
validator = z.array(z.number());
78+
break;
79+
default: {
80+
// This provides compile-time exhaustiveness checking.
81+
const _exhaustiveCheck: never = inputDef.type;
82+
void _exhaustiveCheck;
83+
throw new Error(`Unhandled agent input type: '${inputDef.type}'`);
84+
}
85+
}
86+
87+
if (!inputDef.required) {
88+
validator = validator.optional();
89+
}
90+
91+
inputShape[key] = validator.describe(inputDef.description);
92+
}
93+
94+
// Cast required because Zod can't infer the discriminator from dynamic keys
95+
return z.object(
96+
inputShape,
97+
) as z.ZodDiscriminatedUnionOption<'agent_name'>;
98+
});
99+
100+
// Create the discriminated union
101+
// z.discriminatedUnion requires at least 2 options, so we handle the single agent case
102+
if (agentSchemas.length === 1) {
103+
schema = agentSchemas[0];
104+
} else {
105+
schema = z.discriminatedUnion(
106+
'agent_name',
107+
agentSchemas as [
108+
z.ZodDiscriminatedUnionOption<'agent_name'>,
109+
z.ZodDiscriminatedUnionOption<'agent_name'>,
110+
...Array<z.ZodDiscriminatedUnionOption<'agent_name'>>,
111+
],
112+
);
113+
}
114+
}
115+
116+
super(
117+
DELEGATE_TO_AGENT_TOOL_NAME,
118+
'Delegate to Agent',
119+
registry.getToolDescription(),
120+
Kind.Think,
121+
zodToJsonSchema(schema),
122+
/* isOutputMarkdown */ true,
123+
/* canUpdateOutput */ true,
124+
messageBus,
125+
);
126+
}
127+
128+
protected createInvocation(
129+
params: DelegateParams,
130+
): ToolInvocation<DelegateParams, ToolResult> {
131+
return new DelegateInvocation(
132+
params,
133+
this.registry,
134+
this.config,
135+
this.messageBus,
136+
);
137+
}
138+
}
139+
140+
class DelegateInvocation extends BaseToolInvocation<
141+
DelegateParams,
142+
ToolResult
143+
> {
144+
constructor(
145+
params: DelegateParams,
146+
private readonly registry: AgentRegistry,
147+
private readonly config: Config,
148+
messageBus?: MessageBus,
149+
) {
150+
super(params, messageBus, DELEGATE_TO_AGENT_TOOL_NAME);
151+
}
152+
153+
getDescription(): string {
154+
return `Delegating to agent '${this.params.agent_name}'`;
155+
}
156+
157+
async execute(
158+
signal: AbortSignal,
159+
updateOutput?: (output: string | AnsiOutput) => void,
160+
): Promise<ToolResult> {
161+
const definition = this.registry.getDefinition(this.params.agent_name);
162+
if (!definition) {
163+
throw new Error(
164+
`Agent '${this.params.agent_name}' exists in the tool definition but could not be found in the registry.`,
165+
);
166+
}
167+
168+
// Extract arguments (everything except agent_name)
169+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
170+
const { agent_name, ...agentArgs } = this.params;
171+
172+
// Instantiate the Subagent Loop
173+
const subagentInvocation = new SubagentInvocation(
174+
agentArgs as AgentInputs,
175+
definition,
176+
this.config,
177+
this.messageBus,
178+
);
179+
180+
return subagentInvocation.execute(signal, updateOutput);
181+
}
182+
}

0 commit comments

Comments
 (0)