feat(agents): migrate subagents to event-driven scheduler (#17567)

This commit is contained in:
Abhi
2026-01-26 17:12:55 -05:00
committed by GitHub
parent 13bc5f620c
commit 9d34ae52d6
8 changed files with 741 additions and 335 deletions
+267 -242
View File
@@ -55,6 +55,7 @@ import type {
} from './types.js';
import { AgentTerminateMode } from './types.js';
import type { AnyDeclarativeTool, AnyToolInvocation } from '../tools/tools.js';
import type { ToolCallRequestInfo } from '../scheduler/types.js';
import { CompressionStatus } from '../core/turn.js';
import { ChatCompressionService } from '../services/chatCompressionService.js';
import type {
@@ -67,12 +68,12 @@ import type { ModelRouterService } from '../routing/modelRouterService.js';
const {
mockSendMessageStream,
mockExecuteToolCall,
mockScheduleAgentTools,
mockSetSystemInstruction,
mockCompress,
} = vi.hoisted(() => ({
mockSendMessageStream: vi.fn(),
mockExecuteToolCall: vi.fn(),
mockScheduleAgentTools: vi.fn(),
mockSetSystemInstruction: vi.fn(),
mockCompress: vi.fn(),
}));
@@ -101,8 +102,8 @@ vi.mock('../core/geminiChat.js', async (importOriginal) => {
};
});
vi.mock('../core/nonInteractiveToolExecutor.js', () => ({
executeToolCall: mockExecuteToolCall,
vi.mock('./agent-scheduler.js', () => ({
scheduleAgentTools: mockScheduleAgentTools,
}));
vi.mock('../utils/version.js', () => ({
@@ -275,7 +276,7 @@ describe('LocalAgentExecutor', () => {
mockSetHistory.mockClear();
mockSendMessageStream.mockReset();
mockSetSystemInstruction.mockReset();
mockExecuteToolCall.mockReset();
mockScheduleAgentTools.mockReset();
mockedLogAgentStart.mockReset();
mockedLogAgentFinish.mockReset();
mockedPromptIdContext.getStore.mockReset();
@@ -540,34 +541,36 @@ describe('LocalAgentExecutor', () => {
[{ name: LS_TOOL_NAME, args: { path: '.' }, id: 'call1' }],
'T1: Listing',
);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'file1.txt',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: { result: 'file1.txt' },
id: 'call1',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'file1.txt',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: { result: 'file1.txt' },
id: 'call1',
},
},
},
],
error: undefined,
errorType: undefined,
contentLength: undefined,
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
});
]);
// Turn 2: Model calls complete_task with required output
mockModelResponse(
@@ -686,34 +689,36 @@ describe('LocalAgentExecutor', () => {
mockModelResponse([
{ name: LS_TOOL_NAME, args: { path: '.' }, id: 'call1' },
]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: {},
id: 'call1',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: {},
id: 'call1',
},
},
},
],
error: undefined,
errorType: undefined,
contentLength: undefined,
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
});
]);
mockModelResponse(
[
@@ -759,34 +764,36 @@ describe('LocalAgentExecutor', () => {
mockModelResponse([
{ name: LS_TOOL_NAME, args: { path: '.' }, id: 'call1' },
]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: {},
id: 'call1',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: {},
id: 'call1',
},
},
},
],
error: undefined,
errorType: undefined,
contentLength: undefined,
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
});
]);
// Turn 2 (protocol violation)
mockModelResponse([], 'I think I am done.');
@@ -959,33 +966,40 @@ describe('LocalAgentExecutor', () => {
resolveCalls = r;
});
mockExecuteToolCall.mockImplementation(async (_ctx, reqInfo) => {
callsStarted++;
if (callsStarted === 2) resolveCalls();
await vi.advanceTimersByTimeAsync(100);
return {
status: 'success',
request: reqInfo,
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: reqInfo.callId,
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: reqInfo.name,
response: {},
id: reqInfo.callId,
mockScheduleAgentTools.mockImplementation(
async (_ctx, requests: ToolCallRequestInfo[]) => {
const results = await Promise.all(
requests.map(async (reqInfo) => {
callsStarted++;
if (callsStarted === 2) resolveCalls();
await vi.advanceTimersByTimeAsync(100);
return {
status: 'success',
request: reqInfo,
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: reqInfo.callId,
resultDisplay: 'ok',
responseParts: [
{
functionResponse: {
name: reqInfo.name,
response: {},
id: reqInfo.callId,
},
},
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
};
});
};
}),
);
return results;
},
);
// Turn 2: Completion
mockModelResponse([
@@ -1005,7 +1019,7 @@ describe('LocalAgentExecutor', () => {
const output = await runPromise;
expect(mockExecuteToolCall).toHaveBeenCalledTimes(2);
expect(mockScheduleAgentTools).toHaveBeenCalledTimes(1);
expect(output.terminate_reason).toBe(AgentTerminateMode.GOAL);
// Safe access to message parts
@@ -1059,7 +1073,7 @@ describe('LocalAgentExecutor', () => {
await executor.run({ goal: 'Sec test' }, signal);
// Verify external executor was not called (Security held)
expect(mockExecuteToolCall).not.toHaveBeenCalled();
expect(mockScheduleAgentTools).not.toHaveBeenCalled();
// 2. Verify console warning
expect(consoleWarnSpy).toHaveBeenCalledWith(
@@ -1215,37 +1229,36 @@ describe('LocalAgentExecutor', () => {
mockModelResponse([
{ name: LS_TOOL_NAME, args: { path: '/fake' }, id: 'call1' },
]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'error',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '/fake' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: '',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: { error: toolErrorMessage },
id: 'call1',
},
},
],
error: {
type: 'ToolError',
message: toolErrorMessage,
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'error',
request: {
callId: 'call1',
name: LS_TOOL_NAME,
args: { path: '/fake' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 'call1',
resultDisplay: '',
responseParts: [
{
functionResponse: {
name: LS_TOOL_NAME,
response: { error: toolErrorMessage },
id: 'call1',
},
},
],
error: new Error(toolErrorMessage),
errorType: 'ToolError',
contentLength: 0,
},
errorType: 'ToolError',
contentLength: 0,
},
});
]);
// Turn 2: Model sees the error and completes
mockModelResponse([
@@ -1258,7 +1271,7 @@ describe('LocalAgentExecutor', () => {
const output = await executor.run({ goal: 'Tool failure test' }, signal);
expect(mockExecuteToolCall).toHaveBeenCalledTimes(1);
expect(mockScheduleAgentTools).toHaveBeenCalledTimes(1);
expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
// Verify the error was reported in the activity stream
@@ -1391,28 +1404,30 @@ describe('LocalAgentExecutor', () => {
describe('run (Termination Conditions)', () => {
const mockWorkResponse = (id: string) => {
mockModelResponse([{ name: LS_TOOL_NAME, args: { path: '.' }, id }]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
});
]);
};
it('should terminate when max_turns is reached', async () => {
@@ -1505,23 +1520,27 @@ describe('LocalAgentExecutor', () => {
]);
// Long running tool
mockExecuteToolCall.mockImplementationOnce(async (_ctx, reqInfo) => {
await vi.advanceTimersByTimeAsync(61 * 1000);
return {
status: 'success',
request: reqInfo,
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 't1',
resultDisplay: 'ok',
responseParts: [],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
};
});
mockScheduleAgentTools.mockImplementationOnce(
async (_ctx, requests: ToolCallRequestInfo[]) => {
await vi.advanceTimersByTimeAsync(61 * 1000);
return [
{
status: 'success',
request: requests[0],
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: 't1',
resultDisplay: 'ok',
responseParts: [],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
];
},
);
// Recovery turn
mockModelResponse([], 'I give up');
@@ -1557,28 +1576,30 @@ describe('LocalAgentExecutor', () => {
describe('run (Recovery Turns)', () => {
const mockWorkResponse = (id: string) => {
mockModelResponse([{ name: LS_TOOL_NAME, args: { path: '.' }, id }]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
});
]);
};
it('should recover successfully if complete_task is called during the grace turn after MAX_TURNS', async () => {
@@ -1873,28 +1894,30 @@ describe('LocalAgentExecutor', () => {
describe('Telemetry and Logging', () => {
const mockWorkResponse = (id: string) => {
mockModelResponse([{ name: LS_TOOL_NAME, args: { path: '.' }, id }]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
});
]);
};
beforeEach(() => {
@@ -1960,28 +1983,30 @@ describe('LocalAgentExecutor', () => {
describe('Chat Compression', () => {
const mockWorkResponse = (id: string) => {
mockModelResponse([{ name: LS_TOOL_NAME, args: { path: '.' }, id }]);
mockExecuteToolCall.mockResolvedValueOnce({
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
mockScheduleAgentTools.mockResolvedValueOnce([
{
status: 'success',
request: {
callId: id,
name: LS_TOOL_NAME,
args: { path: '.' },
isClientInitiated: false,
prompt_id: 'test-prompt',
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
},
tool: {} as AnyDeclarativeTool,
invocation: {} as AnyToolInvocation,
response: {
callId: id,
resultDisplay: 'ok',
responseParts: [
{ functionResponse: { name: LS_TOOL_NAME, response: {}, id } },
],
error: undefined,
errorType: undefined,
contentLength: undefined,
},
});
]);
};
it('should attempt to compress chat history on each turn', async () => {