mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-22 19:14:33 -07:00
refactor(agents): Introduce Declarative Agent Framework (#9778)
This commit is contained in:
Abhi
@@ -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);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user