mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 14:10:37 -07:00
refactor(agents): Introduce Declarative Agent Framework (#9778)
This commit is contained in:
Abhi
83
packages/core/src/agents/codebase-investigator.ts
Normal file
83
packages/core/src/agents/codebase-investigator.ts
Normal file
@@ -0,0 +1,83 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { AgentDefinition } from './types.js';
|
||||
import { LSTool } from '../tools/ls.js';
|
||||
import { ReadFileTool } from '../tools/read-file.js';
|
||||
import { GlobTool } from '../tools/glob.js';
|
||||
import { GrepTool } from '../tools/grep.js';
|
||||
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
|
||||
|
||||
/**
|
||||
* A Proof-of-Concept subagent specialized in analyzing codebase structure,
|
||||
* dependencies, and technologies.
|
||||
*/
|
||||
export const CodebaseInvestigatorAgent: AgentDefinition = {
|
||||
name: 'codebase_investigator',
|
||||
displayName: 'Codebase Investigator Agent',
|
||||
description:
|
||||
'A specialized agent used for analyzing and reporting on the structure, technologies, dependencies, and conventions of the current codebase. Use this when asked to understand how the project is set up or how a specific feature is implemented.',
|
||||
|
||||
inputConfig: {
|
||||
inputs: {
|
||||
investigation_focus: {
|
||||
description:
|
||||
'A high-level description of what the agent should investigate (e.g., "frontend framework", "authentication implementation", "testing conventions").',
|
||||
type: 'string',
|
||||
required: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
outputConfig: {
|
||||
description:
|
||||
'A detailed markdown report summarizing the findings of the codebase investigation.',
|
||||
completion_criteria: [
|
||||
'The report must directly address the initial `investigation_focus`.',
|
||||
'Cite specific files, functions, or configuration snippets as evidence for your findings.',
|
||||
'Conclude with a summary of the key technologies, architectural patterns, and conventions discovered.',
|
||||
],
|
||||
},
|
||||
|
||||
modelConfig: {
|
||||
model: DEFAULT_GEMINI_MODEL,
|
||||
temp: 0.2,
|
||||
top_p: 1.0,
|
||||
thinkingBudget: -1,
|
||||
},
|
||||
|
||||
runConfig: {
|
||||
max_time_minutes: 5,
|
||||
max_turns: 15,
|
||||
},
|
||||
|
||||
toolConfig: {
|
||||
// Grant access only to read-only tools.
|
||||
tools: [LSTool.Name, ReadFileTool.Name, GlobTool.Name, GrepTool.Name],
|
||||
},
|
||||
|
||||
promptConfig: {
|
||||
systemPrompt: `You are the Codebase Investigator agent. Your sole purpose is to analyze the provided codebase and generate a detailed report on a specific area of focus.
|
||||
|
||||
# Task
|
||||
Your focus for this investigation is: \${investigation_focus}
|
||||
|
||||
# Methodology
|
||||
1. **Discovery:** Start by looking at high-level configuration files (e.g., package.json, README.md, Cargo.toml, requirements.txt, build.gradle) to understand the project's dependencies and structure.
|
||||
2. **Structure Analysis:** Use '${GlobTool.Name}' and '${LSTool.Name}' to understand the directory layout and identify relevant files/modules related to your focus.
|
||||
3. **Deep Dive:** Use '${ReadFileTool.Name}' and available search tools (Grep/RipGrep) to analyze the contents of relevant files, looking for implementation details, patterns, and conventions.
|
||||
4. **Synthesis:** Synthesize all findings into a coherent markdown report.
|
||||
|
||||
# Rules
|
||||
* You MUST ONLY use the tools provided to you.
|
||||
* You CANNOT modify the codebase.
|
||||
* You must be thorough in your investigation.
|
||||
* Once you have gathered sufficient information, stop calling tools. Your findings will be synthesized into the final report.
|
||||
|
||||
# Report Format
|
||||
The final report should be structured markdown, clearly answering the investigation focus, citing the evidence (files analyzed), and summarizing the technologies/patterns found.
|
||||
`,
|
||||
},
|
||||
};
|
||||
635
packages/core/src/agents/executor.test.ts
Normal file
635
packages/core/src/agents/executor.test.ts
Normal file
@@ -0,0 +1,635 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
vi,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
type MockedClass,
|
||||
} from 'vitest';
|
||||
import { AgentExecutor, type ActivityCallback } from './executor.js';
|
||||
import type {
|
||||
AgentDefinition,
|
||||
AgentInputs,
|
||||
SubagentActivityEvent,
|
||||
} from './types.js';
|
||||
import { AgentTerminateMode } from './types.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import { LSTool } from '../tools/ls.js';
|
||||
import { ReadFileTool } from '../tools/read-file.js';
|
||||
import {
|
||||
GeminiChat,
|
||||
StreamEventType,
|
||||
type StreamEvent,
|
||||
} from '../core/geminiChat.js';
|
||||
import type {
|
||||
FunctionCall,
|
||||
Part,
|
||||
GenerateContentResponse,
|
||||
} from '@google/genai';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { MockTool } from '../test-utils/mock-tool.js';
|
||||
import { getDirectoryContextString } from '../utils/environmentContext.js';
|
||||
|
||||
const { mockSendMessageStream, mockExecuteToolCall } = vi.hoisted(() => ({
|
||||
mockSendMessageStream: vi.fn(),
|
||||
mockExecuteToolCall: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock('../core/geminiChat.js', async (importOriginal) => {
|
||||
const actual = await importOriginal();
|
||||
return {
|
||||
...(actual as object),
|
||||
GeminiChat: vi.fn().mockImplementation(() => ({
|
||||
sendMessageStream: mockSendMessageStream,
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
vi.mock('../core/nonInteractiveToolExecutor.js', () => ({
|
||||
executeToolCall: mockExecuteToolCall,
|
||||
}));
|
||||
|
||||
vi.mock('../utils/environmentContext.js');
|
||||
|
||||
const MockedGeminiChat = GeminiChat as MockedClass<typeof GeminiChat>;
|
||||
|
||||
// A mock tool that is NOT on the NON_INTERACTIVE_TOOL_ALLOWLIST
|
||||
const MOCK_TOOL_NOT_ALLOWED = new MockTool({ name: 'write_file' });
|
||||
|
||||
const createMockResponseChunk = (
|
||||
parts: Part[],
|
||||
functionCalls?: FunctionCall[],
|
||||
): GenerateContentResponse =>
|
||||
({
|
||||
candidates: [{ index: 0, content: { role: 'model', parts } }],
|
||||
functionCalls,
|
||||
}) as unknown as GenerateContentResponse;
|
||||
|
||||
const mockModelResponse = (
|
||||
functionCalls: FunctionCall[],
|
||||
thought?: string,
|
||||
text?: string,
|
||||
) => {
|
||||
const parts: Part[] = [];
|
||||
if (thought) {
|
||||
parts.push({
|
||||
text: `**${thought}** This is the reasoning part.`,
|
||||
thought: true,
|
||||
});
|
||||
}
|
||||
if (text) parts.push({ text });
|
||||
|
||||
const responseChunk = createMockResponseChunk(
|
||||
parts,
|
||||
// Ensure functionCalls is undefined if the array is empty, matching API behavior
|
||||
functionCalls.length > 0 ? functionCalls : undefined,
|
||||
);
|
||||
|
||||
mockSendMessageStream.mockImplementationOnce(async () =>
|
||||
(async function* () {
|
||||
yield {
|
||||
type: StreamEventType.CHUNK,
|
||||
value: responseChunk,
|
||||
} as StreamEvent;
|
||||
})(),
|
||||
);
|
||||
};
|
||||
|
||||
let mockConfig: Config;
|
||||
let parentToolRegistry: ToolRegistry;
|
||||
|
||||
const createTestDefinition = (
|
||||
tools: Array<string | MockTool> = [LSTool.Name],
|
||||
runConfigOverrides: Partial<AgentDefinition['runConfig']> = {},
|
||||
outputConfigOverrides: Partial<AgentDefinition['outputConfig']> = {},
|
||||
): AgentDefinition => ({
|
||||
name: 'TestAgent',
|
||||
description: 'An agent for testing.',
|
||||
inputConfig: {
|
||||
inputs: { goal: { type: 'string', required: true, description: 'goal' } },
|
||||
},
|
||||
modelConfig: { model: 'gemini-test-model', temp: 0, top_p: 1 },
|
||||
runConfig: { max_time_minutes: 5, max_turns: 5, ...runConfigOverrides },
|
||||
promptConfig: { systemPrompt: 'Achieve the goal: ${goal}.' },
|
||||
toolConfig: { tools },
|
||||
outputConfig: { description: 'The final result.', ...outputConfigOverrides },
|
||||
});
|
||||
|
||||
describe('AgentExecutor', () => {
|
||||
let activities: SubagentActivityEvent[];
|
||||
let onActivity: ActivityCallback;
|
||||
let abortController: AbortController;
|
||||
let signal: AbortSignal;
|
||||
|
||||
beforeEach(async () => {
|
||||
mockSendMessageStream.mockClear();
|
||||
mockExecuteToolCall.mockClear();
|
||||
vi.clearAllMocks();
|
||||
// Use fake timers for timeout and concurrency testing
|
||||
vi.useFakeTimers();
|
||||
|
||||
mockConfig = makeFakeConfig();
|
||||
parentToolRegistry = new ToolRegistry(mockConfig);
|
||||
parentToolRegistry.registerTool(new LSTool(mockConfig));
|
||||
parentToolRegistry.registerTool(new ReadFileTool(mockConfig));
|
||||
parentToolRegistry.registerTool(MOCK_TOOL_NOT_ALLOWED);
|
||||
|
||||
vi.spyOn(mockConfig, 'getToolRegistry').mockResolvedValue(
|
||||
parentToolRegistry,
|
||||
);
|
||||
|
||||
vi.mocked(getDirectoryContextString).mockResolvedValue(
|
||||
'Mocked Environment Context',
|
||||
);
|
||||
|
||||
activities = [];
|
||||
onActivity = (activity) => activities.push(activity);
|
||||
abortController = new AbortController();
|
||||
signal = abortController.signal;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
describe('create (Initialization and Validation)', () => {
|
||||
it('should create successfully with allowed tools', async () => {
|
||||
const definition = createTestDefinition([LSTool.Name]);
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
expect(executor).toBeInstanceOf(AgentExecutor);
|
||||
});
|
||||
|
||||
it('SECURITY: should throw if a tool is not on the non-interactive allowlist', async () => {
|
||||
const definition = createTestDefinition([MOCK_TOOL_NOT_ALLOWED.name]);
|
||||
await expect(
|
||||
AgentExecutor.create(definition, mockConfig, onActivity),
|
||||
).rejects.toThrow(
|
||||
`Tool "${MOCK_TOOL_NOT_ALLOWED.name}" is not on the allow-list for non-interactive execution`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should create an isolated ToolRegistry for the agent', async () => {
|
||||
const definition = createTestDefinition([LSTool.Name, ReadFileTool.Name]);
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
// @ts-expect-error - accessing private property for test validation
|
||||
const agentRegistry = executor.toolRegistry as ToolRegistry;
|
||||
|
||||
expect(agentRegistry).not.toBe(parentToolRegistry);
|
||||
expect(agentRegistry.getAllToolNames()).toEqual(
|
||||
expect.arrayContaining([LSTool.Name, ReadFileTool.Name]),
|
||||
);
|
||||
expect(agentRegistry.getAllToolNames()).toHaveLength(2);
|
||||
expect(agentRegistry.getTool(MOCK_TOOL_NOT_ALLOWED.name)).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('run (Execution Loop and Logic)', () => {
|
||||
it('should execute a successful work and extraction phase (Happy Path) and emit activities', async () => {
|
||||
const definition = createTestDefinition();
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
const inputs: AgentInputs = { goal: 'Find files' };
|
||||
|
||||
// Turn 1: Model calls ls
|
||||
mockModelResponse(
|
||||
[{ name: LSTool.Name, args: { path: '.' }, id: 'call1' }],
|
||||
'T1: Listing',
|
||||
);
|
||||
mockExecuteToolCall.mockResolvedValueOnce({
|
||||
callId: 'call1',
|
||||
resultDisplay: 'file1.txt',
|
||||
responseParts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: LSTool.Name,
|
||||
response: { result: 'file1.txt' },
|
||||
id: 'call1',
|
||||
},
|
||||
},
|
||||
],
|
||||
error: undefined,
|
||||
});
|
||||
|
||||
// Turn 2: Model stops
|
||||
mockModelResponse([], 'T2: Done');
|
||||
|
||||
// Extraction Phase
|
||||
mockModelResponse([], undefined, 'Result: file1.txt.');
|
||||
|
||||
const output = await executor.run(inputs, signal);
|
||||
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(3);
|
||||
expect(mockExecuteToolCall).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Verify System Prompt Templating
|
||||
const chatConstructorArgs = MockedGeminiChat.mock.calls[0];
|
||||
const chatConfig = chatConstructorArgs[1];
|
||||
expect(chatConfig?.systemInstruction).toContain(
|
||||
'Achieve the goal: Find files.',
|
||||
);
|
||||
// Verify environment context is appended
|
||||
expect(chatConfig?.systemInstruction).toContain(
|
||||
'# Environment Context\nMocked Environment Context',
|
||||
);
|
||||
// Verify standard rules are appended
|
||||
expect(chatConfig?.systemInstruction).toContain(
|
||||
'You are running in a non-interactive mode.',
|
||||
);
|
||||
// Verify absolute path rule is appended
|
||||
expect(chatConfig?.systemInstruction).toContain(
|
||||
'Always use absolute paths for file operations.',
|
||||
);
|
||||
|
||||
// Verify Extraction Phase Call (Specific arguments)
|
||||
expect(mockSendMessageStream).toHaveBeenCalledWith(
|
||||
'gemini-test-model',
|
||||
expect.objectContaining({
|
||||
// Extraction message should be based on outputConfig.description
|
||||
message: expect.arrayContaining([
|
||||
{
|
||||
text: expect.stringContaining(
|
||||
'Based on your work so far, provide: The final result.',
|
||||
),
|
||||
},
|
||||
]),
|
||||
config: expect.objectContaining({ tools: undefined }), // No tools in extraction
|
||||
}),
|
||||
expect.stringContaining('#extraction'),
|
||||
);
|
||||
|
||||
expect(output.result).toBe('Result: file1.txt.');
|
||||
expect(output.terminate_reason).toBe(AgentTerminateMode.GOAL);
|
||||
|
||||
// Verify Activity Stream (Observability)
|
||||
expect(activities).toEqual(
|
||||
expect.arrayContaining([
|
||||
// Thought subjects are extracted by the executor (parseThought)
|
||||
expect.objectContaining({
|
||||
type: 'THOUGHT_CHUNK',
|
||||
data: { text: 'T1: Listing' },
|
||||
}),
|
||||
expect.objectContaining({
|
||||
type: 'TOOL_CALL_START',
|
||||
data: { name: LSTool.Name, args: { path: '.' } },
|
||||
}),
|
||||
expect.objectContaining({
|
||||
type: 'TOOL_CALL_END',
|
||||
data: { name: LSTool.Name, output: 'file1.txt' },
|
||||
}),
|
||||
expect.objectContaining({
|
||||
type: 'THOUGHT_CHUNK',
|
||||
data: { text: 'T2: Done' },
|
||||
}),
|
||||
]),
|
||||
);
|
||||
});
|
||||
|
||||
it('should execute parallel tool calls concurrently', async () => {
|
||||
const definition = createTestDefinition([LSTool.Name, ReadFileTool.Name]);
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
const call1 = {
|
||||
name: LSTool.Name,
|
||||
args: { path: '/dir1' },
|
||||
id: 'call1',
|
||||
};
|
||||
// Using LSTool twice for simplicity in mocking standardized responses.
|
||||
const call2 = {
|
||||
name: LSTool.Name,
|
||||
args: { path: '/dir2' },
|
||||
id: 'call2',
|
||||
};
|
||||
|
||||
// Turn 1: Model calls two tools simultaneously
|
||||
mockModelResponse([call1, call2], 'T1: Listing both');
|
||||
|
||||
// Use concurrency tracking to ensure parallelism
|
||||
let activeCalls = 0;
|
||||
let maxActiveCalls = 0;
|
||||
|
||||
mockExecuteToolCall.mockImplementation(async (_ctx, reqInfo) => {
|
||||
activeCalls++;
|
||||
maxActiveCalls = Math.max(maxActiveCalls, activeCalls);
|
||||
// Simulate latency. We must advance the fake timers for this to resolve.
|
||||
await new Promise((resolve) => setTimeout(resolve, 100));
|
||||
activeCalls--;
|
||||
return {
|
||||
callId: reqInfo.callId,
|
||||
resultDisplay: `Result for ${reqInfo.name}`,
|
||||
responseParts: [
|
||||
{
|
||||
functionResponse: {
|
||||
name: reqInfo.name,
|
||||
response: {},
|
||||
id: reqInfo.callId,
|
||||
},
|
||||
},
|
||||
],
|
||||
error: undefined,
|
||||
};
|
||||
});
|
||||
|
||||
// Turn 2: Model stops
|
||||
mockModelResponse([]);
|
||||
// Extraction
|
||||
mockModelResponse([], undefined, 'Done.');
|
||||
|
||||
const runPromise = executor.run({ goal: 'Parallel test' }, signal);
|
||||
|
||||
// Advance timers while the parallel calls (Promise.all + setTimeout) are running
|
||||
await vi.advanceTimersByTimeAsync(150);
|
||||
|
||||
await runPromise;
|
||||
|
||||
expect(mockExecuteToolCall).toHaveBeenCalledTimes(2);
|
||||
expect(maxActiveCalls).toBe(2);
|
||||
|
||||
// Verify the input to the next model call (Turn 2) contains both responses
|
||||
// sendMessageStream calls: [0] Turn 1, [1] Turn 2, [2] Extraction
|
||||
const turn2Input = mockSendMessageStream.mock.calls[1][1];
|
||||
const turn2Parts = turn2Input.message as Part[];
|
||||
|
||||
// Promise.all preserves the order of the input array.
|
||||
expect(turn2Parts.length).toBe(2);
|
||||
expect(turn2Parts[0]).toEqual(
|
||||
expect.objectContaining({
|
||||
functionResponse: expect.objectContaining({ id: 'call1' }),
|
||||
}),
|
||||
);
|
||||
expect(turn2Parts[1]).toEqual(
|
||||
expect.objectContaining({
|
||||
functionResponse: expect.objectContaining({ id: 'call2' }),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle tool execution failure gracefully and report error', async () => {
|
||||
const definition = createTestDefinition([LSTool.Name]);
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Turn 1: Model calls ls, but it fails
|
||||
mockModelResponse([
|
||||
{ name: LSTool.Name, args: { path: '/invalid' }, id: 'call1' },
|
||||
]);
|
||||
|
||||
const errorMessage = 'Internal failure.';
|
||||
mockExecuteToolCall.mockResolvedValueOnce({
|
||||
callId: 'call1',
|
||||
resultDisplay: `Error: ${errorMessage}`,
|
||||
responseParts: undefined, // Failed tools might return undefined parts
|
||||
error: { message: errorMessage },
|
||||
});
|
||||
|
||||
// Turn 2: Model stops
|
||||
mockModelResponse([]);
|
||||
mockModelResponse([], undefined, 'Failed.');
|
||||
|
||||
await executor.run({ goal: 'Failure test' }, signal);
|
||||
|
||||
// Verify that the error was reported in the activity stream
|
||||
expect(activities).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: 'ERROR',
|
||||
data: {
|
||||
error: errorMessage,
|
||||
context: 'tool_call',
|
||||
name: LSTool.Name,
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
// Verify the input to the next model call (Turn 2) contains the fallback error message
|
||||
const turn2Input = mockSendMessageStream.mock.calls[1][1];
|
||||
const turn2Parts = turn2Input.message as Part[];
|
||||
expect(turn2Parts).toEqual([
|
||||
{
|
||||
text: 'All tool calls failed. Please analyze the errors and try an alternative approach.',
|
||||
},
|
||||
]);
|
||||
});
|
||||
|
||||
it('SECURITY: should block calls to tools not registered for the agent at runtime', async () => {
|
||||
// Agent definition only includes LSTool
|
||||
const definition = createTestDefinition([LSTool.Name]);
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Turn 1: Model hallucinates a call to ReadFileTool
|
||||
// (ReadFileTool exists in the parent registry but not the agent's isolated registry)
|
||||
mockModelResponse([
|
||||
{
|
||||
name: ReadFileTool.Name,
|
||||
args: { path: 'config.txt' },
|
||||
id: 'call_blocked',
|
||||
},
|
||||
]);
|
||||
|
||||
// Turn 2: Model stops
|
||||
mockModelResponse([]);
|
||||
// Extraction
|
||||
mockModelResponse([], undefined, 'Done.');
|
||||
|
||||
const consoleWarnSpy = vi
|
||||
.spyOn(console, 'warn')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
await executor.run({ goal: 'Security test' }, signal);
|
||||
|
||||
// Verify executeToolCall was NEVER called because the tool was unauthorized
|
||||
expect(mockExecuteToolCall).not.toHaveBeenCalled();
|
||||
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
||||
expect.stringContaining(
|
||||
`attempted to call unauthorized tool '${ReadFileTool.Name}'`,
|
||||
),
|
||||
);
|
||||
|
||||
consoleWarnSpy.mockRestore();
|
||||
|
||||
// Verify the input to the next model call (Turn 2) indicates failure (as the only call was blocked)
|
||||
const turn2Input = mockSendMessageStream.mock.calls[1][1];
|
||||
const turn2Parts = turn2Input.message as Part[];
|
||||
expect(turn2Parts[0].text).toContain('All tool calls failed');
|
||||
});
|
||||
|
||||
it('should use OutputConfig completion_criteria in the extraction message', async () => {
|
||||
const definition = createTestDefinition(
|
||||
[LSTool.Name],
|
||||
{},
|
||||
{
|
||||
description: 'A summary.',
|
||||
completion_criteria: ['Must include file names', 'Must be concise'],
|
||||
},
|
||||
);
|
||||
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Turn 1: Model stops immediately
|
||||
mockModelResponse([]);
|
||||
|
||||
// Extraction Phase
|
||||
mockModelResponse([], undefined, 'Result: Done.');
|
||||
|
||||
await executor.run({ goal: 'Extraction test' }, signal);
|
||||
|
||||
// Verify the extraction call (the second call)
|
||||
const extractionCallArgs = mockSendMessageStream.mock.calls[1][1];
|
||||
const extractionMessageParts = extractionCallArgs.message as Part[];
|
||||
const extractionText = extractionMessageParts[0].text;
|
||||
|
||||
expect(extractionText).toContain(
|
||||
'Based on your work so far, provide: A summary.',
|
||||
);
|
||||
expect(extractionText).toContain('Be sure you have addressed:');
|
||||
expect(extractionText).toContain('- Must include file names');
|
||||
expect(extractionText).toContain('- Must be concise');
|
||||
});
|
||||
});
|
||||
|
||||
describe('run (Termination Conditions)', () => {
|
||||
const mockKeepAliveResponse = () => {
|
||||
mockModelResponse(
|
||||
[{ name: LSTool.Name, args: { path: '.' }, id: 'loop' }],
|
||||
'Looping',
|
||||
);
|
||||
mockExecuteToolCall.mockResolvedValue({
|
||||
callId: 'loop',
|
||||
resultDisplay: 'ok',
|
||||
responseParts: [
|
||||
{ functionResponse: { name: LSTool.Name, response: {}, id: 'loop' } },
|
||||
],
|
||||
error: undefined,
|
||||
});
|
||||
};
|
||||
|
||||
it('should terminate when max_turns is reached', async () => {
|
||||
const MAX_TURNS = 2;
|
||||
const definition = createTestDefinition([LSTool.Name], {
|
||||
max_turns: MAX_TURNS,
|
||||
});
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Turn 1
|
||||
mockKeepAliveResponse();
|
||||
// Turn 2
|
||||
mockKeepAliveResponse();
|
||||
|
||||
const output = await executor.run({ goal: 'Termination test' }, signal);
|
||||
|
||||
expect(output.terminate_reason).toBe(AgentTerminateMode.MAX_TURNS);
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(MAX_TURNS);
|
||||
// Extraction phase should be skipped when termination is forced
|
||||
expect(mockSendMessageStream).not.toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
expect.any(Object),
|
||||
expect.stringContaining('#extraction'),
|
||||
);
|
||||
});
|
||||
|
||||
it('should terminate if timeout is reached', async () => {
|
||||
const definition = createTestDefinition([LSTool.Name], {
|
||||
max_time_minutes: 5,
|
||||
max_turns: 100,
|
||||
});
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Turn 1 setup
|
||||
mockModelResponse(
|
||||
[{ name: LSTool.Name, args: { path: '.' }, id: 'loop' }],
|
||||
'Looping',
|
||||
);
|
||||
|
||||
// Mock a tool call that takes a long time, causing the overall timeout
|
||||
mockExecuteToolCall.mockImplementation(async () => {
|
||||
// Advance time past the 5-minute limit during the tool call execution
|
||||
await vi.advanceTimersByTimeAsync(5 * 60 * 1000 + 1);
|
||||
return {
|
||||
callId: 'loop',
|
||||
resultDisplay: 'ok',
|
||||
responseParts: [
|
||||
{
|
||||
functionResponse: { name: LSTool.Name, response: {}, id: 'loop' },
|
||||
},
|
||||
],
|
||||
error: undefined,
|
||||
};
|
||||
});
|
||||
|
||||
const output = await executor.run({ goal: 'Termination test' }, signal);
|
||||
|
||||
expect(output.terminate_reason).toBe(AgentTerminateMode.TIMEOUT);
|
||||
// Should only have called the model once before the timeout check stopped it
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should terminate when AbortSignal is triggered mid-stream', async () => {
|
||||
const definition = createTestDefinition();
|
||||
const executor = await AgentExecutor.create(
|
||||
definition,
|
||||
mockConfig,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
// Mock the model response stream
|
||||
mockSendMessageStream.mockImplementation(async () =>
|
||||
(async function* () {
|
||||
// Yield the first chunk
|
||||
yield {
|
||||
type: StreamEventType.CHUNK,
|
||||
value: createMockResponseChunk([
|
||||
{ text: '**Thinking** Step 1', thought: true },
|
||||
]),
|
||||
} as StreamEvent;
|
||||
|
||||
// Simulate abort happening mid-stream
|
||||
abortController.abort();
|
||||
// The loop in callModel should break immediately due to signal check.
|
||||
})(),
|
||||
);
|
||||
|
||||
const output = await executor.run({ goal: 'Termination test' }, signal);
|
||||
expect(output.terminate_reason).toBe(AgentTerminateMode.ABORTED);
|
||||
});
|
||||
});
|
||||
});
|
||||
574
packages/core/src/agents/executor.ts
Normal file
574
packages/core/src/agents/executor.ts
Normal file
@@ -0,0 +1,574 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../config/config.js';
|
||||
import { reportError } from '../utils/errorReporting.js';
|
||||
import { GeminiChat, StreamEventType } from '../core/geminiChat.js';
|
||||
import type {
|
||||
Content,
|
||||
Part,
|
||||
FunctionCall,
|
||||
GenerateContentConfig,
|
||||
FunctionDeclaration,
|
||||
} from '@google/genai';
|
||||
import { executeToolCall } from '../core/nonInteractiveToolExecutor.js';
|
||||
import { ToolRegistry } from '../tools/tool-registry.js';
|
||||
import type { ToolCallRequestInfo } from '../core/turn.js';
|
||||
import { getDirectoryContextString } from '../utils/environmentContext.js';
|
||||
import { GlobTool } from '../tools/glob.js';
|
||||
import { GrepTool } from '../tools/grep.js';
|
||||
import { RipGrepTool } from '../tools/ripGrep.js';
|
||||
import { LSTool } from '../tools/ls.js';
|
||||
import { MemoryTool } from '../tools/memoryTool.js';
|
||||
import { ReadFileTool } from '../tools/read-file.js';
|
||||
import { ReadManyFilesTool } from '../tools/read-many-files.js';
|
||||
import { WebSearchTool } from '../tools/web-search.js';
|
||||
import type {
|
||||
AgentDefinition,
|
||||
AgentInputs,
|
||||
OutputObject,
|
||||
SubagentActivityEvent,
|
||||
} from './types.js';
|
||||
import { AgentTerminateMode } from './types.js';
|
||||
import { templateString } from './utils.js';
|
||||
import { parseThought } from '../utils/thoughtUtils.js';
|
||||
|
||||
/** A callback function to report on agent activity. */
|
||||
export type ActivityCallback = (activity: SubagentActivityEvent) => void;
|
||||
|
||||
/**
|
||||
* Executes an agent loop based on an {@link AgentDefinition}.
|
||||
*
|
||||
* This executor uses a simplified two-phase approach:
|
||||
* 1. **Work Phase:** The agent runs in a loop, calling tools until it has
|
||||
* gathered all necessary information to fulfill its goal.
|
||||
* 2. **Extraction Phase:** A final prompt is sent to the model to summarize
|
||||
* the work and extract the final result in the desired format.
|
||||
*/
|
||||
export class AgentExecutor {
|
||||
readonly definition: AgentDefinition;
|
||||
|
||||
private readonly agentId: string;
|
||||
private readonly toolRegistry: ToolRegistry;
|
||||
private readonly runtimeContext: Config;
|
||||
private readonly onActivity?: ActivityCallback;
|
||||
|
||||
/**
|
||||
* Creates and validates a new `AgentExecutor` instance.
|
||||
*
|
||||
* This method ensures that all tools specified in the agent's definition are
|
||||
* safe for non-interactive use before creating the executor.
|
||||
*
|
||||
* @param definition The definition object for the agent.
|
||||
* @param runtimeContext The global runtime configuration.
|
||||
* @param onActivity An optional callback to receive activity events.
|
||||
* @returns A promise that resolves to a new `AgentExecutor` instance.
|
||||
*/
|
||||
static async create(
|
||||
definition: AgentDefinition,
|
||||
runtimeContext: Config,
|
||||
onActivity?: ActivityCallback,
|
||||
): Promise<AgentExecutor> {
|
||||
// Create an isolated tool registry for this agent instance.
|
||||
const agentToolRegistry = new ToolRegistry(runtimeContext);
|
||||
const parentToolRegistry = await runtimeContext.getToolRegistry();
|
||||
|
||||
if (definition.toolConfig) {
|
||||
for (const toolRef of definition.toolConfig.tools) {
|
||||
if (typeof toolRef === 'string') {
|
||||
// If the tool is referenced by name, retrieve it from the parent
|
||||
// registry and register it with the agent's isolated registry.
|
||||
const toolFromParent = parentToolRegistry.getTool(toolRef);
|
||||
if (toolFromParent) {
|
||||
agentToolRegistry.registerTool(toolFromParent);
|
||||
}
|
||||
} else if (
|
||||
typeof toolRef === 'object' &&
|
||||
'name' in toolRef &&
|
||||
'build' in toolRef
|
||||
) {
|
||||
agentToolRegistry.registerTool(toolRef);
|
||||
}
|
||||
// Note: Raw `FunctionDeclaration` objects in the config don't need to be
|
||||
// registered; their schemas are passed directly to the model later.
|
||||
}
|
||||
|
||||
// Validate that all registered tools are safe for non-interactive
|
||||
// execution.
|
||||
await AgentExecutor.validateTools(agentToolRegistry, definition.name);
|
||||
}
|
||||
|
||||
return new AgentExecutor(
|
||||
definition,
|
||||
runtimeContext,
|
||||
agentToolRegistry,
|
||||
onActivity,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructs a new AgentExecutor instance.
|
||||
*
|
||||
* @private This constructor is private. Use the static `create` method to
|
||||
* instantiate the class.
|
||||
*/
|
||||
private constructor(
|
||||
definition: AgentDefinition,
|
||||
runtimeContext: Config,
|
||||
toolRegistry: ToolRegistry,
|
||||
onActivity?: ActivityCallback,
|
||||
) {
|
||||
this.definition = definition;
|
||||
this.runtimeContext = runtimeContext;
|
||||
this.toolRegistry = toolRegistry;
|
||||
this.onActivity = onActivity;
|
||||
|
||||
const randomIdPart = Math.random().toString(36).slice(2, 8);
|
||||
this.agentId = `${this.definition.name}-${randomIdPart}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Runs the agent.
|
||||
*
|
||||
* @param inputs The validated input parameters for this invocation.
|
||||
* @param signal An `AbortSignal` for cancellation.
|
||||
* @returns A promise that resolves to the agent's final output.
|
||||
*/
|
||||
async run(inputs: AgentInputs, signal: AbortSignal): Promise<OutputObject> {
|
||||
const startTime = Date.now();
|
||||
let turnCounter = 0;
|
||||
|
||||
try {
|
||||
const chat = await this.createChatObject(inputs);
|
||||
const tools = this.prepareToolsList();
|
||||
let terminateReason = AgentTerminateMode.GOAL;
|
||||
|
||||
// Phase 1: Work Phase
|
||||
// The agent works in a loop until it stops calling tools.
|
||||
let currentMessages: Content[] = [
|
||||
{ role: 'user', parts: [{ text: 'Get Started!' }] },
|
||||
];
|
||||
|
||||
while (true) {
|
||||
// Check for termination conditions like max turns or timeout.
|
||||
const reason = this.checkTermination(startTime, turnCounter);
|
||||
if (reason) {
|
||||
terminateReason = reason;
|
||||
break;
|
||||
}
|
||||
if (signal.aborted) {
|
||||
terminateReason = AgentTerminateMode.ABORTED;
|
||||
break;
|
||||
}
|
||||
|
||||
// Call model
|
||||
const promptId = `${this.runtimeContext.getSessionId()}#${this.agentId}#${turnCounter++}`;
|
||||
const { functionCalls } = await this.callModel(
|
||||
chat,
|
||||
currentMessages,
|
||||
tools,
|
||||
signal,
|
||||
promptId,
|
||||
);
|
||||
|
||||
if (signal.aborted) {
|
||||
terminateReason = AgentTerminateMode.ABORTED;
|
||||
break;
|
||||
}
|
||||
|
||||
// If the model stops calling tools, the work phase is complete.
|
||||
if (functionCalls.length === 0) {
|
||||
break;
|
||||
}
|
||||
|
||||
currentMessages = await this.processFunctionCalls(
|
||||
functionCalls,
|
||||
signal,
|
||||
promptId,
|
||||
);
|
||||
}
|
||||
|
||||
// If the work phase was terminated early, skip extraction and return.
|
||||
if (terminateReason !== AgentTerminateMode.GOAL) {
|
||||
return {
|
||||
result: 'Agent execution was terminated before completion.',
|
||||
terminate_reason: terminateReason,
|
||||
};
|
||||
}
|
||||
|
||||
// Phase 2: Extraction Phase
|
||||
// A final message is sent to summarize findings and produce the output.
|
||||
const extractionMessage = this.buildExtractionMessage();
|
||||
const extractionMessages: Content[] = [
|
||||
{ role: 'user', parts: [{ text: extractionMessage }] },
|
||||
];
|
||||
|
||||
const extractionPromptId = `${this.runtimeContext.getSessionId()}#${this.agentId}#extraction`;
|
||||
|
||||
// TODO: Consider if we should keep tools to avoid cache reset.
|
||||
const { textResponse } = await this.callModel(
|
||||
chat,
|
||||
extractionMessages,
|
||||
[], // No tools are available in the extraction phase.
|
||||
signal,
|
||||
extractionPromptId,
|
||||
);
|
||||
|
||||
return {
|
||||
result: textResponse || 'No response generated',
|
||||
terminate_reason: terminateReason,
|
||||
};
|
||||
} catch (error) {
|
||||
this.emitActivity('ERROR', { error: String(error) });
|
||||
throw error; // Re-throw the error for the parent context to handle.
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Calls the generative model with the current context and tools.
|
||||
*
|
||||
* @returns The model's response, including any tool calls or text.
|
||||
*/
|
||||
private async callModel(
|
||||
chat: GeminiChat,
|
||||
messages: Content[],
|
||||
tools: FunctionDeclaration[],
|
||||
signal: AbortSignal,
|
||||
promptId: string,
|
||||
): Promise<{ functionCalls: FunctionCall[]; textResponse: string }> {
|
||||
const messageParams = {
|
||||
message: messages[0]?.parts || [],
|
||||
config: {
|
||||
abortSignal: signal,
|
||||
tools: tools.length > 0 ? [{ functionDeclarations: tools }] : undefined,
|
||||
},
|
||||
};
|
||||
|
||||
const responseStream = await chat.sendMessageStream(
|
||||
this.definition.modelConfig.model,
|
||||
messageParams,
|
||||
promptId,
|
||||
);
|
||||
|
||||
const functionCalls: FunctionCall[] = [];
|
||||
let textResponse = '';
|
||||
|
||||
for await (const resp of responseStream) {
|
||||
if (signal.aborted) break;
|
||||
|
||||
if (resp.type === StreamEventType.CHUNK) {
|
||||
const chunk = resp.value;
|
||||
const parts = chunk.candidates?.[0]?.content?.parts;
|
||||
|
||||
// Extract and emit any subject "thought" content from the model.
|
||||
const { subject } = parseThought(
|
||||
parts?.find((p) => p.thought)?.text || '',
|
||||
);
|
||||
if (subject) {
|
||||
this.emitActivity('THOUGHT_CHUNK', { text: subject });
|
||||
}
|
||||
|
||||
// Collect any function calls requested by the model.
|
||||
if (chunk.functionCalls) {
|
||||
functionCalls.push(...chunk.functionCalls);
|
||||
}
|
||||
|
||||
// Handle text response (non-thought text)
|
||||
const text =
|
||||
parts
|
||||
?.filter((p) => !p.thought && p.text)
|
||||
.map((p) => p.text)
|
||||
.join('') || '';
|
||||
|
||||
if (text) {
|
||||
textResponse += text;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { functionCalls, textResponse };
|
||||
}
|
||||
|
||||
/** Initializes a `GeminiChat` instance for the agent run. */
|
||||
private async createChatObject(inputs: AgentInputs): Promise<GeminiChat> {
|
||||
const { promptConfig, modelConfig } = this.definition;
|
||||
|
||||
if (!promptConfig.systemPrompt && !promptConfig.initialMessages) {
|
||||
throw new Error(
|
||||
'PromptConfig must define either `systemPrompt` or `initialMessages`.',
|
||||
);
|
||||
}
|
||||
|
||||
const startHistory = [...(promptConfig.initialMessages ?? [])];
|
||||
|
||||
// Build system instruction from the templated prompt string.
|
||||
const systemInstruction = promptConfig.systemPrompt
|
||||
? await this.buildSystemPrompt(inputs)
|
||||
: undefined;
|
||||
|
||||
try {
|
||||
const generationConfig: GenerateContentConfig = {
|
||||
temperature: modelConfig.temp,
|
||||
topP: modelConfig.top_p,
|
||||
thinkingConfig: {
|
||||
includeThoughts: true,
|
||||
thinkingBudget: modelConfig.thinkingBudget ?? -1,
|
||||
},
|
||||
};
|
||||
|
||||
if (systemInstruction) {
|
||||
generationConfig.systemInstruction = systemInstruction;
|
||||
}
|
||||
|
||||
return new GeminiChat(
|
||||
this.runtimeContext,
|
||||
generationConfig,
|
||||
startHistory,
|
||||
);
|
||||
} catch (error) {
|
||||
await reportError(
|
||||
error,
|
||||
`Error initializing Gemini chat for agent ${this.definition.name}.`,
|
||||
startHistory,
|
||||
'startChat',
|
||||
);
|
||||
// Re-throw as a more specific error after reporting.
|
||||
throw new Error(`Failed to create chat object: ${error}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes function calls requested by the model and returns the results.
|
||||
*
|
||||
* @returns A new `Content` object to be added to the chat history.
|
||||
*/
|
||||
private async processFunctionCalls(
|
||||
functionCalls: FunctionCall[],
|
||||
signal: AbortSignal,
|
||||
promptId: string,
|
||||
): Promise<Content[]> {
|
||||
const allowedToolNames = new Set(this.toolRegistry.getAllToolNames());
|
||||
|
||||
// Filter out any tool calls that are not in the agent's allowed list.
|
||||
const validatedFunctionCalls = functionCalls.filter((call) => {
|
||||
if (!allowedToolNames.has(call.name as string)) {
|
||||
console.warn(
|
||||
`[AgentExecutor] Agent '${this.definition.name}' attempted to call ` +
|
||||
`unauthorized tool '${call.name}'. This call has been blocked.`,
|
||||
);
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
});
|
||||
|
||||
const toolPromises = validatedFunctionCalls.map(async (functionCall) => {
|
||||
const callId = functionCall.id ?? `${functionCall.name}-${Date.now()}`;
|
||||
const args = functionCall.args ?? {};
|
||||
|
||||
this.emitActivity('TOOL_CALL_START', {
|
||||
name: functionCall.name,
|
||||
args,
|
||||
});
|
||||
|
||||
const requestInfo: ToolCallRequestInfo = {
|
||||
callId,
|
||||
name: functionCall.name as string,
|
||||
args: args as Record<string, unknown>,
|
||||
isClientInitiated: true,
|
||||
prompt_id: promptId,
|
||||
};
|
||||
|
||||
const toolResponse = await executeToolCall(
|
||||
this.runtimeContext,
|
||||
requestInfo,
|
||||
signal,
|
||||
);
|
||||
|
||||
if (toolResponse.error) {
|
||||
this.emitActivity('ERROR', {
|
||||
context: 'tool_call',
|
||||
name: functionCall.name,
|
||||
error: toolResponse.error.message,
|
||||
});
|
||||
} else {
|
||||
this.emitActivity('TOOL_CALL_END', {
|
||||
name: functionCall.name,
|
||||
output: toolResponse.resultDisplay,
|
||||
});
|
||||
}
|
||||
|
||||
return toolResponse;
|
||||
});
|
||||
|
||||
const toolResponses = await Promise.all(toolPromises);
|
||||
const toolResponseParts: Part[] = toolResponses
|
||||
.flatMap((response) => response.responseParts)
|
||||
.filter((part): part is Part => part !== undefined);
|
||||
|
||||
// If all authorized tool calls failed, provide a generic error message
|
||||
// to the model so it can try a different approach.
|
||||
if (functionCalls.length > 0 && toolResponseParts.length === 0) {
|
||||
toolResponseParts.push({
|
||||
text: 'All tool calls failed. Please analyze the errors and try an alternative approach.',
|
||||
});
|
||||
}
|
||||
|
||||
return [{ role: 'user', parts: toolResponseParts }];
|
||||
}
|
||||
|
||||
/**
|
||||
* Prepares the list of tool function declarations to be sent to the model.
|
||||
*/
|
||||
private prepareToolsList(): FunctionDeclaration[] {
|
||||
const toolsList: FunctionDeclaration[] = [];
|
||||
const { toolConfig } = this.definition;
|
||||
|
||||
if (toolConfig) {
|
||||
const toolNamesToLoad: string[] = [];
|
||||
for (const toolRef of toolConfig.tools) {
|
||||
if (typeof toolRef === 'string') {
|
||||
toolNamesToLoad.push(toolRef);
|
||||
} else if (typeof toolRef === 'object' && 'schema' in toolRef) {
|
||||
// Tool instance with an explicit schema property.
|
||||
toolsList.push(toolRef.schema as FunctionDeclaration);
|
||||
} else {
|
||||
// Raw `FunctionDeclaration` object.
|
||||
toolsList.push(toolRef as FunctionDeclaration);
|
||||
}
|
||||
}
|
||||
// Add schemas from tools that were registered by name.
|
||||
toolsList.push(
|
||||
...this.toolRegistry.getFunctionDeclarationsFiltered(toolNamesToLoad),
|
||||
);
|
||||
}
|
||||
|
||||
return toolsList;
|
||||
}
|
||||
|
||||
/** Builds the system prompt from the agent definition and inputs. */
|
||||
private async buildSystemPrompt(inputs: AgentInputs): Promise<string> {
|
||||
const { promptConfig, outputConfig } = this.definition;
|
||||
if (!promptConfig.systemPrompt) {
|
||||
return '';
|
||||
}
|
||||
|
||||
// Inject user inputs into the prompt template.
|
||||
let finalPrompt = templateString(promptConfig.systemPrompt, inputs);
|
||||
|
||||
// Append environment context (CWD and folder structure).
|
||||
const dirContext = await getDirectoryContextString(this.runtimeContext);
|
||||
finalPrompt += `\n\n# Environment Context\n${dirContext}`;
|
||||
|
||||
// Append completion criteria to guide the model's output.
|
||||
if (outputConfig?.completion_criteria) {
|
||||
finalPrompt += '\n\nEnsure you complete the following:\n';
|
||||
for (const criteria of outputConfig.completion_criteria) {
|
||||
finalPrompt += `- ${criteria}\n`;
|
||||
}
|
||||
}
|
||||
|
||||
// Append standard rules for non-interactive execution.
|
||||
finalPrompt += `
|
||||
Important Rules:
|
||||
* You are running in a non-interactive mode. You CANNOT ask the user for input or clarification.
|
||||
* Work systematically using available tools to complete your task.
|
||||
* Always use absolute paths for file operations. Construct them using the provided "Environment Context".
|
||||
* When you have completed your analysis and are ready to produce the final answer, stop calling tools.`;
|
||||
|
||||
return finalPrompt;
|
||||
}
|
||||
|
||||
/** Builds the final message for the extraction phase. */
|
||||
private buildExtractionMessage(): string {
|
||||
const { outputConfig } = this.definition;
|
||||
|
||||
if (outputConfig?.description) {
|
||||
let message = `Based on your work so far, provide: ${outputConfig.description}`;
|
||||
|
||||
if (outputConfig.completion_criteria?.length) {
|
||||
message += `\n\nBe sure you have addressed:\n`;
|
||||
for (const criteria of outputConfig.completion_criteria) {
|
||||
message += `- ${criteria}\n`;
|
||||
}
|
||||
}
|
||||
return message;
|
||||
}
|
||||
|
||||
// Fallback to a generic extraction message if no description is provided.
|
||||
return 'Based on your work so far, provide a comprehensive summary of your analysis and findings. Do not perform any more function calls.';
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates that all tools in a registry are safe for non-interactive use.
|
||||
*
|
||||
* @throws An error if a tool is not on the allow-list for non-interactive execution.
|
||||
*/
|
||||
private static async validateTools(
|
||||
toolRegistry: ToolRegistry,
|
||||
agentName: string,
|
||||
): Promise<void> {
|
||||
// Tools that are non-interactive. This is temporary until we have tool
|
||||
// confirmations for subagents.
|
||||
const allowlist = new Set([
|
||||
LSTool.Name,
|
||||
ReadFileTool.Name,
|
||||
GrepTool.Name,
|
||||
RipGrepTool.Name,
|
||||
GlobTool.Name,
|
||||
ReadManyFilesTool.Name,
|
||||
MemoryTool.Name,
|
||||
WebSearchTool.Name,
|
||||
]);
|
||||
for (const tool of toolRegistry.getAllTools()) {
|
||||
if (!allowlist.has(tool.name)) {
|
||||
throw new Error(
|
||||
`Tool "${tool.name}" is not on the allow-list for non-interactive ` +
|
||||
`execution in agent "${agentName}". Only tools that do not require user ` +
|
||||
`confirmation can be used in subagents.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if the agent should terminate due to exceeding configured limits.
|
||||
*
|
||||
* @returns The reason for termination, or `null` if execution can continue.
|
||||
*/
|
||||
private checkTermination(
|
||||
startTime: number,
|
||||
turnCounter: number,
|
||||
): AgentTerminateMode | null {
|
||||
const { runConfig } = this.definition;
|
||||
|
||||
if (runConfig.max_turns && turnCounter >= runConfig.max_turns) {
|
||||
return AgentTerminateMode.MAX_TURNS;
|
||||
}
|
||||
|
||||
const elapsedMinutes = (Date.now() - startTime) / (1000 * 60);
|
||||
if (elapsedMinutes >= runConfig.max_time_minutes) {
|
||||
return AgentTerminateMode.TIMEOUT;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Emits an activity event to the configured callback. */
|
||||
private emitActivity(
|
||||
type: SubagentActivityEvent['type'],
|
||||
data: Record<string, unknown>,
|
||||
): void {
|
||||
if (this.onActivity) {
|
||||
const event: SubagentActivityEvent = {
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: this.definition.name,
|
||||
type,
|
||||
data,
|
||||
};
|
||||
this.onActivity(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
285
packages/core/src/agents/invocation.test.ts
Normal file
285
packages/core/src/agents/invocation.test.ts
Normal file
@@ -0,0 +1,285 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, type Mocked } from 'vitest';
|
||||
import { SubagentInvocation } from './invocation.js';
|
||||
import { AgentExecutor } from './executor.js';
|
||||
import type {
|
||||
AgentDefinition,
|
||||
SubagentActivityEvent,
|
||||
AgentInputs,
|
||||
} from './types.js';
|
||||
import { AgentTerminateMode } from './types.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import { ToolErrorType } from '../tools/tool-error.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
|
||||
vi.mock('./executor.js');
|
||||
|
||||
const MockAgentExecutor = vi.mocked(AgentExecutor);
|
||||
|
||||
let mockConfig: Config;
|
||||
|
||||
const testDefinition: AgentDefinition = {
|
||||
name: 'MockAgent',
|
||||
description: 'A mock agent.',
|
||||
inputConfig: {
|
||||
inputs: {
|
||||
task: { type: 'string', required: true, description: 'task' },
|
||||
priority: { type: 'number', required: false, description: 'prio' },
|
||||
},
|
||||
},
|
||||
modelConfig: { model: 'test', temp: 0, top_p: 1 },
|
||||
runConfig: { max_time_minutes: 1 },
|
||||
promptConfig: { systemPrompt: 'test' },
|
||||
};
|
||||
|
||||
describe('SubagentInvocation', () => {
|
||||
let mockExecutorInstance: Mocked<AgentExecutor>;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
|
||||
mockExecutorInstance = {
|
||||
run: vi.fn(),
|
||||
definition: testDefinition,
|
||||
} as unknown as Mocked<AgentExecutor>;
|
||||
|
||||
MockAgentExecutor.create.mockResolvedValue(mockExecutorInstance);
|
||||
});
|
||||
|
||||
describe('getDescription', () => {
|
||||
it('should format the description with inputs', () => {
|
||||
const params = { task: 'Analyze data', priority: 5 };
|
||||
const invocation = new SubagentInvocation(
|
||||
params,
|
||||
testDefinition,
|
||||
mockConfig,
|
||||
);
|
||||
const description = invocation.getDescription();
|
||||
expect(description).toBe(
|
||||
"Running subagent 'MockAgent' with inputs: { task: Analyze data, priority: 5 }",
|
||||
);
|
||||
});
|
||||
|
||||
it('should truncate long input values', () => {
|
||||
const longTask = 'A'.repeat(100);
|
||||
const params = { task: longTask };
|
||||
const invocation = new SubagentInvocation(
|
||||
params,
|
||||
testDefinition,
|
||||
mockConfig,
|
||||
);
|
||||
const description = invocation.getDescription();
|
||||
// Default INPUT_PREVIEW_MAX_LENGTH is 50
|
||||
expect(description).toBe(
|
||||
`Running subagent 'MockAgent' with inputs: { task: ${'A'.repeat(50)} }`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should truncate the overall description if it exceeds the limit', () => {
|
||||
// Create a definition and inputs that result in a very long description
|
||||
const longNameDef = {
|
||||
...testDefinition,
|
||||
name: 'VeryLongAgentNameThatTakesUpSpace',
|
||||
};
|
||||
const params: AgentInputs = {};
|
||||
for (let i = 0; i < 20; i++) {
|
||||
params[`input${i}`] = `value${i}`;
|
||||
}
|
||||
const invocation = new SubagentInvocation(
|
||||
params,
|
||||
longNameDef,
|
||||
mockConfig,
|
||||
);
|
||||
const description = invocation.getDescription();
|
||||
// Default DESCRIPTION_MAX_LENGTH is 200
|
||||
expect(description.length).toBe(200);
|
||||
expect(
|
||||
description.startsWith(
|
||||
"Running subagent 'VeryLongAgentNameThatTakesUpSpace'",
|
||||
),
|
||||
).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('execute', () => {
|
||||
let signal: AbortSignal;
|
||||
let updateOutput: ReturnType<typeof vi.fn>;
|
||||
const params = { task: 'Execute task' };
|
||||
let invocation: SubagentInvocation;
|
||||
|
||||
beforeEach(() => {
|
||||
signal = new AbortController().signal;
|
||||
updateOutput = vi.fn();
|
||||
invocation = new SubagentInvocation(params, testDefinition, mockConfig);
|
||||
});
|
||||
|
||||
it('should initialize and run the executor successfully', async () => {
|
||||
const mockOutput = {
|
||||
result: 'Analysis complete.',
|
||||
terminate_reason: AgentTerminateMode.GOAL,
|
||||
};
|
||||
mockExecutorInstance.run.mockResolvedValue(mockOutput);
|
||||
|
||||
const result = await invocation.execute(signal, updateOutput);
|
||||
|
||||
expect(MockAgentExecutor.create).toHaveBeenCalledWith(
|
||||
testDefinition,
|
||||
mockConfig,
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
|
||||
|
||||
expect(mockExecutorInstance.run).toHaveBeenCalledWith(params, signal);
|
||||
|
||||
expect(result.llmContent).toEqual([
|
||||
{
|
||||
text: expect.stringContaining(
|
||||
"Subagent 'MockAgent' finished.\nTermination Reason: GOAL\nResult:\nAnalysis complete.",
|
||||
),
|
||||
},
|
||||
]);
|
||||
expect(result.returnDisplay).toContain('Result:\nAnalysis complete.');
|
||||
expect(result.returnDisplay).toContain('Termination Reason:\n GOAL');
|
||||
});
|
||||
|
||||
it('should stream THOUGHT_CHUNK activities from the executor', async () => {
|
||||
mockExecutorInstance.run.mockImplementation(async () => {
|
||||
const onActivity = MockAgentExecutor.create.mock.calls[0][2];
|
||||
|
||||
if (onActivity) {
|
||||
onActivity({
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: 'MockAgent',
|
||||
type: 'THOUGHT_CHUNK',
|
||||
data: { text: 'Analyzing...' },
|
||||
} as SubagentActivityEvent);
|
||||
onActivity({
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: 'MockAgent',
|
||||
type: 'THOUGHT_CHUNK',
|
||||
data: { text: ' Still thinking.' },
|
||||
} as SubagentActivityEvent);
|
||||
}
|
||||
return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
|
||||
});
|
||||
|
||||
await invocation.execute(signal, updateOutput);
|
||||
|
||||
expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
|
||||
expect(updateOutput).toHaveBeenCalledWith('🤖💭 Analyzing...');
|
||||
expect(updateOutput).toHaveBeenCalledWith('🤖💭 Still thinking.');
|
||||
expect(updateOutput).toHaveBeenCalledTimes(3); // Initial message + 2 thoughts
|
||||
});
|
||||
|
||||
it('should NOT stream other activities (e.g., TOOL_CALL_START, ERROR)', async () => {
|
||||
mockExecutorInstance.run.mockImplementation(async () => {
|
||||
const onActivity = MockAgentExecutor.create.mock.calls[0][2];
|
||||
|
||||
if (onActivity) {
|
||||
onActivity({
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: 'MockAgent',
|
||||
type: 'TOOL_CALL_START',
|
||||
data: { name: 'ls' },
|
||||
} as SubagentActivityEvent);
|
||||
onActivity({
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: 'MockAgent',
|
||||
type: 'ERROR',
|
||||
data: { error: 'Failed' },
|
||||
} as SubagentActivityEvent);
|
||||
}
|
||||
return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
|
||||
});
|
||||
|
||||
await invocation.execute(signal, updateOutput);
|
||||
|
||||
// Should only contain the initial "Subagent starting..." message
|
||||
expect(updateOutput).toHaveBeenCalledTimes(1);
|
||||
expect(updateOutput).toHaveBeenCalledWith('Subagent starting...\n');
|
||||
});
|
||||
|
||||
it('should run successfully without an updateOutput callback', async () => {
|
||||
mockExecutorInstance.run.mockImplementation(async () => {
|
||||
const onActivity = MockAgentExecutor.create.mock.calls[0][2];
|
||||
if (onActivity) {
|
||||
// Ensure calling activity doesn't crash when updateOutput is undefined
|
||||
onActivity({
|
||||
isSubagentActivityEvent: true,
|
||||
agentName: 'testAgent',
|
||||
type: 'THOUGHT_CHUNK',
|
||||
data: { text: 'Thinking silently.' },
|
||||
} as SubagentActivityEvent);
|
||||
}
|
||||
return { result: 'Done', terminate_reason: AgentTerminateMode.GOAL };
|
||||
});
|
||||
|
||||
// Execute without the optional callback
|
||||
const result = await invocation.execute(signal);
|
||||
expect(result.error).toBeUndefined();
|
||||
expect(result.returnDisplay).toContain('Result:\nDone');
|
||||
});
|
||||
|
||||
it('should handle executor run failure', async () => {
|
||||
const error = new Error('Model failed during execution.');
|
||||
mockExecutorInstance.run.mockRejectedValue(error);
|
||||
|
||||
const result = await invocation.execute(signal, updateOutput);
|
||||
|
||||
expect(result.error).toEqual({
|
||||
message: error.message,
|
||||
type: ToolErrorType.EXECUTION_FAILED,
|
||||
});
|
||||
expect(result.returnDisplay).toBe(
|
||||
`Subagent Failed: MockAgent\nError: ${error.message}`,
|
||||
);
|
||||
expect(result.llmContent).toBe(
|
||||
`Subagent 'MockAgent' failed. Error: ${error.message}`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle executor creation failure', async () => {
|
||||
const creationError = new Error('Failed to initialize tools.');
|
||||
MockAgentExecutor.create.mockRejectedValue(creationError);
|
||||
|
||||
const result = await invocation.execute(signal, updateOutput);
|
||||
|
||||
expect(mockExecutorInstance.run).not.toHaveBeenCalled();
|
||||
expect(result.error).toEqual({
|
||||
message: creationError.message,
|
||||
type: ToolErrorType.EXECUTION_FAILED,
|
||||
});
|
||||
expect(result.returnDisplay).toContain(`Error: ${creationError.message}`);
|
||||
});
|
||||
|
||||
/**
|
||||
* This test verifies that the AbortSignal is correctly propagated and
|
||||
* that a rejection from the executor due to abortion is handled gracefully.
|
||||
*/
|
||||
it('should handle abortion signal during execution', async () => {
|
||||
const abortError = new Error('Aborted');
|
||||
mockExecutorInstance.run.mockRejectedValue(abortError);
|
||||
|
||||
const controller = new AbortController();
|
||||
const executePromise = invocation.execute(
|
||||
controller.signal,
|
||||
updateOutput,
|
||||
);
|
||||
controller.abort();
|
||||
const result = await executePromise;
|
||||
|
||||
expect(mockExecutorInstance.run).toHaveBeenCalledWith(
|
||||
params,
|
||||
controller.signal,
|
||||
);
|
||||
expect(result.error?.message).toBe('Aborted');
|
||||
expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED);
|
||||
});
|
||||
});
|
||||
});
|
||||
134
packages/core/src/agents/invocation.ts
Normal file
134
packages/core/src/agents/invocation.ts
Normal file
@@ -0,0 +1,134 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../config/config.js';
|
||||
import { AgentExecutor } from './executor.js';
|
||||
import type { AnsiOutput } from '../utils/terminalSerializer.js';
|
||||
import { BaseToolInvocation, type ToolResult } from '../tools/tools.js';
|
||||
import { ToolErrorType } from '../tools/tool-error.js';
|
||||
import type {
|
||||
AgentDefinition,
|
||||
AgentInputs,
|
||||
SubagentActivityEvent,
|
||||
} from './types.js';
|
||||
|
||||
const INPUT_PREVIEW_MAX_LENGTH = 50;
|
||||
const DESCRIPTION_MAX_LENGTH = 200;
|
||||
|
||||
/**
|
||||
* Represents a validated, executable instance of a subagent tool.
|
||||
*
|
||||
* This class orchestrates the execution of a defined agent by:
|
||||
* 1. Initializing the {@link AgentExecutor}.
|
||||
* 2. Running the agent's execution loop.
|
||||
* 3. Bridging the agent's streaming activity (e.g., thoughts) to the tool's
|
||||
* live output stream.
|
||||
* 4. Formatting the final result into a {@link ToolResult}.
|
||||
*/
|
||||
export class SubagentInvocation extends BaseToolInvocation<
|
||||
AgentInputs,
|
||||
ToolResult
|
||||
> {
|
||||
/**
|
||||
* @param params The validated input parameters for the agent.
|
||||
* @param definition The definition object that configures the agent.
|
||||
* @param config The global runtime configuration.
|
||||
*/
|
||||
constructor(
|
||||
params: AgentInputs,
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly config: Config,
|
||||
) {
|
||||
super(params);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a concise, human-readable description of the invocation.
|
||||
* Used for logging and display purposes.
|
||||
*/
|
||||
getDescription(): string {
|
||||
const inputSummary = Object.entries(this.params)
|
||||
.map(
|
||||
([key, value]) =>
|
||||
`${key}: ${String(value).slice(0, INPUT_PREVIEW_MAX_LENGTH)}`,
|
||||
)
|
||||
.join(', ');
|
||||
|
||||
const description = `Running subagent '${this.definition.name}' with inputs: { ${inputSummary} }`;
|
||||
return description.slice(0, DESCRIPTION_MAX_LENGTH);
|
||||
}
|
||||
|
||||
/**
|
||||
* Executes the subagent.
|
||||
*
|
||||
* @param signal An `AbortSignal` to cancel the agent's execution.
|
||||
* @param updateOutput A callback to stream intermediate output, such as the
|
||||
* agent's thoughts, to the user interface.
|
||||
* @returns A `Promise` that resolves with the final `ToolResult`.
|
||||
*/
|
||||
async execute(
|
||||
signal: AbortSignal,
|
||||
updateOutput?: (output: string | AnsiOutput) => void,
|
||||
): Promise<ToolResult> {
|
||||
try {
|
||||
if (updateOutput) {
|
||||
updateOutput('Subagent starting...\n');
|
||||
}
|
||||
|
||||
// Create an activity callback to bridge the executor's events to the
|
||||
// tool's streaming output.
|
||||
const onActivity = (activity: SubagentActivityEvent): void => {
|
||||
if (!updateOutput) return;
|
||||
|
||||
if (
|
||||
activity.type === 'THOUGHT_CHUNK' &&
|
||||
typeof activity.data['text'] === 'string'
|
||||
) {
|
||||
updateOutput(`🤖💭 ${activity.data['text']}`);
|
||||
}
|
||||
};
|
||||
|
||||
const executor = await AgentExecutor.create(
|
||||
this.definition,
|
||||
this.config,
|
||||
onActivity,
|
||||
);
|
||||
|
||||
const output = await executor.run(this.params, signal);
|
||||
|
||||
const resultContent = `Subagent '${this.definition.name}' finished.
|
||||
Termination Reason: ${output.terminate_reason}
|
||||
Result:
|
||||
${output.result}`;
|
||||
|
||||
const displayContent = `
|
||||
Subagent ${this.definition.name} Finished
|
||||
|
||||
Termination Reason:\n ${output.terminate_reason}
|
||||
|
||||
Result:
|
||||
${output.result}
|
||||
`;
|
||||
|
||||
return {
|
||||
llmContent: [{ text: resultContent }],
|
||||
returnDisplay: displayContent,
|
||||
};
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : String(error);
|
||||
|
||||
return {
|
||||
llmContent: `Subagent '${this.definition.name}' failed. Error: ${errorMessage}`,
|
||||
returnDisplay: `Subagent Failed: ${this.definition.name}\nError: ${errorMessage}`,
|
||||
error: {
|
||||
message: errorMessage,
|
||||
type: ToolErrorType.EXECUTION_FAILED,
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
205
packages/core/src/agents/registry.test.ts
Normal file
205
packages/core/src/agents/registry.test.ts
Normal file
@@ -0,0 +1,205 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { AgentRegistry } from './registry.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import type { AgentDefinition } from './types.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
|
||||
// A test-only subclass to expose the protected `registerAgent` method.
|
||||
class TestableAgentRegistry extends AgentRegistry {
|
||||
testRegisterAgent(definition: AgentDefinition): void {
|
||||
this.registerAgent(definition);
|
||||
}
|
||||
}
|
||||
|
||||
// Define mock agent structures for testing registration logic
|
||||
const MOCK_AGENT_V1: AgentDefinition = {
|
||||
name: 'MockAgent',
|
||||
description: 'Mock Description V1',
|
||||
inputConfig: { inputs: {} },
|
||||
modelConfig: { model: 'test', temp: 0, top_p: 1 },
|
||||
runConfig: { max_time_minutes: 1 },
|
||||
promptConfig: { systemPrompt: 'test' },
|
||||
};
|
||||
|
||||
const MOCK_AGENT_V2: AgentDefinition = {
|
||||
...MOCK_AGENT_V1,
|
||||
description: 'Mock Description V2 (Updated)',
|
||||
};
|
||||
|
||||
describe('AgentRegistry', () => {
|
||||
let mockConfig: Config;
|
||||
let registry: TestableAgentRegistry;
|
||||
|
||||
beforeEach(() => {
|
||||
// Default configuration (debugMode: false)
|
||||
mockConfig = makeFakeConfig();
|
||||
registry = new TestableAgentRegistry(mockConfig);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks(); // Restore spies after each test
|
||||
});
|
||||
|
||||
describe('initialize', () => {
|
||||
// TODO: Add this test once we actually have a built-in agent configured.
|
||||
// it('should load built-in agents upon initialization', async () => {
|
||||
// expect(registry.getAllDefinitions()).toHaveLength(0);
|
||||
|
||||
// await registry.initialize();
|
||||
|
||||
// // There are currently no built-in agents.
|
||||
// expect(registry.getAllDefinitions()).toEqual([]);
|
||||
// });
|
||||
|
||||
it('should log the count of loaded agents in debug mode', async () => {
|
||||
const debugConfig = makeFakeConfig({ debugMode: true });
|
||||
const debugRegistry = new TestableAgentRegistry(debugConfig);
|
||||
const consoleLogSpy = vi
|
||||
.spyOn(console, 'log')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
await debugRegistry.initialize();
|
||||
|
||||
const agentCount = debugRegistry.getAllDefinitions().length;
|
||||
expect(consoleLogSpy).toHaveBeenCalledWith(
|
||||
`[AgentRegistry] Initialized with ${agentCount} agents.`,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('registration logic', () => {
|
||||
it('should register a valid agent definition', () => {
|
||||
registry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
expect(registry.getDefinition('MockAgent')).toEqual(MOCK_AGENT_V1);
|
||||
});
|
||||
|
||||
it('should handle special characters in agent names', () => {
|
||||
const specialAgent = {
|
||||
...MOCK_AGENT_V1,
|
||||
name: 'Agent-123_$pecial.v2',
|
||||
};
|
||||
registry.testRegisterAgent(specialAgent);
|
||||
expect(registry.getDefinition('Agent-123_$pecial.v2')).toEqual(
|
||||
specialAgent,
|
||||
);
|
||||
});
|
||||
|
||||
it('should reject an agent definition missing a name', () => {
|
||||
const invalidAgent = { ...MOCK_AGENT_V1, name: '' };
|
||||
const consoleWarnSpy = vi
|
||||
.spyOn(console, 'warn')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
registry.testRegisterAgent(invalidAgent);
|
||||
|
||||
expect(registry.getDefinition('MockAgent')).toBeUndefined();
|
||||
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
||||
'[AgentRegistry] Skipping invalid agent definition. Missing name or description.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should reject an agent definition missing a description', () => {
|
||||
const invalidAgent = { ...MOCK_AGENT_V1, description: '' };
|
||||
const consoleWarnSpy = vi
|
||||
.spyOn(console, 'warn')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
registry.testRegisterAgent(invalidAgent as AgentDefinition);
|
||||
|
||||
expect(registry.getDefinition('MockAgent')).toBeUndefined();
|
||||
expect(consoleWarnSpy).toHaveBeenCalledWith(
|
||||
'[AgentRegistry] Skipping invalid agent definition. Missing name or description.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should overwrite an existing agent definition', () => {
|
||||
registry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
expect(registry.getDefinition('MockAgent')?.description).toBe(
|
||||
'Mock Description V1',
|
||||
);
|
||||
|
||||
registry.testRegisterAgent(MOCK_AGENT_V2);
|
||||
expect(registry.getDefinition('MockAgent')?.description).toBe(
|
||||
'Mock Description V2 (Updated)',
|
||||
);
|
||||
expect(registry.getAllDefinitions()).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('should log overwrites when in debug mode', () => {
|
||||
const debugConfig = makeFakeConfig({ debugMode: true });
|
||||
const debugRegistry = new TestableAgentRegistry(debugConfig);
|
||||
const consoleLogSpy = vi
|
||||
.spyOn(console, 'log')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
debugRegistry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
debugRegistry.testRegisterAgent(MOCK_AGENT_V2);
|
||||
|
||||
expect(consoleLogSpy).toHaveBeenCalledWith(
|
||||
`[AgentRegistry] Overriding agent 'MockAgent'`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should not log overwrites when not in debug mode', () => {
|
||||
const consoleLogSpy = vi
|
||||
.spyOn(console, 'log')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
registry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
registry.testRegisterAgent(MOCK_AGENT_V2);
|
||||
|
||||
expect(consoleLogSpy).not.toHaveBeenCalledWith(
|
||||
`[AgentRegistry] Overriding agent 'MockAgent'`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle bulk registrations correctly', async () => {
|
||||
const promises = Array.from({ length: 100 }, (_, i) =>
|
||||
Promise.resolve(
|
||||
registry.testRegisterAgent({
|
||||
...MOCK_AGENT_V1,
|
||||
name: `Agent${i}`,
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
await Promise.all(promises);
|
||||
expect(registry.getAllDefinitions()).toHaveLength(100);
|
||||
});
|
||||
});
|
||||
|
||||
describe('accessors', () => {
|
||||
const ANOTHER_AGENT: AgentDefinition = {
|
||||
...MOCK_AGENT_V1,
|
||||
name: 'AnotherAgent',
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
registry.testRegisterAgent(MOCK_AGENT_V1);
|
||||
registry.testRegisterAgent(ANOTHER_AGENT);
|
||||
});
|
||||
|
||||
it('getDefinition should return the correct definition', () => {
|
||||
expect(registry.getDefinition('MockAgent')).toEqual(MOCK_AGENT_V1);
|
||||
expect(registry.getDefinition('AnotherAgent')).toEqual(ANOTHER_AGENT);
|
||||
});
|
||||
|
||||
it('getDefinition should return undefined for unknown agents', () => {
|
||||
expect(registry.getDefinition('NonExistentAgent')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('getAllDefinitions should return all registered definitions', () => {
|
||||
const all = registry.getAllDefinitions();
|
||||
expect(all).toHaveLength(2);
|
||||
expect(all).toEqual(
|
||||
expect.arrayContaining([MOCK_AGENT_V1, ANOTHER_AGENT]),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
71
packages/core/src/agents/registry.ts
Normal file
71
packages/core/src/agents/registry.ts
Normal file
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { AgentDefinition } from './types.js';
|
||||
import { CodebaseInvestigatorAgent } from './codebase-investigator.js';
|
||||
|
||||
/**
|
||||
* Manages the discovery, loading, validation, and registration of
|
||||
* AgentDefinitions.
|
||||
*/
|
||||
export class AgentRegistry {
|
||||
private readonly agents = new Map<string, AgentDefinition>();
|
||||
|
||||
constructor(private readonly config: Config) {}
|
||||
|
||||
/**
|
||||
* Discovers and loads agents.
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
this.loadBuiltInAgents();
|
||||
|
||||
if (this.config.getDebugMode()) {
|
||||
console.log(
|
||||
`[AgentRegistry] Initialized with ${this.agents.size} agents.`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private loadBuiltInAgents(): void {
|
||||
this.registerAgent(CodebaseInvestigatorAgent);
|
||||
}
|
||||
|
||||
/**
|
||||
* Registers an agent definition. If an agent with the same name exists,
|
||||
* it will be overwritten, respecting the precedence established by the
|
||||
* initialization order.
|
||||
*/
|
||||
protected registerAgent(definition: AgentDefinition): void {
|
||||
// Basic validation
|
||||
if (!definition.name || !definition.description) {
|
||||
console.warn(
|
||||
`[AgentRegistry] Skipping invalid agent definition. Missing name or description.`,
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.agents.has(definition.name) && this.config.getDebugMode()) {
|
||||
console.log(`[AgentRegistry] Overriding agent '${definition.name}'`);
|
||||
}
|
||||
|
||||
this.agents.set(definition.name, definition);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves an agent definition by name.
|
||||
*/
|
||||
getDefinition(name: string): AgentDefinition | undefined {
|
||||
return this.agents.get(name);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns all active agent definitions.
|
||||
*/
|
||||
getAllDefinitions(): AgentDefinition[] {
|
||||
return Array.from(this.agents.values());
|
||||
}
|
||||
}
|
||||
165
packages/core/src/agents/schema-utils.test.ts
Normal file
165
packages/core/src/agents/schema-utils.test.ts
Normal file
@@ -0,0 +1,165 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { convertInputConfigToJsonSchema } from './schema-utils.js';
|
||||
import type { InputConfig } from './types.js';
|
||||
|
||||
const PRIMITIVE_TYPES_CONFIG: InputConfig = {
|
||||
inputs: {
|
||||
goal: {
|
||||
type: 'string',
|
||||
description: 'The primary objective',
|
||||
required: true,
|
||||
},
|
||||
max_retries: {
|
||||
type: 'integer',
|
||||
description: 'Maximum number of retries',
|
||||
required: false,
|
||||
},
|
||||
temperature: {
|
||||
type: 'number',
|
||||
description: 'The model temperature',
|
||||
required: true,
|
||||
},
|
||||
verbose: {
|
||||
type: 'boolean',
|
||||
description: 'Enable verbose logging',
|
||||
required: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const ARRAY_TYPES_CONFIG: InputConfig = {
|
||||
inputs: {
|
||||
filenames: {
|
||||
type: 'string[]',
|
||||
description: 'A list of file paths',
|
||||
required: true,
|
||||
},
|
||||
scores: {
|
||||
type: 'number[]',
|
||||
description: 'A list of scores',
|
||||
required: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const NO_REQUIRED_FIELDS_CONFIG: InputConfig = {
|
||||
inputs: {
|
||||
optional_param: {
|
||||
type: 'string',
|
||||
description: 'An optional parameter',
|
||||
required: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const ALL_REQUIRED_FIELDS_CONFIG: InputConfig = {
|
||||
inputs: {
|
||||
paramA: { type: 'string', description: 'Parameter A', required: true },
|
||||
paramB: { type: 'boolean', description: 'Parameter B', required: true },
|
||||
},
|
||||
};
|
||||
|
||||
const EMPTY_CONFIG: InputConfig = {
|
||||
inputs: {},
|
||||
};
|
||||
|
||||
const UNSUPPORTED_TYPE_CONFIG: InputConfig = {
|
||||
inputs: {
|
||||
invalid_param: {
|
||||
// @ts-expect-error - Intentionally testing an invalid type
|
||||
type: 'date',
|
||||
description: 'This type is not supported',
|
||||
required: true,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
describe('convertInputConfigToJsonSchema', () => {
|
||||
describe('type conversion', () => {
|
||||
it('should correctly convert an InputConfig with various primitive types', () => {
|
||||
const result = convertInputConfigToJsonSchema(PRIMITIVE_TYPES_CONFIG);
|
||||
|
||||
expect(result).toEqual({
|
||||
type: 'object',
|
||||
properties: {
|
||||
goal: { type: 'string', description: 'The primary objective' },
|
||||
max_retries: {
|
||||
type: 'integer',
|
||||
description: 'Maximum number of retries',
|
||||
},
|
||||
temperature: { type: 'number', description: 'The model temperature' },
|
||||
verbose: { type: 'boolean', description: 'Enable verbose logging' },
|
||||
},
|
||||
required: ['goal', 'temperature'],
|
||||
});
|
||||
});
|
||||
|
||||
it('should correctly handle array types for strings and numbers', () => {
|
||||
const result = convertInputConfigToJsonSchema(ARRAY_TYPES_CONFIG);
|
||||
|
||||
expect(result).toEqual({
|
||||
type: 'object',
|
||||
properties: {
|
||||
filenames: {
|
||||
type: 'array',
|
||||
description: 'A list of file paths',
|
||||
items: { type: 'string' },
|
||||
},
|
||||
scores: {
|
||||
type: 'array',
|
||||
description: 'A list of scores',
|
||||
items: { type: 'number' },
|
||||
},
|
||||
},
|
||||
required: ['filenames'],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('required field handling', () => {
|
||||
it('should produce an undefined `required` field when no inputs are required', () => {
|
||||
const result = convertInputConfigToJsonSchema(NO_REQUIRED_FIELDS_CONFIG);
|
||||
|
||||
expect(result.properties['optional_param']).toBeDefined();
|
||||
// Per the implementation and JSON Schema spec, the `required` field
|
||||
// should be omitted if no properties are required.
|
||||
expect(result.required).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should list all properties in `required` when all are marked as required', () => {
|
||||
const result = convertInputConfigToJsonSchema(ALL_REQUIRED_FIELDS_CONFIG);
|
||||
expect(result.required).toHaveLength(2);
|
||||
expect(result.required).toEqual(
|
||||
expect.arrayContaining(['paramA', 'paramB']),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('edge cases', () => {
|
||||
it('should return a valid, empty schema for an empty input config', () => {
|
||||
const result = convertInputConfigToJsonSchema(EMPTY_CONFIG);
|
||||
|
||||
expect(result).toEqual({
|
||||
type: 'object',
|
||||
properties: {},
|
||||
required: undefined,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('error handling', () => {
|
||||
it('should throw an informative error for an unsupported input type', () => {
|
||||
const action = () =>
|
||||
convertInputConfigToJsonSchema(UNSUPPORTED_TYPE_CONFIG);
|
||||
|
||||
expect(action).toThrow(/Unsupported input type 'date'/);
|
||||
expect(action).toThrow(/parameter 'invalid_param'/);
|
||||
});
|
||||
});
|
||||
});
|
||||
90
packages/core/src/agents/schema-utils.ts
Normal file
90
packages/core/src/agents/schema-utils.ts
Normal file
@@ -0,0 +1,90 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { InputConfig } from './types.js';
|
||||
|
||||
/**
|
||||
* Defines the structure for a JSON Schema object, used for tool function
|
||||
* declarations.
|
||||
*/
|
||||
interface JsonSchemaObject {
|
||||
type: 'object';
|
||||
properties: Record<string, JsonSchemaProperty>;
|
||||
required?: string[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Defines the structure for a property within a {@link JsonSchemaObject}.
|
||||
*/
|
||||
interface JsonSchemaProperty {
|
||||
type: 'string' | 'number' | 'integer' | 'boolean' | 'array';
|
||||
description: string;
|
||||
items?: { type: 'string' | 'number' };
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts an internal `InputConfig` definition into a standard JSON Schema
|
||||
* object suitable for a tool's `FunctionDeclaration`.
|
||||
*
|
||||
* This utility ensures that the configuration for a subagent's inputs is
|
||||
* correctly translated into the format expected by the generative model.
|
||||
*
|
||||
* @param inputConfig The internal `InputConfig` to convert.
|
||||
* @returns A JSON Schema object representing the inputs.
|
||||
* @throws An `Error` if an unsupported input type is encountered, ensuring
|
||||
* configuration errors are caught early.
|
||||
*/
|
||||
export function convertInputConfigToJsonSchema(
|
||||
inputConfig: InputConfig,
|
||||
): JsonSchemaObject {
|
||||
const properties: Record<string, JsonSchemaProperty> = {};
|
||||
const required: string[] = [];
|
||||
|
||||
for (const [name, definition] of Object.entries(inputConfig.inputs)) {
|
||||
const schemaProperty: Partial<JsonSchemaProperty> = {
|
||||
description: definition.description,
|
||||
};
|
||||
|
||||
switch (definition.type) {
|
||||
case 'string':
|
||||
case 'number':
|
||||
case 'integer':
|
||||
case 'boolean':
|
||||
schemaProperty.type = definition.type;
|
||||
break;
|
||||
|
||||
case 'string[]':
|
||||
schemaProperty.type = 'array';
|
||||
schemaProperty.items = { type: 'string' };
|
||||
break;
|
||||
|
||||
case 'number[]':
|
||||
schemaProperty.type = 'array';
|
||||
schemaProperty.items = { type: 'number' };
|
||||
break;
|
||||
|
||||
default: {
|
||||
const exhaustiveCheck: never = definition.type;
|
||||
throw new Error(
|
||||
`Unsupported input type '${exhaustiveCheck}' for parameter '${name}'. ` +
|
||||
'Supported types: string, number, integer, boolean, string[], number[]',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
properties[name] = schemaProperty as JsonSchemaProperty;
|
||||
|
||||
if (definition.required) {
|
||||
required.push(name);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
type: 'object',
|
||||
properties,
|
||||
required: required.length > 0 ? required : undefined,
|
||||
};
|
||||
}
|
||||
138
packages/core/src/agents/subagent-tool-wrapper.test.ts
Normal file
138
packages/core/src/agents/subagent-tool-wrapper.test.ts
Normal file
@@ -0,0 +1,138 @@
|
||||
/**
|
||||
* @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 { SubagentInvocation } from './invocation.js';
|
||||
import { convertInputConfigToJsonSchema } from './schema-utils.js';
|
||||
import { makeFakeConfig } from '../test-utils/config.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import { Kind } from '../tools/tools.js';
|
||||
|
||||
// Mock dependencies to isolate the SubagentToolWrapper class
|
||||
vi.mock('./invocation.js');
|
||||
vi.mock('./schema-utils.js');
|
||||
|
||||
const MockedSubagentInvocation = vi.mocked(SubagentInvocation);
|
||||
const mockConvertInputConfigToJsonSchema = vi.mocked(
|
||||
convertInputConfigToJsonSchema,
|
||||
);
|
||||
|
||||
// Define reusable test data
|
||||
let mockConfig: Config;
|
||||
|
||||
const mockDefinition: AgentDefinition = {
|
||||
name: 'TestAgent',
|
||||
displayName: 'Test Agent Display Name',
|
||||
description: 'An agent for testing.',
|
||||
inputConfig: {
|
||||
inputs: {
|
||||
goal: { type: 'string', required: true, description: 'The goal.' },
|
||||
priority: {
|
||||
type: 'number',
|
||||
required: false,
|
||||
description: 'The priority.',
|
||||
},
|
||||
},
|
||||
},
|
||||
modelConfig: { model: 'gemini-test-model', temp: 0, top_p: 1 },
|
||||
runConfig: { max_time_minutes: 5 },
|
||||
promptConfig: { systemPrompt: 'You are a test agent.' },
|
||||
};
|
||||
|
||||
const mockSchema = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
goal: { type: 'string', description: 'The goal.' },
|
||||
priority: { type: 'number', description: 'The priority.' },
|
||||
},
|
||||
required: ['goal'],
|
||||
};
|
||||
|
||||
describe('SubagentToolWrapper', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = makeFakeConfig();
|
||||
// Provide a mock implementation for the schema conversion utility
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
mockConvertInputConfigToJsonSchema.mockReturnValue(mockSchema as any);
|
||||
});
|
||||
|
||||
describe('constructor', () => {
|
||||
it('should call convertInputConfigToJsonSchema with the correct agent inputConfig', () => {
|
||||
new SubagentToolWrapper(mockDefinition, mockConfig);
|
||||
|
||||
expect(convertInputConfigToJsonSchema).toHaveBeenCalledOnce();
|
||||
expect(convertInputConfigToJsonSchema).toHaveBeenCalledWith(
|
||||
mockDefinition.inputConfig,
|
||||
);
|
||||
});
|
||||
|
||||
it('should correctly configure the tool properties from the agent definition', () => {
|
||||
const wrapper = new SubagentToolWrapper(mockDefinition, mockConfig);
|
||||
|
||||
expect(wrapper.name).toBe(mockDefinition.name);
|
||||
expect(wrapper.displayName).toBe(mockDefinition.displayName);
|
||||
expect(wrapper.description).toBe(mockDefinition.description);
|
||||
expect(wrapper.kind).toBe(Kind.Think);
|
||||
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,
|
||||
);
|
||||
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);
|
||||
const schema = wrapper.schema;
|
||||
|
||||
expect(schema.name).toBe(mockDefinition.name);
|
||||
expect(schema.description).toBe(mockDefinition.description);
|
||||
expect(schema.parametersJsonSchema).toEqual(mockSchema);
|
||||
});
|
||||
});
|
||||
|
||||
describe('createInvocation', () => {
|
||||
it('should create a SubagentInvocation with the correct parameters', () => {
|
||||
const wrapper = new SubagentToolWrapper(mockDefinition, mockConfig);
|
||||
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(SubagentInvocation);
|
||||
expect(MockedSubagentInvocation).toHaveBeenCalledOnce();
|
||||
expect(MockedSubagentInvocation).toHaveBeenCalledWith(
|
||||
params,
|
||||
mockDefinition,
|
||||
mockConfig,
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw a validation error for invalid parameters before creating an invocation', () => {
|
||||
const wrapper = new SubagentToolWrapper(mockDefinition, mockConfig);
|
||||
// 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(MockedSubagentInvocation).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
69
packages/core/src/agents/subagent-tool-wrapper.ts
Normal file
69
packages/core/src/agents/subagent-tool-wrapper.ts
Normal file
@@ -0,0 +1,69 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
BaseDeclarativeTool,
|
||||
Kind,
|
||||
type ToolInvocation,
|
||||
type ToolResult,
|
||||
} from '../tools/tools.js';
|
||||
import type { Config } from '../config/config.js';
|
||||
import type { AgentDefinition, AgentInputs } from './types.js';
|
||||
import { convertInputConfigToJsonSchema } from './schema-utils.js';
|
||||
import { SubagentInvocation } from './invocation.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 config The runtime configuration, passed down to the subagent.
|
||||
*/
|
||||
constructor(
|
||||
private readonly definition: AgentDefinition,
|
||||
private readonly config: Config,
|
||||
) {
|
||||
// Dynamically generate the JSON schema required for the tool definition.
|
||||
const parameterSchema = convertInputConfigToJsonSchema(
|
||||
definition.inputConfig,
|
||||
);
|
||||
|
||||
super(
|
||||
definition.name,
|
||||
definition.displayName ?? definition.name,
|
||||
definition.description,
|
||||
Kind.Think,
|
||||
parameterSchema,
|
||||
/* 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,
|
||||
): ToolInvocation<AgentInputs, ToolResult> {
|
||||
return new SubagentInvocation(params, this.definition, this.config);
|
||||
}
|
||||
}
|
||||
139
packages/core/src/agents/types.ts
Normal file
139
packages/core/src/agents/types.ts
Normal file
@@ -0,0 +1,139 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
/**
|
||||
* @fileoverview Defines the core configuration interfaces and types for the agent architecture.
|
||||
*/
|
||||
|
||||
import type { Content, FunctionDeclaration } from '@google/genai';
|
||||
import type { AnyDeclarativeTool } from '../tools/tools.js';
|
||||
|
||||
/**
|
||||
* Describes the possible termination modes for an agent.
|
||||
*/
|
||||
export enum AgentTerminateMode {
|
||||
ERROR = 'ERROR',
|
||||
TIMEOUT = 'TIMEOUT',
|
||||
GOAL = 'GOAL',
|
||||
MAX_TURNS = 'MAX_TURNS',
|
||||
ABORTED = 'ABORTED',
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents the output structure of an agent's execution.
|
||||
*/
|
||||
export interface OutputObject {
|
||||
result: string;
|
||||
terminate_reason: AgentTerminateMode;
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents the validated input parameters passed to an agent upon invocation.
|
||||
* Used primarily for templating the system prompt. (Replaces ContextState)
|
||||
*/
|
||||
export type AgentInputs = Record<string, unknown>;
|
||||
|
||||
/**
|
||||
* Structured events emitted during subagent execution for user observability.
|
||||
*/
|
||||
export interface SubagentActivityEvent {
|
||||
isSubagentActivityEvent: true;
|
||||
agentName: string;
|
||||
type: 'TOOL_CALL_START' | 'TOOL_CALL_END' | 'THOUGHT_CHUNK' | 'ERROR';
|
||||
data: Record<string, unknown>;
|
||||
}
|
||||
|
||||
/**
|
||||
* The definition for an agent.
|
||||
*/
|
||||
export interface AgentDefinition {
|
||||
/** Unique identifier for the agent. */
|
||||
name: string;
|
||||
displayName?: string;
|
||||
description: string;
|
||||
promptConfig: PromptConfig;
|
||||
modelConfig: ModelConfig;
|
||||
runConfig: RunConfig;
|
||||
toolConfig?: ToolConfig;
|
||||
outputConfig?: OutputConfig;
|
||||
inputConfig: InputConfig;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the initial prompt for the agent.
|
||||
*/
|
||||
export interface PromptConfig {
|
||||
/**
|
||||
* A single system prompt string. Supports templating using `${input_name}` syntax.
|
||||
*/
|
||||
systemPrompt?: string;
|
||||
/**
|
||||
* An array of user/model content pairs for few-shot prompting.
|
||||
*/
|
||||
initialMessages?: Content[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the tools available to the agent during its execution.
|
||||
*/
|
||||
export interface ToolConfig {
|
||||
tools: Array<string | FunctionDeclaration | AnyDeclarativeTool>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the expected inputs (parameters) for the agent.
|
||||
*/
|
||||
export interface InputConfig {
|
||||
/**
|
||||
* Defines the parameters the agent accepts.
|
||||
* This is vital for generating the tool wrapper schema.
|
||||
*/
|
||||
inputs: Record<
|
||||
string,
|
||||
{
|
||||
description: string;
|
||||
type:
|
||||
| 'string'
|
||||
| 'number'
|
||||
| 'boolean'
|
||||
| 'integer'
|
||||
| 'string[]'
|
||||
| 'number[]';
|
||||
required: boolean;
|
||||
}
|
||||
>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the expected outputs for the agent.
|
||||
*/
|
||||
export interface OutputConfig {
|
||||
/** Description of what the agent should return when finished. */
|
||||
description: string;
|
||||
/** Optional criteria that must be completed before the agent finishes. */
|
||||
completion_criteria?: string[];
|
||||
// TODO(abhipatel12): Add required_outputs if natural completion insufficient
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the generative model parameters for the agent.
|
||||
*/
|
||||
export interface ModelConfig {
|
||||
model: string;
|
||||
temp: number;
|
||||
top_p: number;
|
||||
thinkingBudget?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the execution environment and constraints for the agent.
|
||||
*/
|
||||
export interface RunConfig {
|
||||
/** The maximum execution time for the agent in minutes. */
|
||||
max_time_minutes: number;
|
||||
/** The maximum number of conversational turns. */
|
||||
max_turns?: number;
|
||||
}
|
||||
115
packages/core/src/agents/utils.test.ts
Normal file
115
packages/core/src/agents/utils.test.ts
Normal file
@@ -0,0 +1,115 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { templateString } from './utils.js';
|
||||
import type { AgentInputs } from './types.js';
|
||||
|
||||
describe('templateString', () => {
|
||||
it('should replace a single placeholder with a string value', () => {
|
||||
const template = 'Hello, ${name}!';
|
||||
const inputs: AgentInputs = { name: 'World' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('Hello, World!');
|
||||
});
|
||||
|
||||
it('should replace multiple unique placeholders', () => {
|
||||
const template = 'User: ${user}, Role: ${role}';
|
||||
const inputs: AgentInputs = { user: 'Alex', role: 'Admin' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('User: Alex, Role: Admin');
|
||||
});
|
||||
|
||||
it('should replace multiple instances of the same placeholder', () => {
|
||||
const template = '${greeting}, ${user}. Welcome, ${user}!';
|
||||
const inputs: AgentInputs = { greeting: 'Hi', user: 'Sam' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('Hi, Sam. Welcome, Sam!');
|
||||
});
|
||||
|
||||
it('should handle various data types for input values', () => {
|
||||
const template =
|
||||
'Name: ${name}, Age: ${age}, Active: ${isActive}, Plan: ${plan}, Score: ${score}';
|
||||
const inputs: AgentInputs = {
|
||||
name: 'Jo',
|
||||
age: 30,
|
||||
isActive: true,
|
||||
plan: null,
|
||||
score: undefined,
|
||||
};
|
||||
const result = templateString(template, inputs);
|
||||
// All values are converted to their string representations
|
||||
expect(result).toBe(
|
||||
'Name: Jo, Age: 30, Active: true, Plan: null, Score: undefined',
|
||||
);
|
||||
});
|
||||
|
||||
it('should return the original string if no placeholders are present', () => {
|
||||
const template = 'This is a plain string with no placeholders.';
|
||||
const inputs: AgentInputs = { key: 'value' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('This is a plain string with no placeholders.');
|
||||
});
|
||||
|
||||
it('should correctly handle an empty template string', () => {
|
||||
const template = '';
|
||||
const inputs: AgentInputs = { key: 'value' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('');
|
||||
});
|
||||
|
||||
it('should ignore extra keys in the inputs object that are not in the template', () => {
|
||||
const template = 'Hello, ${name}.';
|
||||
const inputs: AgentInputs = { name: 'Alice', extra: 'ignored' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('Hello, Alice.');
|
||||
});
|
||||
|
||||
it('should throw an error if a required key is missing from the inputs', () => {
|
||||
const template = 'The goal is ${goal}.';
|
||||
const inputs: AgentInputs = { other_input: 'some value' };
|
||||
|
||||
expect(() => templateString(template, inputs)).toThrow(
|
||||
'Template validation failed: Missing required input parameters: goal. Available inputs: other_input',
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw an error listing all missing keys if multiple are missing', () => {
|
||||
const template = 'Analyze ${file} with ${tool}.';
|
||||
const inputs: AgentInputs = { an_available_key: 'foo' };
|
||||
|
||||
// Using a regex to allow for any order of missing keys in the error message
|
||||
expect(() => templateString(template, inputs)).toThrow(
|
||||
/Missing required input parameters: (file, tool|tool, file)/,
|
||||
);
|
||||
});
|
||||
|
||||
it('should be case-sensitive with placeholder keys', () => {
|
||||
const template = 'Value: ${Key}';
|
||||
const inputs: AgentInputs = { key: 'some value' }; // 'key' is lowercase
|
||||
|
||||
expect(() => templateString(template, inputs)).toThrow(
|
||||
'Template validation failed: Missing required input parameters: Key. Available inputs: key',
|
||||
);
|
||||
});
|
||||
|
||||
it('should not replace malformed or incomplete placeholders', () => {
|
||||
const template =
|
||||
'This is {not_a_placeholder} and this is $$escaped. Test: ${valid}';
|
||||
const inputs: AgentInputs = { valid: 'works' };
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe(
|
||||
'This is {not_a_placeholder} and this is $$escaped. Test: works',
|
||||
);
|
||||
});
|
||||
|
||||
it('should work correctly with an empty inputs object if the template has no placeholders', () => {
|
||||
const template = 'Static text.';
|
||||
const inputs: AgentInputs = {};
|
||||
const result = templateString(template, inputs);
|
||||
expect(result).toBe('Static text.');
|
||||
});
|
||||
});
|
||||
43
packages/core/src/agents/utils.ts
Normal file
43
packages/core/src/agents/utils.ts
Normal file
@@ -0,0 +1,43 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type { AgentInputs } from './types.js';
|
||||
|
||||
/**
|
||||
* Replaces `${...}` placeholders in a template string with values from AgentInputs.
|
||||
*
|
||||
* @param template The template string containing placeholders.
|
||||
* @param inputs The AgentInputs object providing placeholder values.
|
||||
* @returns The populated string with all placeholders replaced.
|
||||
* @throws {Error} if any placeholder key is not found in the inputs.
|
||||
*/
|
||||
export function templateString(template: string, inputs: AgentInputs): string {
|
||||
const placeholderRegex = /\$\{(\w+)\}/g;
|
||||
|
||||
// First, find all unique keys required by the template.
|
||||
const requiredKeys = new Set(
|
||||
Array.from(template.matchAll(placeholderRegex), (match) => match[1]),
|
||||
);
|
||||
|
||||
// Check if all required keys exist in the inputs.
|
||||
const inputKeys = new Set(Object.keys(inputs));
|
||||
const missingKeys = Array.from(requiredKeys).filter(
|
||||
(key) => !inputKeys.has(key),
|
||||
);
|
||||
|
||||
if (missingKeys.length > 0) {
|
||||
// Enhanced error message showing both missing and available keys
|
||||
throw new Error(
|
||||
`Template validation failed: Missing required input parameters: ${missingKeys.join(', ')}. ` +
|
||||
`Available inputs: ${Object.keys(inputs).join(', ')}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Perform the replacement using a replacer function.
|
||||
return template.replace(placeholderRegex, (_match, key) =>
|
||||
String(inputs[key]),
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user