mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-15 07:41:03 -07:00
refactor(core): remove legacy subagent wrapping tools (#25053)
This commit is contained in:
Abhi
@@ -1,187 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import { LocalSubagentInvocation } from './local-invocation.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import type { LocalAgentDefinition, AgentInputs } from './types.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { Kind } from '../tools/tools.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
|
||||
// Mock dependencies to isolate the SubagentToolWrapper class
|
||||
vi.mock('./local-invocation.js');
|
||||
|
||||
const MockedLocalSubagentInvocation = vi.mocked(LocalSubagentInvocation);
|
||||
|
||||
// Define reusable test data
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
const mockDefinition: LocalAgentDefinition = {
|
||||
kind: 'local',
|
||||
name: 'TestAgent',
|
||||
displayName: 'Test Agent Display Name',
|
||||
description: 'An agent for testing.',
|
||||
inputConfig: {
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
goal: { type: 'string', description: 'The goal.' },
|
||||
priority: {
|
||||
type: 'number',
|
||||
description: 'The priority.',
|
||||
},
|
||||
},
|
||||
required: ['goal'],
|
||||
},
|
||||
},
|
||||
modelConfig: {
|
||||
model: 'gemini-test-model',
|
||||
generateContentConfig: {
|
||||
temperature: 0,
|
||||
topP: 1,
|
||||
},
|
||||
},
|
||||
runConfig: { maxTimeMinutes: 5 },
|
||||
promptConfig: { systemPrompt: 'You are a test agent.' },
|
||||
};
|
||||
|
||||
describe('SubagentToolWrapper', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should correctly configure the tool properties from the agent definition', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
|
||||
expect(wrapper.name).toBe(mockDefinition.name);
|
||||
expect(wrapper.displayName).toBe(mockDefinition.displayName);
|
||||
expect(wrapper.description).toBe(mockDefinition.description);
|
||||
expect(wrapper.kind).toBe(Kind.Agent);
|
||||
expect(wrapper.isOutputMarkdown).toBe(true);
|
||||
expect(wrapper.canUpdateOutput).toBe(true);
|
||||
});
|
||||
|
||||
it('should fall back to the agent name for displayName if it is not provided', () => {
|
||||
const definitionWithoutDisplayName = {
|
||||
...mockDefinition,
|
||||
displayName: undefined,
|
||||
};
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
definitionWithoutDisplayName,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
expect(wrapper.displayName).toBe(definitionWithoutDisplayName.name);
|
||||
});
|
||||
|
||||
it('should generate a valid tool schema using the definition and converted schema', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const schema = wrapper.schema;
|
||||
|
||||
expect(schema.name).toBe(mockDefinition.name);
|
||||
expect(schema.description).toBe(mockDefinition.description);
|
||||
expect(schema.parametersJsonSchema).toEqual({
|
||||
...(mockDefinition.inputConfig.inputSchema as Record<string, unknown>),
|
||||
properties: {
|
||||
...((
|
||||
mockDefinition.inputConfig.inputSchema as Record<string, unknown>
|
||||
)['properties'] as Record<string, unknown>),
|
||||
wait_for_previous: {
|
||||
type: 'boolean',
|
||||
description:
|
||||
'Set to true to wait for all previously requested tools in this turn to complete before starting. Set to false (or omit) to run in parallel. Use true when this tool depends on the output of previous tools.',
|
||||
},
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('createInvocation', () => {
|
||||
it('should create a LocalSubagentInvocation with the correct parameters', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params: AgentInputs = { goal: 'Test the invocation', priority: 1 };
|
||||
|
||||
// The public `build` method calls the protected `createInvocation` after validation
|
||||
const invocation = wrapper.build(params);
|
||||
|
||||
expect(invocation).toBeInstanceOf(LocalSubagentInvocation);
|
||||
expect(MockedLocalSubagentInvocation).toHaveBeenCalledExactlyOnceWith(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
params,
|
||||
mockMessageBus,
|
||||
mockDefinition.name,
|
||||
mockDefinition.displayName,
|
||||
);
|
||||
});
|
||||
|
||||
it('should pass the messageBus to the LocalSubagentInvocation constructor', () => {
|
||||
const specificMessageBus = {
|
||||
publish: vi.fn(),
|
||||
subscribe: vi.fn(),
|
||||
unsubscribe: vi.fn(),
|
||||
} as unknown as MessageBus;
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
specificMessageBus,
|
||||
);
|
||||
const params: AgentInputs = { goal: 'Test the invocation', priority: 1 };
|
||||
|
||||
wrapper.build(params);
|
||||
|
||||
expect(MockedLocalSubagentInvocation).toHaveBeenCalledWith(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
params,
|
||||
specificMessageBus,
|
||||
mockDefinition.name,
|
||||
mockDefinition.displayName,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw a validation error for invalid parameters before creating an invocation', () => {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
// Missing the required 'goal' parameter
|
||||
const invalidParams = { priority: 1 };
|
||||
|
||||
// The `build` method in the base class performs JSON schema validation
|
||||
// before calling the protected `createInvocation` method.
|
||||
expect(() => wrapper.build(invalidParams)).toThrow(
|
||||
"params must have required property 'goal'",
|
||||
);
|
||||
expect(MockedLocalSubagentInvocation).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,106 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
} from '../tools/tools.js';
|
||||
|
||||
import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import { LocalSubagentInvocation } from './local-invocation.js';
|
||||
import { RemoteAgentInvocation } from './remote-invocation.js';
|
||||
import { BrowserAgentInvocation } from './browser/browserAgentInvocation.js';
|
||||
import { BROWSER_AGENT_NAME } from './browser/browserAgentDefinition.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
|
||||
/**
|
||||
* A tool wrapper that dynamically exposes a subagent as a standard,
|
||||
* strongly-typed `DeclarativeTool`.
|
||||
*/
|
||||
export class SubagentToolWrapper extends BaseDeclarativeTool<
|
||||
AgentInputs,
|
||||
ToolResult
|
||||
> {
|
||||
/**
|
||||
* Constructs the tool wrapper.
|
||||
*
|
||||
* The constructor dynamically generates the JSON schema for the tool's
|
||||
* parameters based on the subagent's input configuration.
|
||||
*
|
||||
* @param definition The `AgentDefinition` of the subagent to wrap.
|
||||
* @param context The execution context.
|
||||
* @param messageBus Optional message bus for policy enforcement.
|
||||
*/
|
||||
constructor(
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
) {
|
||||
super(
|
||||
definition.name,
|
||||
definition.displayName ?? definition.name,
|
||||
definition.description,
|
||||
Kind.Agent,
|
||||
definition.inputConfig.inputSchema,
|
||||
messageBus,
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates an invocation instance for executing the subagent.
|
||||
*
|
||||
* This method is called by the tool framework when the parent agent decides
|
||||
* to use this tool.
|
||||
*
|
||||
* @param params The validated input parameters from the parent agent's call.
|
||||
* @returns A `ToolInvocation` instance ready for execution.
|
||||
*/
|
||||
protected createInvocation(
|
||||
params: AgentInputs,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
const definition = this.definition;
|
||||
const effectiveMessageBus = messageBus;
|
||||
|
||||
if (definition.kind === 'remote') {
|
||||
return new RemoteAgentInvocation(
|
||||
definition,
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
|
||||
// Special handling for browser agent - needs async MCP setup
|
||||
if (definition.name === BROWSER_AGENT_NAME) {
|
||||
return new BrowserAgentInvocation(
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
|
||||
return new LocalSubagentInvocation(
|
||||
definition,
|
||||
this.context,
|
||||
params,
|
||||
effectiveMessageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,424 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { SubagentTool } from './subagent-tool.js';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import {
|
||||
Kind,
|
||||
type DeclarativeTool,
|
||||
type ToolCallConfirmationDetails,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
} from '../tools/tools.js';
|
||||
import type {
|
||||
LocalAgentDefinition,
|
||||
RemoteAgentDefinition,
|
||||
AgentInputs,
|
||||
} from './types.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import {
|
||||
GeminiCliOperation,
|
||||
GEN_AI_AGENT_DESCRIPTION,
|
||||
GEN_AI_AGENT_NAME,
|
||||
} from '../telemetry/constants.js';
|
||||
import type { ToolRegistry } from 'src/tools/tool-registry.js';
|
||||
|
||||
vi.mock('./subagent-tool-wrapper.js');
|
||||
|
||||
// Mock runInDevTraceSpan
|
||||
const runInDevTraceSpan = vi.hoisted(() =>
|
||||
vi.fn(async (opts, fn) => {
|
||||
const metadata = { attributes: opts.attributes || {} };
|
||||
return fn({
|
||||
metadata,
|
||||
});
|
||||
}),
|
||||
);
|
||||
|
||||
vi.mock('../telemetry/trace.js', () => ({
|
||||
runInDevTraceSpan,
|
||||
}));
|
||||
|
||||
const MockSubagentToolWrapper = vi.mocked(SubagentToolWrapper);
|
||||
|
||||
const testDefinition: LocalAgentDefinition = {
|
||||
kind: 'local',
|
||||
name: 'LocalAgent',
|
||||
description: 'A local agent.',
|
||||
inputConfig: { inputSchema: { type: 'object', properties: {} } },
|
||||
modelConfig: { model: 'test', generateContentConfig: {} },
|
||||
runConfig: { maxTimeMinutes: 1 },
|
||||
promptConfig: { systemPrompt: 'test' },
|
||||
};
|
||||
|
||||
const testRemoteDefinition: RemoteAgentDefinition = {
|
||||
kind: 'remote',
|
||||
name: 'RemoteAgent',
|
||||
description: 'A remote agent.',
|
||||
inputConfig: {
|
||||
inputSchema: { type: 'object', properties: { query: { type: 'string' } } },
|
||||
},
|
||||
agentCardUrl: 'http://example.com/agent',
|
||||
};
|
||||
|
||||
describe('SubAgentInvocation', () => {
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
let mockInnerInvocation: ToolInvocation<AgentInputs, ToolResult>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
mockInnerInvocation = {
|
||||
shouldConfirmExecute: vi.fn(),
|
||||
execute: vi.fn(),
|
||||
params: {},
|
||||
getDescription: vi.fn(),
|
||||
toolLocations: vi.fn(),
|
||||
};
|
||||
|
||||
MockSubagentToolWrapper.prototype.build = vi
|
||||
.fn()
|
||||
.mockReturnValue(mockInnerInvocation);
|
||||
});
|
||||
|
||||
it('should have Kind.Agent', () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
expect(tool.kind).toBe(Kind.Agent);
|
||||
});
|
||||
|
||||
it('should delegate shouldConfirmExecute to the inner sub-invocation (local)', async () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
vi.mocked(mockInnerInvocation.shouldConfirmExecute).mockResolvedValue(
|
||||
false,
|
||||
);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const result = await invocation.shouldConfirmExecute(abortSignal);
|
||||
|
||||
expect(result).toBe(false);
|
||||
expect(mockInnerInvocation.shouldConfirmExecute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
);
|
||||
expect(MockSubagentToolWrapper).toHaveBeenCalledWith(
|
||||
testDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
});
|
||||
|
||||
it('should return the correct description', () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
expect(invocation.getDescription()).toBe(
|
||||
"Delegating to agent 'LocalAgent'",
|
||||
);
|
||||
});
|
||||
|
||||
it('should delegate shouldConfirmExecute to the inner sub-invocation (remote)', async () => {
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'test' };
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
const confirmationDetails = {
|
||||
type: 'info',
|
||||
title: 'Confirm',
|
||||
prompt: 'Prompt',
|
||||
onConfirm: vi.fn(),
|
||||
} as const;
|
||||
vi.mocked(mockInnerInvocation.shouldConfirmExecute).mockResolvedValue(
|
||||
confirmationDetails as unknown as ToolCallConfirmationDetails,
|
||||
);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const result = await invocation.shouldConfirmExecute(abortSignal);
|
||||
|
||||
expect(result).toBe(confirmationDetails);
|
||||
expect(mockInnerInvocation.shouldConfirmExecute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
);
|
||||
expect(MockSubagentToolWrapper).toHaveBeenCalledWith(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
});
|
||||
|
||||
it('should delegate execute to the inner sub-invocation', async () => {
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = {};
|
||||
// @ts-expect-error - accessing protected method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
const mockResult: ToolResult = {
|
||||
llmContent: 'success',
|
||||
returnDisplay: 'success',
|
||||
};
|
||||
vi.mocked(mockInnerInvocation.execute).mockResolvedValue(mockResult);
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const updateOutput = vi.fn();
|
||||
const result = await invocation.execute(abortSignal, updateOutput);
|
||||
|
||||
expect(result).toBe(mockResult);
|
||||
expect(mockInnerInvocation.execute).toHaveBeenCalledWith(
|
||||
abortSignal,
|
||||
updateOutput,
|
||||
);
|
||||
|
||||
expect(runInDevTraceSpan).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
operation: GeminiCliOperation.AgentCall,
|
||||
attributes: expect.objectContaining({
|
||||
[GEN_AI_AGENT_NAME]: testDefinition.name,
|
||||
[GEN_AI_AGENT_DESCRIPTION]: testDefinition.description,
|
||||
}),
|
||||
}),
|
||||
expect.any(Function),
|
||||
);
|
||||
|
||||
// Verify metadata was set on the span
|
||||
const spanCallback = vi.mocked(runInDevTraceSpan).mock.calls[0][1];
|
||||
const mockMetadata = { input: undefined, output: undefined };
|
||||
const mockSpan = { metadata: mockMetadata };
|
||||
await spanCallback(mockSpan as Parameters<typeof spanCallback>[0]);
|
||||
expect(mockMetadata.input).toBe(params);
|
||||
expect(mockMetadata.output).toBe(mockResult);
|
||||
});
|
||||
|
||||
describe('withUserHints', () => {
|
||||
it('should NOT modify query for local agents', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Test Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(testDefinition, mockConfig, mockMessageBus);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should NOT modify query for remote agents if model steering is disabled', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: false });
|
||||
mockConfig.injectionService.addInjection('Test Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should NOT modify query for remote agents if there are no hints', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
});
|
||||
|
||||
it('should prepend hints to query for remote agents when hints exist and steering is enabled', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
mockConfig.injectionService.addInjection('Hint 1', 'user_steering');
|
||||
mockConfig.injectionService.addInjection('Hint 2', 'user_steering');
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toContain('Hint 1');
|
||||
expect(hintedParams.query).toContain('Hint 2');
|
||||
expect(hintedParams.query).toMatch(/original query$/);
|
||||
});
|
||||
|
||||
it('should NOT include legacy hints added before the invocation was created', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Legacy Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { query: 'original query' };
|
||||
|
||||
// Creation of invocation captures the current hint state
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// Verify no hints are present yet
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
let hintedParams = invocation.withUserHints(params);
|
||||
expect(hintedParams.query).toBe('original query');
|
||||
|
||||
// Add a new hint after creation
|
||||
mockConfig.injectionService.addInjection('New Hint', 'user_steering');
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams.query).toContain('New Hint');
|
||||
expect(hintedParams.query).not.toContain('Legacy Hint');
|
||||
});
|
||||
|
||||
it('should NOT modify query if query is missing or not a string', async () => {
|
||||
mockConfig = makeFakeConfig({ modelSteering: true });
|
||||
mockConfig.injectionService.addInjection('Hint', 'user_steering');
|
||||
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
const params = { other: 'param' };
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const invocation = tool.createInvocation(params, mockMessageBus);
|
||||
|
||||
// @ts-expect-error - accessing private method for testing
|
||||
const hintedParams = invocation.withUserHints(params);
|
||||
|
||||
expect(hintedParams).toEqual(params);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('SubagentTool Read-Only logic', () => {
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// .config is already set correctly by the getter on the instance.
|
||||
Object.defineProperty(mockConfig, 'promptId', {
|
||||
get: () => 'test-prompt-id',
|
||||
configurable: true,
|
||||
});
|
||||
mockMessageBus = createMockMessageBus();
|
||||
});
|
||||
|
||||
it('should be false for remote agents', () => {
|
||||
const tool = new SubagentTool(
|
||||
testRemoteDefinition,
|
||||
mockConfig,
|
||||
mockMessageBus,
|
||||
);
|
||||
expect(tool.isReadOnly).toBe(false);
|
||||
});
|
||||
|
||||
it('should be true for local agent with only read-only tools', () => {
|
||||
const readOnlyTool = {
|
||||
name: 'read',
|
||||
isReadOnly: true,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const registry = {
|
||||
getTool: (name: string) => (name === 'read' ? readOnlyTool : undefined),
|
||||
};
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defWithTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: ['read'] },
|
||||
};
|
||||
const tool = new SubagentTool(defWithTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(true);
|
||||
});
|
||||
|
||||
it('should be false for local agent with at least one non-read-only tool', () => {
|
||||
const readOnlyTool = {
|
||||
name: 'read',
|
||||
isReadOnly: true,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const mutatorTool = {
|
||||
name: 'write',
|
||||
isReadOnly: false,
|
||||
} as unknown as DeclarativeTool<object, ToolResult>;
|
||||
const registry = {
|
||||
getTool: (name: string) => {
|
||||
if (name === 'read') return readOnlyTool;
|
||||
if (name === 'write') return mutatorTool;
|
||||
return undefined;
|
||||
},
|
||||
};
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defWithTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: ['read', 'write'] },
|
||||
};
|
||||
const tool = new SubagentTool(defWithTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(false);
|
||||
});
|
||||
|
||||
it('should be true for local agent with no tools', () => {
|
||||
const registry = { getTool: () => undefined };
|
||||
vi.spyOn(mockConfig, 'toolRegistry', 'get').mockReturnValue(
|
||||
registry as unknown as ToolRegistry,
|
||||
);
|
||||
|
||||
const defNoTools: LocalAgentDefinition = {
|
||||
...testDefinition,
|
||||
toolConfig: { tools: [] },
|
||||
};
|
||||
const tool = new SubagentTool(defNoTools, mockConfig, mockMessageBus);
|
||||
expect(tool.isReadOnly).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -1,237 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
BaseToolInvocation,
|
||||
type ToolCallConfirmationDetails,
|
||||
isTool,
|
||||
type ToolLiveOutput,
|
||||
} from '../tools/tools.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { type AgentLoopContext } from '../config/agent-loop-context.js';
|
||||
import type { MessageBus } from '../confirmation-bus/message-bus.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import { SubagentToolWrapper } from './subagent-tool-wrapper.js';
|
||||
import { SchemaValidator } from '../utils/schemaValidator.js';
|
||||
import { formatUserHintsForModel } from '../utils/fastAckHelper.js';
|
||||
import { runInDevTraceSpan } from '../telemetry/trace.js';
|
||||
import {
|
||||
GeminiCliOperation,
|
||||
GEN_AI_AGENT_DESCRIPTION,
|
||||
GEN_AI_AGENT_NAME,
|
||||
} from '../telemetry/constants.js';
|
||||
|
||||
export class SubagentTool extends BaseDeclarativeTool<AgentInputs, ToolResult> {
|
||||
constructor(
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
) {
|
||||
const inputSchema = definition.inputConfig.inputSchema;
|
||||
|
||||
// Validate schema on construction
|
||||
const schemaError = SchemaValidator.validateSchema(inputSchema);
|
||||
if (schemaError) {
|
||||
throw new Error(
|
||||
`Invalid schema for agent ${definition.name}: ${schemaError}`,
|
||||
);
|
||||
}
|
||||
|
||||
super(
|
||||
definition.name,
|
||||
definition.displayName ?? definition.name,
|
||||
definition.description,
|
||||
Kind.Agent,
|
||||
inputSchema,
|
||||
messageBus,
|
||||
/* isOutputMarkdown */ true,
|
||||
/* canUpdateOutput */ true,
|
||||
);
|
||||
}
|
||||
|
||||
private _memoizedIsReadOnly: boolean | undefined;
|
||||
|
||||
override get isReadOnly(): boolean {
|
||||
if (this._memoizedIsReadOnly !== undefined) {
|
||||
return this._memoizedIsReadOnly;
|
||||
}
|
||||
// No try-catch here. If getToolRegistry() throws, we let it throw.
|
||||
// This is an invariant: you can't check read-only status if the system isn't initialized.
|
||||
this._memoizedIsReadOnly = SubagentTool.checkIsReadOnly(
|
||||
this.definition,
|
||||
this.context,
|
||||
);
|
||||
return this._memoizedIsReadOnly;
|
||||
}
|
||||
|
||||
private static checkIsReadOnly(
|
||||
definition: AgentDefinition,
|
||||
context: AgentLoopContext,
|
||||
): boolean {
|
||||
if (definition.kind === 'remote') {
|
||||
return false;
|
||||
}
|
||||
const tools = definition.toolConfig?.tools ?? [];
|
||||
const registry = context.toolRegistry;
|
||||
|
||||
if (!registry) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const tool of tools) {
|
||||
if (typeof tool === 'string') {
|
||||
const resolvedTool = registry.getTool(tool);
|
||||
if (!resolvedTool || !resolvedTool.isReadOnly) {
|
||||
return false;
|
||||
}
|
||||
} else if (isTool(tool)) {
|
||||
if (!tool.isReadOnly) {
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
// FunctionDeclaration - we don't know, so assume NOT read-only
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
protected createInvocation(
|
||||
params: AgentInputs,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
return new SubAgentInvocation(
|
||||
params,
|
||||
this.definition,
|
||||
this.context,
|
||||
messageBus,
|
||||
_toolName,
|
||||
_toolDisplayName,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
class SubAgentInvocation extends BaseToolInvocation<AgentInputs, ToolResult> {
|
||||
private readonly startIndex: number;
|
||||
|
||||
constructor(
|
||||
params: AgentInputs,
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly context: AgentLoopContext,
|
||||
messageBus: MessageBus,
|
||||
_toolName?: string,
|
||||
_toolDisplayName?: string,
|
||||
) {
|
||||
super(
|
||||
params,
|
||||
messageBus,
|
||||
_toolName ?? definition.name,
|
||||
_toolDisplayName ?? definition.displayName ?? definition.name,
|
||||
);
|
||||
this.startIndex = context.config.injectionService.getLatestInjectionIndex();
|
||||
}
|
||||
|
||||
private get config(): Config {
|
||||
return this.context.config;
|
||||
}
|
||||
|
||||
getDescription(): string {
|
||||
return `Delegating to agent '${this.definition.name}'`;
|
||||
}
|
||||
|
||||
override async shouldConfirmExecute(
|
||||
abortSignal: AbortSignal,
|
||||
): Promise<ToolCallConfirmationDetails | false> {
|
||||
const invocation = this.buildSubInvocation(
|
||||
this.definition,
|
||||
this.withUserHints(this.params),
|
||||
);
|
||||
return invocation.shouldConfirmExecute(abortSignal);
|
||||
}
|
||||
|
||||
async execute(
|
||||
signal: AbortSignal,
|
||||
updateOutput?: (output: ToolLiveOutput) => void,
|
||||
): Promise<ToolResult> {
|
||||
const validationError = SchemaValidator.validate(
|
||||
this.definition.inputConfig.inputSchema,
|
||||
this.params,
|
||||
);
|
||||
|
||||
if (validationError) {
|
||||
throw new Error(
|
||||
`Invalid arguments for agent '${this.definition.name}': ${validationError}. Input schema: ${JSON.stringify(this.definition.inputConfig.inputSchema)}.`,
|
||||
);
|
||||
}
|
||||
|
||||
const invocation = this.buildSubInvocation(
|
||||
this.definition,
|
||||
this.withUserHints(this.params),
|
||||
);
|
||||
|
||||
return runInDevTraceSpan(
|
||||
{
|
||||
operation: GeminiCliOperation.AgentCall,
|
||||
logPrompts: this.context.config.getTelemetryLogPromptsEnabled(),
|
||||
sessionId: this.context.config.getSessionId(),
|
||||
attributes: {
|
||||
[GEN_AI_AGENT_NAME]: this.definition.name,
|
||||
[GEN_AI_AGENT_DESCRIPTION]: this.definition.description,
|
||||
},
|
||||
},
|
||||
async ({ metadata }) => {
|
||||
metadata.input = this.params;
|
||||
const result = await invocation.execute(signal, updateOutput);
|
||||
metadata.output = result;
|
||||
return result;
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
private withUserHints(agentArgs: AgentInputs): AgentInputs {
|
||||
if (this.definition.kind !== 'remote') {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
const userHints = this.config.injectionService.getInjectionsAfter(
|
||||
this.startIndex,
|
||||
'user_steering',
|
||||
);
|
||||
const formattedHints = formatUserHintsForModel(userHints);
|
||||
if (!formattedHints) {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
const query = agentArgs['query'];
|
||||
if (typeof query !== 'string' || query.trim().length === 0) {
|
||||
return agentArgs;
|
||||
}
|
||||
|
||||
return {
|
||||
...agentArgs,
|
||||
query: `${formattedHints}\n\n${query}`,
|
||||
};
|
||||
}
|
||||
|
||||
private buildSubInvocation(
|
||||
definition: AgentDefinition,
|
||||
agentArgs: AgentInputs,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
const wrapper = new SubagentToolWrapper(
|
||||
definition,
|
||||
this.context,
|
||||
this.messageBus,
|
||||
);
|
||||
|
||||
return wrapper.build(agentArgs);
|
||||
}
|
||||
}
|
||||
@@ -192,10 +192,6 @@ vi.mock('../agents/registry.js', () => {
|
||||
return { AgentRegistry: AgentRegistryMock };
|
||||
});
|
||||
|
||||
vi.mock('../agents/subagent-tool.js', () => ({
|
||||
SubagentTool: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../resources/resource-registry.js', () => ({
|
||||
ResourceRegistry: vi.fn(),
|
||||
}));
|
||||
|
||||
Reference in New Issue
Block a user