mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-22 02:54:31 -07:00
Merge branch 'main' into galzahavi/add/sandboxed-tool
This commit is contained in:
galz10
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@google/gemini-cli-a2a-server",
|
||||
"version": "0.34.0-nightly.20260304.28af4e127",
|
||||
"version": "0.35.0-nightly.20260311.657f19c1f",
|
||||
"description": "Gemini CLI A2A Server",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
@@ -25,7 +25,7 @@
|
||||
"dist"
|
||||
],
|
||||
"dependencies": {
|
||||
"@a2a-js/sdk": "^0.3.8",
|
||||
"@a2a-js/sdk": "0.3.11",
|
||||
"@google-cloud/storage": "^7.16.0",
|
||||
"@google/gemini-cli-core": "file:../core",
|
||||
"express": "^5.1.0",
|
||||
@@ -36,7 +36,7 @@
|
||||
"winston": "^3.17.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@google/genai": "^1.30.0",
|
||||
"@google/genai": "1.30.0",
|
||||
"@types/express": "^5.0.3",
|
||||
"@types/fs-extra": "^11.0.4",
|
||||
"@types/supertest": "^6.0.3",
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||
import { CoderAgentExecutor } from './executor.js';
|
||||
import type {
|
||||
ExecutionEventBus,
|
||||
RequestContext,
|
||||
TaskStore,
|
||||
} from '@a2a-js/sdk/server';
|
||||
import { EventEmitter } from 'node:events';
|
||||
import { requestStorage } from '../http/requestStorage.js';
|
||||
|
||||
// Mocks for constructor dependencies
|
||||
vi.mock('../config/config.js', () => ({
|
||||
loadConfig: vi.fn().mockReturnValue({
|
||||
getSessionId: () => 'test-session',
|
||||
getTargetDir: () => '/tmp',
|
||||
getCheckpointingEnabled: () => false,
|
||||
}),
|
||||
loadEnvironment: vi.fn(),
|
||||
setTargetDir: vi.fn().mockReturnValue('/tmp'),
|
||||
}));
|
||||
|
||||
vi.mock('../config/settings.js', () => ({
|
||||
loadSettings: vi.fn().mockReturnValue({}),
|
||||
}));
|
||||
|
||||
vi.mock('../config/extension.js', () => ({
|
||||
loadExtensions: vi.fn().mockReturnValue([]),
|
||||
}));
|
||||
|
||||
vi.mock('../http/requestStorage.js', () => ({
|
||||
requestStorage: {
|
||||
getStore: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('./task.js', () => {
|
||||
const mockTaskInstance = (taskId: string, contextId: string) => ({
|
||||
id: taskId,
|
||||
contextId,
|
||||
taskState: 'working',
|
||||
acceptUserMessage: vi
|
||||
.fn()
|
||||
.mockImplementation(async function* (context, aborted) {
|
||||
const isConfirmation = (
|
||||
context.userMessage.parts as Array<{ kind: string }>
|
||||
).some((p) => p.kind === 'confirmation');
|
||||
// Hang only for main user messages (text), allow confirmations to finish quickly
|
||||
if (!isConfirmation && aborted) {
|
||||
await new Promise((resolve) => {
|
||||
aborted.addEventListener('abort', resolve, { once: true });
|
||||
});
|
||||
}
|
||||
yield { type: 'content', value: 'hello' };
|
||||
}),
|
||||
acceptAgentMessage: vi.fn().mockResolvedValue(undefined),
|
||||
scheduleToolCalls: vi.fn().mockResolvedValue(undefined),
|
||||
waitForPendingTools: vi.fn().mockResolvedValue(undefined),
|
||||
getAndClearCompletedTools: vi.fn().mockReturnValue([]),
|
||||
addToolResponsesToHistory: vi.fn(),
|
||||
sendCompletedToolsToLlm: vi.fn().mockImplementation(async function* () {}),
|
||||
cancelPendingTools: vi.fn(),
|
||||
setTaskStateAndPublishUpdate: vi.fn(),
|
||||
dispose: vi.fn(),
|
||||
getMetadata: vi.fn().mockResolvedValue({}),
|
||||
geminiClient: {
|
||||
initialize: vi.fn().mockResolvedValue(undefined),
|
||||
},
|
||||
toSDKTask: () => ({
|
||||
id: taskId,
|
||||
contextId,
|
||||
kind: 'task',
|
||||
status: { state: 'working', timestamp: new Date().toISOString() },
|
||||
metadata: {},
|
||||
history: [],
|
||||
artifacts: [],
|
||||
}),
|
||||
});
|
||||
|
||||
const MockTask = vi.fn().mockImplementation(mockTaskInstance);
|
||||
(MockTask as unknown as { create: Mock }).create = vi
|
||||
.fn()
|
||||
.mockImplementation(async (taskId: string, contextId: string) =>
|
||||
mockTaskInstance(taskId, contextId),
|
||||
);
|
||||
|
||||
return { Task: MockTask };
|
||||
});
|
||||
|
||||
describe('CoderAgentExecutor', () => {
|
||||
let executor: CoderAgentExecutor;
|
||||
let mockTaskStore: TaskStore;
|
||||
let mockEventBus: ExecutionEventBus;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockTaskStore = {
|
||||
save: vi.fn().mockResolvedValue(undefined),
|
||||
load: vi.fn().mockResolvedValue(undefined),
|
||||
delete: vi.fn().mockResolvedValue(undefined),
|
||||
list: vi.fn().mockResolvedValue([]),
|
||||
} as unknown as TaskStore;
|
||||
|
||||
mockEventBus = new EventEmitter() as unknown as ExecutionEventBus;
|
||||
mockEventBus.publish = vi.fn();
|
||||
mockEventBus.finished = vi.fn();
|
||||
|
||||
executor = new CoderAgentExecutor(mockTaskStore);
|
||||
});
|
||||
|
||||
it('should distinguish between primary and secondary execution', async () => {
|
||||
const taskId = 'test-task';
|
||||
const contextId = 'test-context';
|
||||
|
||||
const mockSocket = new EventEmitter();
|
||||
const requestContext = {
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
taskId,
|
||||
contextId,
|
||||
parts: [{ kind: 'text', text: 'hi' }],
|
||||
metadata: {
|
||||
coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
|
||||
},
|
||||
},
|
||||
} as unknown as RequestContext;
|
||||
|
||||
// Mock requestStorage for primary
|
||||
(requestStorage.getStore as Mock).mockReturnValue({
|
||||
req: { socket: mockSocket },
|
||||
});
|
||||
|
||||
// First execution (Primary)
|
||||
const primaryPromise = executor.execute(requestContext, mockEventBus);
|
||||
|
||||
// Give it enough time to reach line 490 in executor.ts
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
expect(
|
||||
(
|
||||
executor as unknown as { executingTasks: Set<string> }
|
||||
).executingTasks.has(taskId),
|
||||
).toBe(true);
|
||||
const wrapper = executor.getTask(taskId);
|
||||
expect(wrapper).toBeDefined();
|
||||
|
||||
// Mock requestStorage for secondary
|
||||
const secondarySocket = new EventEmitter();
|
||||
(requestStorage.getStore as Mock).mockReturnValue({
|
||||
req: { socket: secondarySocket },
|
||||
});
|
||||
|
||||
const secondaryRequestContext = {
|
||||
userMessage: {
|
||||
messageId: 'msg-2',
|
||||
taskId,
|
||||
contextId,
|
||||
parts: [{ kind: 'confirmation', callId: '1', outcome: 'proceed' }],
|
||||
metadata: {
|
||||
coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
|
||||
},
|
||||
},
|
||||
} as unknown as RequestContext;
|
||||
|
||||
const secondaryPromise = executor.execute(
|
||||
secondaryRequestContext,
|
||||
mockEventBus,
|
||||
);
|
||||
|
||||
// Secondary execution should NOT add to executingTasks (already there)
|
||||
// and should return early after its loop
|
||||
await secondaryPromise;
|
||||
|
||||
// Task should still be in executingTasks and NOT disposed
|
||||
expect(
|
||||
(
|
||||
executor as unknown as { executingTasks: Set<string> }
|
||||
).executingTasks.has(taskId),
|
||||
).toBe(true);
|
||||
expect(wrapper?.task.dispose).not.toHaveBeenCalled();
|
||||
|
||||
// Now simulate secondary socket closure - it should NOT affect primary
|
||||
secondarySocket.emit('end');
|
||||
expect(
|
||||
(
|
||||
executor as unknown as { executingTasks: Set<string> }
|
||||
).executingTasks.has(taskId),
|
||||
).toBe(true);
|
||||
expect(wrapper?.task.dispose).not.toHaveBeenCalled();
|
||||
|
||||
// Set to terminal state to verify disposal on finish
|
||||
wrapper!.task.taskState = 'completed';
|
||||
|
||||
// Now close primary socket
|
||||
mockSocket.emit('end');
|
||||
|
||||
await primaryPromise;
|
||||
|
||||
expect(
|
||||
(
|
||||
executor as unknown as { executingTasks: Set<string> }
|
||||
).executingTasks.has(taskId),
|
||||
).toBe(false);
|
||||
expect(wrapper?.task.dispose).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should evict task from cache when it reaches terminal state', async () => {
|
||||
const taskId = 'test-task-terminal';
|
||||
const contextId = 'test-context';
|
||||
|
||||
const mockSocket = new EventEmitter();
|
||||
(requestStorage.getStore as Mock).mockReturnValue({
|
||||
req: { socket: mockSocket },
|
||||
});
|
||||
|
||||
const requestContext = {
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
taskId,
|
||||
contextId,
|
||||
parts: [{ kind: 'text', text: 'hi' }],
|
||||
metadata: {
|
||||
coderAgent: { kind: 'agent-settings', workspacePath: '/tmp' },
|
||||
},
|
||||
},
|
||||
} as unknown as RequestContext;
|
||||
|
||||
const primaryPromise = executor.execute(requestContext, mockEventBus);
|
||||
await new Promise((resolve) => setTimeout(resolve, 50));
|
||||
|
||||
const wrapper = executor.getTask(taskId)!;
|
||||
expect(wrapper).toBeDefined();
|
||||
// Simulate terminal state
|
||||
wrapper.task.taskState = 'completed';
|
||||
|
||||
// Finish primary execution
|
||||
mockSocket.emit('end');
|
||||
await primaryPromise;
|
||||
|
||||
expect(executor.getTask(taskId)).toBeUndefined();
|
||||
expect(wrapper.task.dispose).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -252,6 +252,10 @@ export class CoderAgentExecutor implements AgentExecutor {
|
||||
);
|
||||
await this.taskStore?.save(wrapper.toSDKTask());
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} state CANCELED saved.`);
|
||||
|
||||
// Cleanup listener subscriptions to avoid memory leaks.
|
||||
wrapper.task.dispose();
|
||||
this.tasks.delete(taskId);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
@@ -320,23 +324,26 @@ export class CoderAgentExecutor implements AgentExecutor {
|
||||
if (store) {
|
||||
// Grab the raw socket from the request object
|
||||
const socket = store.req.socket;
|
||||
const onClientEnd = () => {
|
||||
const onSocketEnd = () => {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Client socket closed for task ${taskId}. Cancelling execution.`,
|
||||
`[CoderAgentExecutor] Socket ended for message ${userMessage.messageId} (task ${taskId}). Aborting execution loop.`,
|
||||
);
|
||||
if (!abortController.signal.aborted) {
|
||||
abortController.abort();
|
||||
}
|
||||
// Clean up the listener to prevent memory leaks
|
||||
socket.removeListener('close', onClientEnd);
|
||||
socket.removeListener('end', onSocketEnd);
|
||||
};
|
||||
|
||||
// Listen on the socket's 'end' event (remote closed the connection)
|
||||
socket.on('end', onClientEnd);
|
||||
socket.on('end', onSocketEnd);
|
||||
socket.once('close', () => {
|
||||
socket.removeListener('end', onSocketEnd);
|
||||
});
|
||||
|
||||
// It's also good practice to remove the listener if the task completes successfully
|
||||
abortSignal.addEventListener('abort', () => {
|
||||
socket.removeListener('end', onClientEnd);
|
||||
socket.removeListener('end', onSocketEnd);
|
||||
});
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Socket close handler set up for task ${taskId}.`,
|
||||
@@ -457,6 +464,26 @@ export class CoderAgentExecutor implements AgentExecutor {
|
||||
return;
|
||||
}
|
||||
|
||||
// Check if this is the primary/initial execution for this task
|
||||
const isPrimaryExecution = !this.executingTasks.has(taskId);
|
||||
|
||||
if (!isPrimaryExecution) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Primary execution already active for task ${taskId}. Starting secondary loop for message ${userMessage.messageId}.`,
|
||||
);
|
||||
currentTask.eventBus = eventBus;
|
||||
for await (const _ of currentTask.acceptUserMessage(
|
||||
requestContext,
|
||||
abortController.signal,
|
||||
)) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Processing user message ${userMessage.messageId} in secondary execution loop for task ${taskId}.`,
|
||||
);
|
||||
}
|
||||
// End this execution-- the original/source will be resumed.
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Starting main execution for message ${userMessage.messageId} for task ${taskId}.`,
|
||||
);
|
||||
@@ -598,18 +625,30 @@ export class CoderAgentExecutor implements AgentExecutor {
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
this.executingTasks.delete(taskId);
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Saving final state for task ${taskId}.`,
|
||||
);
|
||||
try {
|
||||
await this.taskStore?.save(wrapper.toSDKTask());
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} state saved.`);
|
||||
} catch (saveError) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Failed to save task ${taskId} state in finally block:`,
|
||||
saveError,
|
||||
if (isPrimaryExecution) {
|
||||
this.executingTasks.delete(taskId);
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Saving final state for task ${taskId}.`,
|
||||
);
|
||||
try {
|
||||
await this.taskStore?.save(wrapper.toSDKTask());
|
||||
logger.info(`[CoderAgentExecutor] Task ${taskId} state saved.`);
|
||||
} catch (saveError) {
|
||||
logger.error(
|
||||
`[CoderAgentExecutor] Failed to save task ${taskId} state in finally block:`,
|
||||
saveError,
|
||||
);
|
||||
}
|
||||
|
||||
if (
|
||||
['canceled', 'failed', 'completed'].includes(currentTask.taskState)
|
||||
) {
|
||||
logger.info(
|
||||
`[CoderAgentExecutor] Task ${taskId} reached terminal state ${currentTask.taskState}. Evicting and disposing.`,
|
||||
);
|
||||
wrapper.task.dispose();
|
||||
this.tasks.delete(taskId);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,655 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||
import { Task } from './task.js';
|
||||
import {
|
||||
type Config,
|
||||
MessageBusType,
|
||||
ToolConfirmationOutcome,
|
||||
ApprovalMode,
|
||||
Scheduler,
|
||||
type MessageBus,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { createMockConfig } from '../utils/testing_utils.js';
|
||||
import type { ExecutionEventBus } from '@a2a-js/sdk/server';
|
||||
|
||||
describe('Task Event-Driven Scheduler', () => {
|
||||
let mockConfig: Config;
|
||||
let mockEventBus: ExecutionEventBus;
|
||||
let messageBus: MessageBus;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockConfig = createMockConfig({
|
||||
isEventDrivenSchedulerEnabled: () => true,
|
||||
}) as Config;
|
||||
messageBus = mockConfig.getMessageBus();
|
||||
mockEventBus = {
|
||||
publish: vi.fn(),
|
||||
on: vi.fn(),
|
||||
off: vi.fn(),
|
||||
once: vi.fn(),
|
||||
removeAllListeners: vi.fn(),
|
||||
finished: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
it('should instantiate Scheduler when enabled', () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
expect(task.scheduler).toBeInstanceOf(Scheduler);
|
||||
});
|
||||
|
||||
it('should subscribe to TOOL_CALLS_UPDATE and map status changes', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'executing',
|
||||
};
|
||||
|
||||
// Simulate MessageBus event
|
||||
// Simulate MessageBus event
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
|
||||
if (!handler) {
|
||||
throw new Error('TOOL_CALLS_UPDATE handler not found');
|
||||
}
|
||||
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall],
|
||||
});
|
||||
|
||||
expect(mockEventBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
status: expect.objectContaining({
|
||||
state: 'submitted', // initial task state
|
||||
}),
|
||||
metadata: expect.objectContaining({
|
||||
coderAgent: expect.objectContaining({
|
||||
kind: 'tool-call-update',
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle tool confirmations by publishing to MessageBus', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-1',
|
||||
confirmationDetails: { type: 'info', title: 'test', prompt: 'test' },
|
||||
};
|
||||
|
||||
// Simulate MessageBus event to stash the correlationId
|
||||
// Simulate MessageBus event
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
|
||||
if (!handler) {
|
||||
throw new Error('TOOL_CALLS_UPDATE handler not found');
|
||||
}
|
||||
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall],
|
||||
});
|
||||
|
||||
// Simulate A2A client confirmation
|
||||
const part = {
|
||||
kind: 'data',
|
||||
data: {
|
||||
callId: '1',
|
||||
outcome: 'proceed_once',
|
||||
},
|
||||
};
|
||||
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart(part);
|
||||
expect(handled).toBe(true);
|
||||
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-1',
|
||||
confirmed: true,
|
||||
outcome: ToolConfirmationOutcome.ProceedOnce,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle Rejection (Cancel) and Modification (ModifyWithEditor)', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-1',
|
||||
confirmationDetails: { type: 'info', title: 'test', prompt: 'test' },
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
// Simulate Rejection (Cancel)
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'cancel' },
|
||||
});
|
||||
expect(handled).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-1',
|
||||
confirmed: false,
|
||||
}),
|
||||
);
|
||||
|
||||
const toolCall2 = {
|
||||
request: { callId: '2', name: 'ls', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-2',
|
||||
confirmationDetails: { type: 'info', title: 'test', prompt: 'test' },
|
||||
};
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall2] });
|
||||
|
||||
// Simulate ModifyWithEditor
|
||||
const handled2 = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '2', outcome: 'modify_with_editor' },
|
||||
});
|
||||
expect(handled2).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-2',
|
||||
confirmed: false,
|
||||
outcome: ToolConfirmationOutcome.ModifyWithEditor,
|
||||
payload: undefined,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle MCP Server tool operations correctly', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'call_mcp_tool', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-mcp-1',
|
||||
confirmationDetails: {
|
||||
type: 'mcp',
|
||||
title: 'MCP Server Operation',
|
||||
prompt: 'test_mcp',
|
||||
},
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
// Simulate ProceedOnce for MCP
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'proceed_once' },
|
||||
});
|
||||
expect(handled).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-mcp-1',
|
||||
confirmed: true,
|
||||
outcome: ToolConfirmationOutcome.ProceedOnce,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle MCP Server tool ProceedAlwaysServer outcome', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'call_mcp_tool', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-mcp-2',
|
||||
confirmationDetails: {
|
||||
type: 'mcp',
|
||||
title: 'MCP Server Operation',
|
||||
prompt: 'test_mcp',
|
||||
},
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'proceed_always_server' },
|
||||
});
|
||||
expect(handled).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-mcp-2',
|
||||
confirmed: true,
|
||||
outcome: ToolConfirmationOutcome.ProceedAlwaysServer,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle MCP Server tool ProceedAlwaysTool outcome', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'call_mcp_tool', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-mcp-3',
|
||||
confirmationDetails: {
|
||||
type: 'mcp',
|
||||
title: 'MCP Server Operation',
|
||||
prompt: 'test_mcp',
|
||||
},
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'proceed_always_tool' },
|
||||
});
|
||||
expect(handled).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-mcp-3',
|
||||
confirmed: true,
|
||||
outcome: ToolConfirmationOutcome.ProceedAlwaysTool,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle MCP Server tool ProceedAlwaysAndSave outcome', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'call_mcp_tool', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-mcp-4',
|
||||
confirmationDetails: {
|
||||
type: 'mcp',
|
||||
title: 'MCP Server Operation',
|
||||
prompt: 'test_mcp',
|
||||
},
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'proceed_always_and_save' },
|
||||
});
|
||||
expect(handled).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-mcp-4',
|
||||
confirmed: true,
|
||||
outcome: ToolConfirmationOutcome.ProceedAlwaysAndSave,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should execute without confirmation in YOLO mode and not transition to input-required', async () => {
|
||||
// Enable YOLO mode
|
||||
const yoloConfig = createMockConfig({
|
||||
isEventDrivenSchedulerEnabled: () => true,
|
||||
getApprovalMode: () => ApprovalMode.YOLO,
|
||||
}) as Config;
|
||||
const yoloMessageBus = yoloConfig.getMessageBus();
|
||||
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', yoloConfig, mockEventBus);
|
||||
task.setTaskStateAndPublishUpdate = vi.fn();
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-1',
|
||||
confirmationDetails: { type: 'info', title: 'test', prompt: 'test' },
|
||||
};
|
||||
|
||||
const handler = (yoloMessageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
// Should NOT auto-publish ProceedOnce anymore, because PolicyEngine handles it directly
|
||||
expect(yoloMessageBus.publish).not.toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
}),
|
||||
);
|
||||
|
||||
// Should NOT transition to input-required since it was auto-approved
|
||||
expect(task.setTaskStateAndPublishUpdate).not.toHaveBeenCalledWith(
|
||||
'input-required',
|
||||
expect.anything(),
|
||||
undefined,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle output updates via the message bus', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'executing',
|
||||
liveOutput: 'chunk1',
|
||||
};
|
||||
|
||||
// Simulate MessageBus event
|
||||
// Simulate MessageBus event
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
|
||||
if (!handler) {
|
||||
throw new Error('TOOL_CALLS_UPDATE handler not found');
|
||||
}
|
||||
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall],
|
||||
});
|
||||
|
||||
// Should publish artifact update for output
|
||||
expect(mockEventBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
kind: 'artifact-update',
|
||||
artifact: expect.objectContaining({
|
||||
artifactId: 'tool-1-output',
|
||||
parts: [{ kind: 'text', text: 'chunk1' }],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should complete artifact creation without hanging', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCallId = 'create-file-123';
|
||||
task['_registerToolCall'](toolCallId, 'executing');
|
||||
|
||||
const toolCall = {
|
||||
request: {
|
||||
callId: toolCallId,
|
||||
name: 'writeFile',
|
||||
args: { path: 'test.sh' },
|
||||
},
|
||||
status: 'success',
|
||||
result: { ok: true },
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
handler({ type: MessageBusType.TOOL_CALLS_UPDATE, toolCalls: [toolCall] });
|
||||
|
||||
// The tool should be complete and registered appropriately, eventually
|
||||
// triggering the toolCompletionPromise resolution when all clear.
|
||||
const internalTask = task as unknown as {
|
||||
completedToolCalls: unknown[];
|
||||
pendingToolCalls: Map<string, string>;
|
||||
};
|
||||
expect(internalTask.completedToolCalls.length).toBe(1);
|
||||
expect(internalTask.pendingToolCalls.size).toBe(0);
|
||||
});
|
||||
|
||||
it('should preserve messageId across multiple text chunks to prevent UI duplication', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
// Initialize the ID for the first turn (happens internally upon LLM stream)
|
||||
task.currentAgentMessageId = 'test-id-123';
|
||||
|
||||
// Simulate sending multiple text chunks
|
||||
task._sendTextContent('chunk 1');
|
||||
task._sendTextContent('chunk 2');
|
||||
|
||||
// Both text contents should have been published with the same messageId
|
||||
const textCalls = (mockEventBus.publish as Mock).mock.calls.filter(
|
||||
(call) => call[0].status?.message?.kind === 'message',
|
||||
);
|
||||
expect(textCalls.length).toBe(2);
|
||||
expect(textCalls[0][0].status.message.messageId).toBe('test-id-123');
|
||||
expect(textCalls[1][0].status.message.messageId).toBe('test-id-123');
|
||||
|
||||
// Simulate starting a new turn by calling getAndClearCompletedTools
|
||||
// (which precedes sendCompletedToolsToLlm where a new ID is minted)
|
||||
task.getAndClearCompletedTools();
|
||||
|
||||
// sendCompletedToolsToLlm internally rolls the ID forward.
|
||||
// Simulate what sendCompletedToolsToLlm does:
|
||||
const internalTask = task as unknown as {
|
||||
setTaskStateAndPublishUpdate: (state: string, change: unknown) => void;
|
||||
};
|
||||
internalTask.setTaskStateAndPublishUpdate('working', {});
|
||||
|
||||
// Simulate what sendCompletedToolsToLlm does: generate a new UUID for the next turn
|
||||
task.currentAgentMessageId = 'test-id-456';
|
||||
|
||||
task._sendTextContent('chunk 3');
|
||||
|
||||
const secondTurnCalls = (mockEventBus.publish as Mock).mock.calls.filter(
|
||||
(call) => call[0].status?.message?.messageId === 'test-id-456',
|
||||
);
|
||||
expect(secondTurnCalls.length).toBe(1);
|
||||
expect(secondTurnCalls[0][0].status.message.parts[0].text).toBe('chunk 3');
|
||||
});
|
||||
|
||||
it('should handle parallel tool calls correctly', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const toolCall1 = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-1',
|
||||
confirmationDetails: { type: 'info', title: 'test 1', prompt: 'test 1' },
|
||||
};
|
||||
|
||||
const toolCall2 = {
|
||||
request: { callId: '2', name: 'pwd', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-2',
|
||||
confirmationDetails: { type: 'info', title: 'test 2', prompt: 'test 2' },
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
|
||||
// Publish update for both tool calls simultaneously
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall1, toolCall2],
|
||||
});
|
||||
|
||||
// Confirm first tool call
|
||||
const handled1 = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '1', outcome: 'proceed_once' },
|
||||
});
|
||||
expect(handled1).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-1',
|
||||
confirmed: true,
|
||||
}),
|
||||
);
|
||||
|
||||
// Confirm second tool call
|
||||
const handled2 = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: '2', outcome: 'cancel' },
|
||||
});
|
||||
expect(handled2).toBe(true);
|
||||
expect(messageBus.publish).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-2',
|
||||
confirmed: false,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should wait for executing tools before transitioning to input-required state', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
task.setTaskStateAndPublishUpdate = vi.fn();
|
||||
|
||||
// Register tool 1 as executing
|
||||
task['_registerToolCall']('1', 'executing');
|
||||
|
||||
const toolCall1 = {
|
||||
request: { callId: '1', name: 'ls', args: {} },
|
||||
status: 'executing',
|
||||
};
|
||||
|
||||
const toolCall2 = {
|
||||
request: { callId: '2', name: 'pwd', args: {} },
|
||||
status: 'awaiting_approval',
|
||||
correlationId: 'corr-2',
|
||||
confirmationDetails: { type: 'info', title: 'test 2', prompt: 'test 2' },
|
||||
};
|
||||
|
||||
const handler = (messageBus.subscribe as Mock).mock.calls.find(
|
||||
(call: unknown[]) => call[0] === MessageBusType.TOOL_CALLS_UPDATE,
|
||||
)?.[1];
|
||||
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall1, toolCall2],
|
||||
});
|
||||
|
||||
// Should NOT transition to input-required yet
|
||||
expect(task.setTaskStateAndPublishUpdate).not.toHaveBeenCalledWith(
|
||||
'input-required',
|
||||
expect.anything(),
|
||||
undefined,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
|
||||
// Complete tool 1
|
||||
const toolCall1Complete = {
|
||||
...toolCall1,
|
||||
status: 'success',
|
||||
result: { ok: true },
|
||||
};
|
||||
|
||||
handler({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [toolCall1Complete, toolCall2],
|
||||
});
|
||||
|
||||
// Now it should transition
|
||||
expect(task.setTaskStateAndPublishUpdate).toHaveBeenCalledWith(
|
||||
'input-required',
|
||||
expect.anything(),
|
||||
undefined,
|
||||
undefined,
|
||||
true,
|
||||
);
|
||||
});
|
||||
|
||||
it('should ignore confirmations for unknown tool calls', async () => {
|
||||
// @ts-expect-error - Calling private constructor
|
||||
const task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
const handled = await (
|
||||
task as unknown as {
|
||||
_handleToolConfirmationPart: (part: unknown) => Promise<boolean>;
|
||||
}
|
||||
)._handleToolConfirmationPart({
|
||||
kind: 'data',
|
||||
data: { callId: 'unknown-id', outcome: 'proceed_once' },
|
||||
});
|
||||
|
||||
// Should return false for unhandled tool call
|
||||
expect(handled).toBe(false);
|
||||
|
||||
// Should not publish anything to the message bus
|
||||
expect(messageBus.publish).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -4,25 +4,14 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
describe,
|
||||
it,
|
||||
expect,
|
||||
vi,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
type Mock,
|
||||
} from 'vitest';
|
||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||
import { Task } from './task.js';
|
||||
import {
|
||||
GeminiEventType,
|
||||
ApprovalMode,
|
||||
ToolConfirmationOutcome,
|
||||
type Config,
|
||||
type ToolCallRequestInfo,
|
||||
type GitService,
|
||||
type CompletedToolCall,
|
||||
type ToolCall,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { createMockConfig } from '../utils/testing_utils.js';
|
||||
import type { ExecutionEventBus, RequestContext } from '@a2a-js/sdk/server';
|
||||
@@ -389,188 +378,6 @@ describe('Task', () => {
|
||||
);
|
||||
});
|
||||
|
||||
describe('_schedulerToolCallsUpdate', () => {
|
||||
let task: Task;
|
||||
type SpyInstance = ReturnType<typeof vi.spyOn>;
|
||||
let setTaskStateAndPublishUpdateSpy: SpyInstance;
|
||||
let mockConfig: Config;
|
||||
let mockEventBus: ExecutionEventBus;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = createMockConfig() as Config;
|
||||
mockEventBus = {
|
||||
publish: vi.fn(),
|
||||
on: vi.fn(),
|
||||
off: vi.fn(),
|
||||
once: vi.fn(),
|
||||
removeAllListeners: vi.fn(),
|
||||
finished: vi.fn(),
|
||||
};
|
||||
|
||||
// @ts-expect-error - Calling private constructor
|
||||
task = new Task('task-id', 'context-id', mockConfig, mockEventBus);
|
||||
|
||||
// Spy on the method we want to check calls for
|
||||
setTaskStateAndPublishUpdateSpy = vi.spyOn(
|
||||
task,
|
||||
'setTaskStateAndPublishUpdate',
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('should set state to input-required when a tool is awaiting approval and none are executing', () => {
|
||||
const toolCalls = [
|
||||
{ request: { callId: '1' }, status: 'awaiting_approval' },
|
||||
] as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
// The last call should be the final state update
|
||||
expect(setTaskStateAndPublishUpdateSpy).toHaveBeenLastCalledWith(
|
||||
'input-required',
|
||||
{ kind: 'state-change' },
|
||||
undefined,
|
||||
undefined,
|
||||
true, // final: true
|
||||
);
|
||||
});
|
||||
|
||||
it('should NOT set state to input-required if a tool is awaiting approval but another is executing', () => {
|
||||
const toolCalls = [
|
||||
{ request: { callId: '1' }, status: 'awaiting_approval' },
|
||||
{ request: { callId: '2' }, status: 'executing' },
|
||||
] as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
// It will be called for status updates, but not with final: true
|
||||
const finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find(
|
||||
(call) => call[4] === true,
|
||||
);
|
||||
expect(finalCall).toBeUndefined();
|
||||
});
|
||||
|
||||
it('should set state to input-required once an executing tool finishes, leaving one awaiting approval', () => {
|
||||
const initialToolCalls = [
|
||||
{ request: { callId: '1' }, status: 'awaiting_approval' },
|
||||
{ request: { callId: '2' }, status: 'executing' },
|
||||
] as ToolCall[];
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(initialToolCalls);
|
||||
|
||||
// No final call yet
|
||||
let finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find(
|
||||
(call) => call[4] === true,
|
||||
);
|
||||
expect(finalCall).toBeUndefined();
|
||||
|
||||
// Now, the executing tool finishes. The scheduler would call _resolveToolCall for it.
|
||||
// @ts-expect-error - Calling private method
|
||||
task._resolveToolCall('2');
|
||||
|
||||
// Then another update comes in for the awaiting tool (e.g., a re-check)
|
||||
const subsequentToolCalls = [
|
||||
{ request: { callId: '1' }, status: 'awaiting_approval' },
|
||||
] as ToolCall[];
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(subsequentToolCalls);
|
||||
|
||||
// NOW we should get the final call
|
||||
finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find(
|
||||
(call) => call[4] === true,
|
||||
);
|
||||
expect(finalCall).toBeDefined();
|
||||
expect(finalCall?.[0]).toBe('input-required');
|
||||
});
|
||||
|
||||
it('should NOT set state to input-required if skipFinalTrueAfterInlineEdit is true', () => {
|
||||
task.skipFinalTrueAfterInlineEdit = true;
|
||||
const toolCalls = [
|
||||
{ request: { callId: '1' }, status: 'awaiting_approval' },
|
||||
] as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
const finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find(
|
||||
(call) => call[4] === true,
|
||||
);
|
||||
expect(finalCall).toBeUndefined();
|
||||
});
|
||||
|
||||
describe('auto-approval', () => {
|
||||
it('should auto-approve tool calls when autoExecute is true', () => {
|
||||
task.autoExecute = true;
|
||||
const onConfirmSpy = vi.fn();
|
||||
const toolCalls = [
|
||||
{
|
||||
request: { callId: '1' },
|
||||
status: 'awaiting_approval',
|
||||
confirmationDetails: {
|
||||
type: 'edit',
|
||||
onConfirm: onConfirmSpy,
|
||||
},
|
||||
},
|
||||
] as unknown as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
expect(onConfirmSpy).toHaveBeenCalledWith(
|
||||
ToolConfirmationOutcome.ProceedOnce,
|
||||
);
|
||||
});
|
||||
|
||||
it('should auto-approve tool calls when approval mode is YOLO', () => {
|
||||
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO);
|
||||
task.autoExecute = false;
|
||||
const onConfirmSpy = vi.fn();
|
||||
const toolCalls = [
|
||||
{
|
||||
request: { callId: '1' },
|
||||
status: 'awaiting_approval',
|
||||
confirmationDetails: {
|
||||
type: 'edit',
|
||||
onConfirm: onConfirmSpy,
|
||||
},
|
||||
},
|
||||
] as unknown as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
expect(onConfirmSpy).toHaveBeenCalledWith(
|
||||
ToolConfirmationOutcome.ProceedOnce,
|
||||
);
|
||||
});
|
||||
|
||||
it('should NOT auto-approve when autoExecute is false and mode is not YOLO', () => {
|
||||
task.autoExecute = false;
|
||||
(mockConfig.getApprovalMode as Mock).mockReturnValue(
|
||||
ApprovalMode.DEFAULT,
|
||||
);
|
||||
const onConfirmSpy = vi.fn();
|
||||
const toolCalls = [
|
||||
{
|
||||
request: { callId: '1' },
|
||||
status: 'awaiting_approval',
|
||||
confirmationDetails: { onConfirm: onConfirmSpy },
|
||||
},
|
||||
] as unknown as ToolCall[];
|
||||
|
||||
// @ts-expect-error - Calling private method
|
||||
task._schedulerToolCallsUpdate(toolCalls);
|
||||
|
||||
expect(onConfirmSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('currentPromptId and promptCount', () => {
|
||||
it('should correctly initialize and update promptId and promptCount', async () => {
|
||||
const mockConfig = createMockConfig();
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
*/
|
||||
|
||||
import {
|
||||
CoreToolScheduler,
|
||||
Scheduler,
|
||||
type GeminiClient,
|
||||
GeminiEventType,
|
||||
ToolConfirmationOutcome,
|
||||
@@ -34,6 +34,8 @@ import {
|
||||
isSubagentProgress,
|
||||
EDIT_TOOL_NAMES,
|
||||
processRestorableToolCalls,
|
||||
MessageBusType,
|
||||
type ToolCallsUpdateMessage,
|
||||
} from '@google/gemini-cli-core';
|
||||
import {
|
||||
type ExecutionEventBus,
|
||||
@@ -66,51 +68,33 @@ import type { PartUnion, Part as genAiPart } from '@google/genai';
|
||||
|
||||
type UnionKeys<T> = T extends T ? keyof T : never;
|
||||
|
||||
type ConfirmationType = ToolCallConfirmationDetails['type'];
|
||||
|
||||
const VALID_CONFIRMATION_TYPES: readonly ConfirmationType[] = [
|
||||
'edit',
|
||||
'exec',
|
||||
'mcp',
|
||||
'info',
|
||||
'ask_user',
|
||||
'exit_plan_mode',
|
||||
] as const;
|
||||
|
||||
function isToolCallConfirmationDetails(
|
||||
value: unknown,
|
||||
): value is ToolCallConfirmationDetails {
|
||||
if (
|
||||
typeof value !== 'object' ||
|
||||
value === null ||
|
||||
!('onConfirm' in value) ||
|
||||
typeof value.onConfirm !== 'function' ||
|
||||
!('type' in value) ||
|
||||
typeof value.type !== 'string'
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
return (VALID_CONFIRMATION_TYPES as readonly string[]).includes(value.type);
|
||||
}
|
||||
|
||||
export class Task {
|
||||
id: string;
|
||||
contextId: string;
|
||||
scheduler: CoreToolScheduler;
|
||||
scheduler: Scheduler;
|
||||
config: Config;
|
||||
geminiClient: GeminiClient;
|
||||
pendingToolConfirmationDetails: Map<string, ToolCallConfirmationDetails>;
|
||||
pendingCorrelationIds: Map<string, string> = new Map();
|
||||
taskState: TaskState;
|
||||
eventBus?: ExecutionEventBus;
|
||||
completedToolCalls: CompletedToolCall[];
|
||||
processedToolCallIds: Set<string> = new Set();
|
||||
skipFinalTrueAfterInlineEdit = false;
|
||||
modelInfo?: string;
|
||||
currentPromptId: string | undefined;
|
||||
currentAgentMessageId = uuidv4();
|
||||
promptCount = 0;
|
||||
autoExecute: boolean;
|
||||
private get isYoloMatch(): boolean {
|
||||
return (
|
||||
this.autoExecute || this.config.getApprovalMode() === ApprovalMode.YOLO
|
||||
);
|
||||
}
|
||||
|
||||
// For tool waiting logic
|
||||
private pendingToolCalls: Map<string, string> = new Map(); //toolCallId --> status
|
||||
private toolsAlreadyConfirmed: Set<string> = new Set();
|
||||
private toolCompletionPromise?: Promise<void>;
|
||||
private toolCompletionNotifier?: {
|
||||
resolve: () => void;
|
||||
@@ -127,7 +111,9 @@ export class Task {
|
||||
this.id = id;
|
||||
this.contextId = contextId;
|
||||
this.config = config;
|
||||
this.scheduler = this.createScheduler();
|
||||
|
||||
this.scheduler = this.setupEventDrivenScheduler();
|
||||
|
||||
this.geminiClient = this.config.getGeminiClient();
|
||||
this.pendingToolConfirmationDetails = new Map();
|
||||
this.taskState = 'submitted';
|
||||
@@ -227,7 +213,7 @@ export class Task {
|
||||
logger.info(
|
||||
`[Task] Waiting for ${this.pendingToolCalls.size} pending tool(s)...`,
|
||||
);
|
||||
return this.toolCompletionPromise;
|
||||
await this.toolCompletionPromise;
|
||||
}
|
||||
|
||||
cancelPendingTools(reason: string): void {
|
||||
@@ -240,6 +226,9 @@ export class Task {
|
||||
this.toolCompletionNotifier.reject(new Error(reason));
|
||||
}
|
||||
this.pendingToolCalls.clear();
|
||||
this.pendingCorrelationIds.clear();
|
||||
|
||||
this.scheduler.cancelAll();
|
||||
// Reset the promise for any future operations, ensuring it's in a clean state.
|
||||
this._resetToolCompletionPromise();
|
||||
}
|
||||
@@ -252,7 +241,7 @@ export class Task {
|
||||
kind: 'message',
|
||||
role,
|
||||
parts: [{ kind: 'text', text }],
|
||||
messageId: uuidv4(),
|
||||
messageId: role === 'agent' ? this.currentAgentMessageId : uuidv4(),
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
};
|
||||
@@ -384,104 +373,153 @@ export class Task {
|
||||
this.eventBus?.publish(artifactEvent);
|
||||
}
|
||||
|
||||
private async _schedulerAllToolCallsComplete(
|
||||
completedToolCalls: CompletedToolCall[],
|
||||
): Promise<void> {
|
||||
logger.info(
|
||||
'[Task] All tool calls completed by scheduler (batch):',
|
||||
completedToolCalls.map((tc) => tc.request.callId),
|
||||
);
|
||||
this.completedToolCalls.push(...completedToolCalls);
|
||||
completedToolCalls.forEach((tc) => {
|
||||
this._resolveToolCall(tc.request.callId);
|
||||
private messageBusListener?: (message: ToolCallsUpdateMessage) => void;
|
||||
|
||||
private setupEventDrivenScheduler(): Scheduler {
|
||||
const messageBus = this.config.getMessageBus();
|
||||
const scheduler = new Scheduler({
|
||||
schedulerId: this.id,
|
||||
context: this.config,
|
||||
messageBus,
|
||||
getPreferredEditor: () => DEFAULT_GUI_EDITOR,
|
||||
});
|
||||
|
||||
this.messageBusListener = this.handleEventDrivenToolCallsUpdate.bind(this);
|
||||
messageBus.subscribe<ToolCallsUpdateMessage>(
|
||||
MessageBusType.TOOL_CALLS_UPDATE,
|
||||
this.messageBusListener,
|
||||
);
|
||||
|
||||
return scheduler;
|
||||
}
|
||||
|
||||
private _schedulerToolCallsUpdate(toolCalls: ToolCall[]): void {
|
||||
logger.info(
|
||||
'[Task] Scheduler tool calls updated:',
|
||||
toolCalls.map((tc) => `${tc.request.callId} (${tc.status})`),
|
||||
);
|
||||
dispose(): void {
|
||||
if (this.messageBusListener) {
|
||||
this.config
|
||||
.getMessageBus()
|
||||
.unsubscribe(MessageBusType.TOOL_CALLS_UPDATE, this.messageBusListener);
|
||||
this.messageBusListener = undefined;
|
||||
}
|
||||
|
||||
// Update state and send continuous, non-final updates
|
||||
toolCalls.forEach((tc) => {
|
||||
const previousStatus = this.pendingToolCalls.get(tc.request.callId);
|
||||
const hasChanged = previousStatus !== tc.status;
|
||||
this.scheduler.dispose();
|
||||
}
|
||||
|
||||
// Resolve tool call if it has reached a terminal state
|
||||
if (['success', 'error', 'cancelled'].includes(tc.status)) {
|
||||
this._resolveToolCall(tc.request.callId);
|
||||
} else {
|
||||
// This will update the map
|
||||
this._registerToolCall(tc.request.callId, tc.status);
|
||||
}
|
||||
|
||||
if (tc.status === 'awaiting_approval' && tc.confirmationDetails) {
|
||||
const details = tc.confirmationDetails;
|
||||
if (isToolCallConfirmationDetails(details)) {
|
||||
this.pendingToolConfirmationDetails.set(tc.request.callId, details);
|
||||
}
|
||||
}
|
||||
|
||||
// Only send an update if the status has actually changed.
|
||||
if (hasChanged) {
|
||||
const coderAgentMessage: CoderAgentMessage =
|
||||
tc.status === 'awaiting_approval'
|
||||
? { kind: CoderAgentEvent.ToolCallConfirmationEvent }
|
||||
: { kind: CoderAgentEvent.ToolCallUpdateEvent };
|
||||
const message = this.toolStatusMessage(tc, this.id, this.contextId);
|
||||
|
||||
const event = this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
coderAgentMessage,
|
||||
message,
|
||||
false, // Always false for these continuous updates
|
||||
);
|
||||
this.eventBus?.publish(event);
|
||||
}
|
||||
});
|
||||
|
||||
if (
|
||||
this.autoExecute ||
|
||||
this.config.getApprovalMode() === ApprovalMode.YOLO
|
||||
) {
|
||||
logger.info(
|
||||
'[Task] ' +
|
||||
(this.autoExecute ? '' : 'YOLO mode enabled. ') +
|
||||
'Auto-approving all tool calls.',
|
||||
);
|
||||
toolCalls.forEach((tc: ToolCall) => {
|
||||
if (tc.status === 'awaiting_approval' && tc.confirmationDetails) {
|
||||
const details = tc.confirmationDetails;
|
||||
if (isToolCallConfirmationDetails(details)) {
|
||||
// eslint-disable-next-line @typescript-eslint/no-floating-promises
|
||||
details.onConfirm(ToolConfirmationOutcome.ProceedOnce);
|
||||
this.pendingToolConfirmationDetails.delete(tc.request.callId);
|
||||
}
|
||||
}
|
||||
});
|
||||
private handleEventDrivenToolCallsUpdate(
|
||||
event: ToolCallsUpdateMessage,
|
||||
): void {
|
||||
if (event.type !== MessageBusType.TOOL_CALLS_UPDATE) {
|
||||
return;
|
||||
}
|
||||
|
||||
const allPendingStatuses = Array.from(this.pendingToolCalls.values());
|
||||
const isAwaitingApproval = allPendingStatuses.some(
|
||||
(status) => status === 'awaiting_approval',
|
||||
);
|
||||
const isExecuting = allPendingStatuses.some(
|
||||
(status) => status === 'executing',
|
||||
);
|
||||
const toolCalls = event.toolCalls;
|
||||
|
||||
toolCalls.forEach((tc) => {
|
||||
this.handleEventDrivenToolCall(tc);
|
||||
});
|
||||
|
||||
this.checkInputRequiredState();
|
||||
}
|
||||
|
||||
private handleEventDrivenToolCall(tc: ToolCall): void {
|
||||
const callId = tc.request.callId;
|
||||
|
||||
// Do not process events for tools that have already been finalized.
|
||||
// This prevents duplicate completions if the state manager emits a snapshot containing
|
||||
// already resolved tools whose IDs were removed from pendingToolCalls.
|
||||
if (
|
||||
this.processedToolCallIds.has(callId) ||
|
||||
this.completedToolCalls.some((c) => c.request.callId === callId)
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
const previousStatus = this.pendingToolCalls.get(callId);
|
||||
const hasChanged = previousStatus !== tc.status;
|
||||
|
||||
// 1. Handle Output
|
||||
if (tc.status === 'executing' && tc.liveOutput) {
|
||||
this._schedulerOutputUpdate(callId, tc.liveOutput);
|
||||
}
|
||||
|
||||
// 2. Handle terminal states
|
||||
if (
|
||||
tc.status === 'success' ||
|
||||
tc.status === 'error' ||
|
||||
tc.status === 'cancelled'
|
||||
) {
|
||||
this.toolsAlreadyConfirmed.delete(callId);
|
||||
if (hasChanged) {
|
||||
logger.info(
|
||||
`[Task] Tool call ${callId} completed with status: ${tc.status}`,
|
||||
);
|
||||
this.completedToolCalls.push(tc);
|
||||
this._resolveToolCall(callId);
|
||||
}
|
||||
} else {
|
||||
// Keep track of pending tools
|
||||
this._registerToolCall(callId, tc.status);
|
||||
}
|
||||
|
||||
// 3. Handle Confirmation Stash
|
||||
if (tc.status === 'awaiting_approval' && tc.confirmationDetails) {
|
||||
const details = tc.confirmationDetails;
|
||||
|
||||
if (tc.correlationId) {
|
||||
this.pendingCorrelationIds.set(callId, tc.correlationId);
|
||||
}
|
||||
|
||||
this.pendingToolConfirmationDetails.set(callId, {
|
||||
...details,
|
||||
onConfirm: async () => {},
|
||||
} as ToolCallConfirmationDetails);
|
||||
}
|
||||
|
||||
// 4. Publish Status Updates to A2A event bus
|
||||
if (hasChanged) {
|
||||
const coderAgentMessage: CoderAgentMessage =
|
||||
tc.status === 'awaiting_approval'
|
||||
? { kind: CoderAgentEvent.ToolCallConfirmationEvent }
|
||||
: { kind: CoderAgentEvent.ToolCallUpdateEvent };
|
||||
|
||||
const message = this.toolStatusMessage(tc, this.id, this.contextId);
|
||||
const statusUpdate = this._createStatusUpdateEvent(
|
||||
this.taskState,
|
||||
coderAgentMessage,
|
||||
message,
|
||||
false,
|
||||
);
|
||||
this.eventBus?.publish(statusUpdate);
|
||||
}
|
||||
}
|
||||
|
||||
private checkInputRequiredState(): void {
|
||||
if (this.isYoloMatch) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 6. Handle Input Required State
|
||||
let isAwaitingApproval = false;
|
||||
let isExecuting = false;
|
||||
|
||||
for (const [callId, status] of this.pendingToolCalls.entries()) {
|
||||
if (status === 'executing' || status === 'scheduled') {
|
||||
isExecuting = true;
|
||||
} else if (
|
||||
status === 'awaiting_approval' &&
|
||||
!this.toolsAlreadyConfirmed.has(callId)
|
||||
) {
|
||||
isAwaitingApproval = true;
|
||||
}
|
||||
}
|
||||
|
||||
// The turn is complete and requires user input if at least one tool
|
||||
// is waiting for the user's decision, and no other tool is actively
|
||||
// running in the background.
|
||||
if (
|
||||
isAwaitingApproval &&
|
||||
!isExecuting &&
|
||||
!this.skipFinalTrueAfterInlineEdit
|
||||
) {
|
||||
this.skipFinalTrueAfterInlineEdit = false;
|
||||
const wasAlreadyInputRequired = this.taskState === 'input-required';
|
||||
|
||||
// We don't need to send another message, just a final status update.
|
||||
this.setTaskStateAndPublishUpdate(
|
||||
'input-required',
|
||||
{ kind: CoderAgentEvent.StateChangeEvent },
|
||||
@@ -489,18 +527,13 @@ export class Task {
|
||||
undefined,
|
||||
/*final*/ true,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private createScheduler(): CoreToolScheduler {
|
||||
const scheduler = new CoreToolScheduler({
|
||||
outputUpdateHandler: this._schedulerOutputUpdate.bind(this),
|
||||
onAllToolCallsComplete: this._schedulerAllToolCallsComplete.bind(this),
|
||||
onToolCallsUpdate: this._schedulerToolCallsUpdate.bind(this),
|
||||
getPreferredEditor: () => DEFAULT_GUI_EDITOR,
|
||||
config: this.config,
|
||||
});
|
||||
return scheduler;
|
||||
// Unblock waitForPendingTools to correctly end the executor loop and release the HTTP response stream.
|
||||
// The IDE client will open a new stream with the confirmation reply.
|
||||
if (!wasAlreadyInputRequired && this.toolCompletionNotifier) {
|
||||
this.toolCompletionNotifier.resolve();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private _pickFields<
|
||||
@@ -713,7 +746,16 @@ export class Task {
|
||||
};
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange);
|
||||
|
||||
await this.scheduler.schedule(updatedRequests, abortSignal);
|
||||
// Pre-register tools to ensure waitForPendingTools sees them as pending
|
||||
// before the async scheduler enqueues them and fires the event bus update.
|
||||
for (const req of updatedRequests) {
|
||||
if (!this.pendingToolCalls.has(req.callId)) {
|
||||
this._registerToolCall(req.callId, 'scheduled');
|
||||
}
|
||||
}
|
||||
|
||||
// Fire and forget so we don't block the executor loop before waitForPendingTools can be called
|
||||
void this.scheduler.schedule(updatedRequests, abortSignal);
|
||||
}
|
||||
|
||||
async acceptAgentMessage(event: ServerGeminiStreamEvent): Promise<void> {
|
||||
@@ -839,9 +881,15 @@ export class Task {
|
||||
) {
|
||||
return false;
|
||||
}
|
||||
if (!part.data['outcome']) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const callId = part.data['callId'];
|
||||
const outcomeString = part.data['outcome'];
|
||||
|
||||
this.toolsAlreadyConfirmed.add(callId);
|
||||
|
||||
let confirmationOutcome: ToolConfirmationOutcome | undefined;
|
||||
|
||||
if (outcomeString === 'proceed_once') {
|
||||
@@ -854,6 +902,8 @@ export class Task {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlwaysServer;
|
||||
} else if (outcomeString === 'proceed_always_tool') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlwaysTool;
|
||||
} else if (outcomeString === 'proceed_always_and_save') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ProceedAlwaysAndSave;
|
||||
} else if (outcomeString === 'modify_with_editor') {
|
||||
confirmationOutcome = ToolConfirmationOutcome.ModifyWithEditor;
|
||||
} else {
|
||||
@@ -864,8 +914,9 @@ export class Task {
|
||||
}
|
||||
|
||||
const confirmationDetails = this.pendingToolConfirmationDetails.get(callId);
|
||||
const correlationId = this.pendingCorrelationIds.get(callId);
|
||||
|
||||
if (!confirmationDetails) {
|
||||
if (!confirmationDetails && !correlationId) {
|
||||
logger.warn(
|
||||
`[Task] Received tool confirmation for unknown or already processed callId: ${callId}`,
|
||||
);
|
||||
@@ -887,24 +938,35 @@ export class Task {
|
||||
// This will trigger the scheduler to continue or cancel the specific tool.
|
||||
// The scheduler's onToolCallsUpdate will then reflect the new state (e.g., executing or cancelled).
|
||||
|
||||
// If `edit` tool call, pass updated payload if presesent
|
||||
if (confirmationDetails.type === 'edit') {
|
||||
const newContent = part.data['newContent'];
|
||||
const payload =
|
||||
typeof newContent === 'string'
|
||||
? ({ newContent } as ToolConfirmationPayload)
|
||||
: undefined;
|
||||
this.skipFinalTrueAfterInlineEdit = !!payload;
|
||||
try {
|
||||
// If `edit` tool call, pass updated payload if present
|
||||
const newContent = part.data['newContent'];
|
||||
const payload =
|
||||
confirmationDetails?.type === 'edit' && typeof newContent === 'string'
|
||||
? ({ newContent } as ToolConfirmationPayload)
|
||||
: undefined;
|
||||
this.skipFinalTrueAfterInlineEdit = !!payload;
|
||||
|
||||
try {
|
||||
if (correlationId) {
|
||||
await this.config.getMessageBus().publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId,
|
||||
confirmed:
|
||||
confirmationOutcome !== ToolConfirmationOutcome.Cancel &&
|
||||
confirmationOutcome !==
|
||||
ToolConfirmationOutcome.ModifyWithEditor,
|
||||
outcome: confirmationOutcome,
|
||||
payload,
|
||||
});
|
||||
} else if (confirmationDetails?.onConfirm) {
|
||||
// Fallback for legacy callback-based confirmation
|
||||
await confirmationDetails.onConfirm(confirmationOutcome, payload);
|
||||
} finally {
|
||||
// Once confirmationDetails.onConfirm finishes (or fails) with a payload,
|
||||
// reset skipFinalTrueAfterInlineEdit so that external callers receive
|
||||
// their call has been completed.
|
||||
this.skipFinalTrueAfterInlineEdit = false;
|
||||
}
|
||||
} else {
|
||||
await confirmationDetails.onConfirm(confirmationOutcome);
|
||||
} finally {
|
||||
// Once confirmation payload is sent or callback finishes,
|
||||
// reset skipFinalTrueAfterInlineEdit so that external callers receive
|
||||
// their call has been completed.
|
||||
this.skipFinalTrueAfterInlineEdit = false;
|
||||
}
|
||||
} finally {
|
||||
if (gcpProject) {
|
||||
@@ -920,6 +982,7 @@ export class Task {
|
||||
// Note !== ToolConfirmationOutcome.ModifyWithEditor does not work!
|
||||
if (confirmationOutcome !== 'modify_with_editor') {
|
||||
this.pendingToolConfirmationDetails.delete(callId);
|
||||
this.pendingCorrelationIds.delete(callId);
|
||||
}
|
||||
|
||||
// If outcome is Cancel, scheduler should update status to 'cancelled', which then resolves the tool.
|
||||
@@ -953,6 +1016,9 @@ export class Task {
|
||||
|
||||
getAndClearCompletedTools(): CompletedToolCall[] {
|
||||
const tools = [...this.completedToolCalls];
|
||||
for (const tool of tools) {
|
||||
this.processedToolCallIds.add(tool.request.callId);
|
||||
}
|
||||
this.completedToolCalls = [];
|
||||
return tools;
|
||||
}
|
||||
@@ -1013,6 +1079,7 @@ export class Task {
|
||||
};
|
||||
// Set task state to working as we are about to call LLM
|
||||
this.setTaskStateAndPublishUpdate('working', stateChange);
|
||||
this.currentAgentMessageId = uuidv4();
|
||||
yield* this.geminiClient.sendMessageStream(
|
||||
llmParts,
|
||||
aborted,
|
||||
@@ -1034,6 +1101,10 @@ export class Task {
|
||||
if (confirmationHandled) {
|
||||
anyConfirmationHandled = true;
|
||||
// If a confirmation was handled, the scheduler will now run the tool (or cancel it).
|
||||
// We resolve the toolCompletionPromise manually in checkInputRequiredState
|
||||
// to break the original execution loop, so we must reset it here so the
|
||||
// new loop correctly awaits the tool's final execution.
|
||||
this._resetToolCompletionPromise();
|
||||
// We don't send anything to the LLM for this part.
|
||||
// The subsequent tool execution will eventually lead to resolveToolCall.
|
||||
continue;
|
||||
@@ -1048,6 +1119,7 @@ export class Task {
|
||||
if (hasContentForLlm) {
|
||||
this.currentPromptId =
|
||||
this.config.getSessionId() + '########' + this.promptCount++;
|
||||
this.currentAgentMessageId = uuidv4();
|
||||
logger.info('[Task] Sending new parts to LLM.');
|
||||
const stateChange: StateChange = {
|
||||
kind: CoderAgentEvent.StateChangeEvent,
|
||||
@@ -1093,7 +1165,6 @@ export class Task {
|
||||
if (content === '') {
|
||||
return;
|
||||
}
|
||||
logger.info('[Task] Sending text content to event bus.');
|
||||
const message = this._createTextMessage(content);
|
||||
const textContent: TextContent = {
|
||||
kind: CoderAgentEvent.TextContentEvent,
|
||||
@@ -1125,7 +1196,7 @@ export class Task {
|
||||
data: content,
|
||||
} as Part,
|
||||
],
|
||||
messageId: uuidv4(),
|
||||
messageId: this.currentAgentMessageId,
|
||||
taskId: this.id,
|
||||
contextId: this.contextId,
|
||||
};
|
||||
|
||||
@@ -91,6 +91,15 @@ describe('loadConfig', () => {
|
||||
expect(fetchAdminControlsOnce).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should pass clientName as a2a-server to Config', async () => {
|
||||
await loadConfig(mockSettings, mockExtensionLoader, taskId);
|
||||
expect(Config).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
clientName: 'a2a-server',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
describe('when admin controls experiment is enabled', () => {
|
||||
beforeEach(() => {
|
||||
// We need to cast to any here to modify the mock implementation
|
||||
|
||||
@@ -62,6 +62,7 @@ export async function loadConfig(
|
||||
|
||||
const configParams: ConfigParameters = {
|
||||
sessionId: taskId,
|
||||
clientName: 'a2a-server',
|
||||
model: PREVIEW_GEMINI_MODEL,
|
||||
embeddingModel: DEFAULT_GEMINI_EMBEDDING_MODEL,
|
||||
sandbox: undefined, // Sandbox might not be relevant for a server-side agent
|
||||
|
||||
@@ -37,6 +37,9 @@ export interface Settings {
|
||||
showMemoryUsage?: boolean;
|
||||
checkpointing?: CheckpointingSettings;
|
||||
folderTrust?: boolean;
|
||||
general?: {
|
||||
previewFeatures?: boolean;
|
||||
};
|
||||
|
||||
// Git-aware file filtering settings
|
||||
fileFiltering?: {
|
||||
|
||||
@@ -65,7 +65,12 @@ vi.mock('../utils/logger.js', () => ({
|
||||
}));
|
||||
|
||||
let config: Config;
|
||||
const getToolRegistrySpy = vi.fn().mockReturnValue(ApprovalMode.DEFAULT);
|
||||
const getToolRegistrySpy = vi.fn().mockReturnValue({
|
||||
getTool: vi.fn(),
|
||||
getAllToolNames: vi.fn().mockReturnValue([]),
|
||||
getAllTools: vi.fn().mockReturnValue([]),
|
||||
getToolsByServer: vi.fn().mockReturnValue([]),
|
||||
});
|
||||
const getApprovalModeSpy = vi.fn();
|
||||
const getShellExecutionConfigSpy = vi.fn();
|
||||
const getExtensionsSpy = vi.fn();
|
||||
|
||||
@@ -21,6 +21,7 @@ import {
|
||||
type Config,
|
||||
type Storage,
|
||||
NoopSandboxManager,
|
||||
type ToolRegistry,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js';
|
||||
import { expect, vi } from 'vitest';
|
||||
@@ -31,6 +32,10 @@ export function createMockConfig(
|
||||
const tmpDir = tmpdir();
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const mockConfig = {
|
||||
get toolRegistry(): ToolRegistry {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
return (this as unknown as Config).getToolRegistry();
|
||||
},
|
||||
getToolRegistry: vi.fn().mockReturnValue({
|
||||
getTool: vi.fn(),
|
||||
getAllToolNames: vi.fn().mockReturnValue([]),
|
||||
|
||||
Reference in New Issue
Block a user