mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-29 06:25:16 -07:00
refactor: implement DelegateToAgentTool with discriminated union (#14769)
This commit is contained in:
Abhi
@@ -0,0 +1,178 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { DelegateToAgentTool } from './delegate-to-agent-tool.js';
|
||||
import { AgentRegistry } from './registry.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { AgentDefinition } from './types.js';
|
||||
import { SubagentInvocation } from './invocation.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { MessageBusType } from '../confirmation-bus/types.js';
|
||||
import { DELEGATE_TO_AGENT_TOOL_NAME } from '../tools/tool-names.js';
|
||||
|
||||
vi.mock('./invocation.js', () => ({
|
||||
SubagentInvocation: vi.fn().mockImplementation(() => ({
|
||||
execute: vi
|
||||
.fn()
|
||||
.mockResolvedValue({ content: [{ type: 'text', text: 'Success' }] }),
|
||||
})),
|
||||
}));
|
||||
|
||||
describe('DelegateToAgentTool', () => {
|
||||
let registry: AgentRegistry;
|
||||
let config: Config;
|
||||
let tool: DelegateToAgentTool;
|
||||
let messageBus: MessageBus;
|
||||
|
||||
const mockAgentDef: AgentDefinition = {
|
||||
name: 'test_agent',
|
||||
description: 'A test agent',
|
||||
promptConfig: {},
|
||||
modelConfig: { model: 'test-model', temp: 0, top_p: 0 },
|
||||
inputConfig: {
|
||||
inputs: {
|
||||
arg1: { type: 'string', description: 'Argument 1', required: true },
|
||||
arg2: { type: 'number', description: 'Argument 2', required: false },
|
||||
},
|
||||
},
|
||||
runConfig: { max_turns: 1, max_time_minutes: 1 },
|
||||
toolConfig: { tools: [] },
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
config = {
|
||||
getDebugMode: () => false,
|
||||
modelConfigService: {
|
||||
registerRuntimeModelConfig: vi.fn(),
|
||||
},
|
||||
} as unknown as Config;
|
||||
|
||||
registry = new AgentRegistry(config);
|
||||
// Manually register the mock agent (bypassing protected method for testing)
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(registry as any).agents.set(mockAgentDef.name, mockAgentDef);
|
||||
|
||||
messageBus = {
|
||||
publish: vi.fn(),
|
||||
subscribe: vi.fn(),
|
||||
unsubscribe: vi.fn(),
|
||||
} as unknown as MessageBus;
|
||||
|
||||
tool = new DelegateToAgentTool(registry, config, messageBus);
|
||||
});
|
||||
|
||||
it('should use dynamic description from registry', () => {
|
||||
// registry has mockAgentDef registered in beforeEach
|
||||
expect(tool.description).toContain(
|
||||
'Delegates a task to a specialized sub-agent',
|
||||
);
|
||||
expect(tool.description).toContain(
|
||||
`- **${mockAgentDef.name}**: ${mockAgentDef.description}`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should validate agent_name exists in registry', async () => {
|
||||
// Zod validation happens at build time now (or rather, build validates the schema)
|
||||
// Since we use discriminated union, an invalid agent_name won't match any option.
|
||||
expect(() =>
|
||||
tool.build({
|
||||
agent_name: 'non_existent_agent',
|
||||
}),
|
||||
).toThrow();
|
||||
});
|
||||
|
||||
it('should validate correct arguments', async () => {
|
||||
const invocation = tool.build({
|
||||
agent_name: 'test_agent',
|
||||
arg1: 'valid',
|
||||
});
|
||||
|
||||
const result = await invocation.execute(new AbortController().signal);
|
||||
expect(result).toEqual({ content: [{ type: 'text', text: 'Success' }] });
|
||||
expect(SubagentInvocation).toHaveBeenCalledWith(
|
||||
{ arg1: 'valid' },
|
||||
mockAgentDef,
|
||||
config,
|
||||
messageBus,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw error for missing required argument', async () => {
|
||||
// Missing arg1 should fail Zod validation
|
||||
expect(() =>
|
||||
tool.build({
|
||||
agent_name: 'test_agent',
|
||||
arg2: 123,
|
||||
}),
|
||||
).toThrow();
|
||||
});
|
||||
|
||||
it('should throw error for invalid argument type', async () => {
|
||||
// arg1 should be string, passing number
|
||||
expect(() =>
|
||||
tool.build({
|
||||
agent_name: 'test_agent',
|
||||
arg1: 123,
|
||||
}),
|
||||
).toThrow();
|
||||
});
|
||||
|
||||
it('should allow optional arguments to be omitted', async () => {
|
||||
const invocation = tool.build({
|
||||
agent_name: 'test_agent',
|
||||
arg1: 'valid',
|
||||
// arg2 is optional
|
||||
});
|
||||
|
||||
await expect(
|
||||
invocation.execute(new AbortController().signal),
|
||||
).resolves.toBeDefined();
|
||||
});
|
||||
|
||||
it('should throw error if an agent has an input named "agent_name"', () => {
|
||||
const invalidAgentDef: AgentDefinition = {
|
||||
...mockAgentDef,
|
||||
name: 'invalid_agent',
|
||||
inputConfig: {
|
||||
inputs: {
|
||||
agent_name: {
|
||||
type: 'string',
|
||||
description: 'Conflict',
|
||||
required: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(registry as any).agents.set(invalidAgentDef.name, invalidAgentDef);
|
||||
|
||||
expect(() => new DelegateToAgentTool(registry, config)).toThrow(
|
||||
"Agent 'invalid_agent' cannot have an input parameter named 'agent_name' as it is a reserved parameter for delegation.",
|
||||
);
|
||||
});
|
||||
|
||||
it('should use correct tool name "delegate_to_agent" when requesting confirmation', async () => {
|
||||
const invocation = tool.build({
|
||||
agent_name: 'test_agent',
|
||||
arg1: 'valid',
|
||||
});
|
||||
|
||||
// Trigger confirmation check
|
||||
const p = invocation.shouldConfirmExecute(new AbortController().signal);
|
||||
void p;
|
||||
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_REQUEST,
|
||||
toolCall: expect.objectContaining({
|
||||
name: DELEGATE_TO_AGENT_TOOL_NAME,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { z } from 'zod';
|
||||
import { zodToJsonSchema } from 'zod-to-json-schema';
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
BaseToolInvocation,
|
||||
} from '../tools/tools.js';
|
||||
import type { AnsiOutput } from '../utils/terminalSerializer.js';
|
||||
import { DELEGATE_TO_AGENT_TOOL_NAME } from '../tools/tool-names.js';
|
||||
import type { AgentRegistry } from './registry.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { SubagentInvocation } from './invocation.js';
|
||||
import type { AgentInputs } from './types.js';
|
||||
|
||||
type DelegateParams = { agent_name: string } & Record<string, unknown>;
|
||||
|
||||
export class DelegateToAgentTool extends BaseDeclarativeTool<
|
||||
DelegateParams,
|
||||
ToolResult
|
||||
> {
|
||||
constructor(
|
||||
private readonly registry: AgentRegistry,
|
||||
private readonly config: Config,
|
||||
messageBus?: MessageBus,
|
||||
) {
|
||||
const definitions = registry.getAllDefinitions();
|
||||
|
||||
let schema: z.ZodTypeAny;
|
||||
|
||||
if (definitions.length === 0) {
|
||||
// Fallback if no agents are registered (mostly for testing/safety)
|
||||
schema = z.object({
|
||||
agent_name: z.string().describe('No agents are currently available.'),
|
||||
});
|
||||
} else {
|
||||
const agentSchemas = definitions.map((def) => {
|
||||
const inputShape: Record<string, z.ZodTypeAny> = {
|
||||
agent_name: z.literal(def.name).describe(def.description),
|
||||
};
|
||||
|
||||
for (const [key, inputDef] of Object.entries(def.inputConfig.inputs)) {
|
||||
if (key === 'agent_name') {
|
||||
throw new Error(
|
||||
`Agent '${def.name}' cannot have an input parameter named 'agent_name' as it is a reserved parameter for delegation.`,
|
||||
);
|
||||
}
|
||||
|
||||
let validator: z.ZodTypeAny;
|
||||
|
||||
// Map input types to Zod
|
||||
switch (inputDef.type) {
|
||||
case 'string':
|
||||
validator = z.string();
|
||||
break;
|
||||
case 'number':
|
||||
validator = z.number();
|
||||
break;
|
||||
case 'boolean':
|
||||
validator = z.boolean();
|
||||
break;
|
||||
case 'integer':
|
||||
validator = z.number().int();
|
||||
break;
|
||||
case 'string[]':
|
||||
validator = z.array(z.string());
|
||||
break;
|
||||
case 'number[]':
|
||||
validator = z.array(z.number());
|
||||
break;
|
||||
default: {
|
||||
// This provides compile-time exhaustiveness checking.
|
||||
const _exhaustiveCheck: never = inputDef.type;
|
||||
void _exhaustiveCheck;
|
||||
throw new Error(`Unhandled agent input type: '${inputDef.type}'`);
|
||||
}
|
||||
}
|
||||
|
||||
if (!inputDef.required) {
|
||||
validator = validator.optional();
|
||||
}
|
||||
|
||||
inputShape[key] = validator.describe(inputDef.description);
|
||||
}
|
||||
|
||||
// Cast required because Zod can't infer the discriminator from dynamic keys
|
||||
return z.object(
|
||||
inputShape,
|
||||
) as z.ZodDiscriminatedUnionOption<'agent_name'>;
|
||||
});
|
||||
|
||||
// Create the discriminated union
|
||||
// z.discriminatedUnion requires at least 2 options, so we handle the single agent case
|
||||
if (agentSchemas.length === 1) {
|
||||
schema = agentSchemas[0];
|
||||
} else {
|
||||
schema = z.discriminatedUnion(
|
||||
'agent_name',
|
||||
agentSchemas as [
|
||||
z.ZodDiscriminatedUnionOption<'agent_name'>,
|
||||
z.ZodDiscriminatedUnionOption<'agent_name'>,
|
||||
...Array<z.ZodDiscriminatedUnionOption<'agent_name'>>,
|
||||
],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
super(
|
||||
DELEGATE_TO_AGENT_TOOL_NAME,
|
||||
'Delegate to Agent',
|
||||
registry.getToolDescription(),
|
||||
Kind.Think,
|
||||
zodToJsonSchema(schema),
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
messageBus,
|
||||
);
|
||||
}
|
||||
|
||||
protected createInvocation(
|
||||
params: DelegateParams,
|
||||
): ToolInvocation<DelegateParams, ToolResult> {
|
||||
return new DelegateInvocation(
|
||||
params,
|
||||
this.registry,
|
||||
this.config,
|
||||
this.messageBus,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
class DelegateInvocation extends BaseToolInvocation<
|
||||
DelegateParams,
|
||||
ToolResult
|
||||
> {
|
||||
constructor(
|
||||
params: DelegateParams,
|
||||
private readonly registry: AgentRegistry,
|
||||
private readonly config: Config,
|
||||
messageBus?: MessageBus,
|
||||
) {
|
||||
super(params, messageBus, DELEGATE_TO_AGENT_TOOL_NAME);
|
||||
}
|
||||
|
||||
getDescription(): string {
|
||||
return `Delegating to agent '${this.params.agent_name}'`;
|
||||
}
|
||||
|
||||
async execute(
|
||||
signal: AbortSignal,
|
||||
updateOutput?: (output: string | AnsiOutput) => void,
|
||||
): Promise<ToolResult> {
|
||||
const definition = this.registry.getDefinition(this.params.agent_name);
|
||||
if (!definition) {
|
||||
throw new Error(
|
||||
`Agent '${this.params.agent_name}' exists in the tool definition but could not be found in the registry.`,
|
||||
);
|
||||
}
|
||||
|
||||
// Extract arguments (everything except agent_name)
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
const { agent_name, ...agentArgs } = this.params;
|
||||
|
||||
// Instantiate the Subagent Loop
|
||||
const subagentInvocation = new SubagentInvocation(
|
||||
agentArgs as AgentInputs,
|
||||
definition,
|
||||
this.config,
|
||||
this.messageBus,
|
||||
);
|
||||
|
||||
return subagentInvocation.execute(signal, updateOutput);
|
||||
}
|
||||
}
|
||||
@@ -237,4 +237,33 @@ describe('AgentRegistry', () => {
|
||||
);
|
||||
});
|
||||
});
|
||||
describe('getToolDescription', () => {
|
||||
it('should return default message when no agents are registered', () => {
|
||||
expect(registry.getToolDescription()).toContain(
|
||||
'No agents are currently available',
|
||||
);
|
||||
});
|
||||
|
||||
it('should return formatted list of agents when agents are available', () => {
|
||||
registry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
registry.testRegisterAgent({
|
||||
...MOCK_AGENT_V2,
|
||||
name: 'AnotherAgent',
|
||||
description: 'Another agent description',
|
||||
});
|
||||
|
||||
const description = registry.getToolDescription();
|
||||
|
||||
expect(description).toContain(
|
||||
'Delegates a task to a specialized sub-agent',
|
||||
);
|
||||
expect(description).toContain('Available agents:');
|
||||
expect(description).toContain(
|
||||
`- **${MOCK_AGENT_V1.name}**: ${MOCK_AGENT_V1.description}`,
|
||||
);
|
||||
expect(description).toContain(
|
||||
`- **AnotherAgent**: Another agent description`,
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -150,4 +150,50 @@ export class AgentRegistry {
|
||||
getAllDefinitions(): AgentDefinition[] {
|
||||
return Array.from(this.agents.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a list of all registered agent names.
|
||||
*/
|
||||
getAllAgentNames(): string[] {
|
||||
return Array.from(this.agents.keys());
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a description for the delegate_to_agent tool.
|
||||
* Unlike getDirectoryContext() which is for system prompts,
|
||||
* this is formatted for tool descriptions.
|
||||
*/
|
||||
getToolDescription(): string {
|
||||
if (this.agents.size === 0) {
|
||||
return 'Delegates a task to a specialized sub-agent. No agents are currently available.';
|
||||
}
|
||||
|
||||
const agentDescriptions = Array.from(this.agents.entries())
|
||||
.map(([name, def]) => `- **${name}**: ${def.description}`)
|
||||
.join('\n');
|
||||
|
||||
return `Delegates a task to a specialized sub-agent.
|
||||
|
||||
Available agents:
|
||||
${agentDescriptions}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a markdown "Phone Book" of available agents and their schemas.
|
||||
* This MUST be injected into the System Prompt of the parent agent.
|
||||
*/
|
||||
getDirectoryContext(): string {
|
||||
if (this.agents.size === 0) {
|
||||
return 'No sub-agents are currently available.';
|
||||
}
|
||||
|
||||
let context = '## Available Sub-Agents\n';
|
||||
context +=
|
||||
'Use `delegate_to_agent` for complex tasks requiring specialized analysis.\n\n';
|
||||
|
||||
for (const [name, def] of this.agents.entries()) {
|
||||
context += `- **${name}**: ${def.description}\n`;
|
||||
}
|
||||
return context;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user