diff --git a/packages/a2a-server/src/agent.test.ts b/packages/a2a-server/src/agent.test.ts index 4c563179..04160d3f 100644 --- a/packages/a2a-server/src/agent.test.ts +++ b/packages/a2a-server/src/agent.test.ts @@ -32,8 +32,8 @@ import { assertUniqueFinalEventIsLast, assertTaskCreationAndWorkingStatus, createStreamMessageRequest, + MockTool, } from './testing_utils.js'; -import { MockTool } from '@google/gemini-cli-core'; const mockToolConfirmationFn = async () => ({}) as unknown as ToolCallConfirmationDetails; @@ -189,10 +189,13 @@ describe('E2E Tests', () => { yield* []; }); - const mockTool = new MockTool({ - name: 'test-tool', - shouldConfirmExecute: vi.fn(mockToolConfirmationFn), - }); + const mockTool = new MockTool( + 'test-tool', + 'Test Tool', + true, + false, + mockToolConfirmationFn, + ); getToolRegistrySpy.mockReturnValue({ getAllTools: vi.fn().mockReturnValue([mockTool]), @@ -281,16 +284,20 @@ describe('E2E Tests', () => { yield* []; }); - const mockTool1 = new MockTool({ - name: 'test-tool-1', - displayName: 'Test Tool 1', - shouldConfirmExecute: vi.fn(mockToolConfirmationFn), - }); - const mockTool2 = new MockTool({ - name: 'test-tool-2', - displayName: 'Test Tool 2', - shouldConfirmExecute: vi.fn(mockToolConfirmationFn), - }); + const mockTool1 = new MockTool( + 'test-tool-1', + 'Test Tool 1', + false, + false, + mockToolConfirmationFn, + ); + const mockTool2 = new MockTool( + 'test-tool-2', + 'Test Tool 2', + false, + false, + mockToolConfirmationFn, + ); getToolRegistrySpy.mockReturnValue({ getAllTools: vi.fn().mockReturnValue([mockTool1, mockTool2]), @@ -397,13 +404,13 @@ describe('E2E Tests', () => { yield* [{ type: 'content', value: 'Tool executed successfully.' }]; }); - const mockTool = new MockTool({ - name: 'test-tool-no-approval', - displayName: 'Test Tool No Approval', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Tool executed successfully.', - returnDisplay: 'Tool executed successfully.', - }), + const mockTool = new MockTool( + 'test-tool-no-approval', + 'Test Tool No Approval', + ); + mockTool.execute.mockResolvedValue({ + llmContent: 'Tool executed successfully.', + returnDisplay: 'Tool executed successfully.', }); getToolRegistrySpy.mockReturnValue({ @@ -528,13 +535,15 @@ describe('E2E Tests', () => { // Set approval mode to yolo getApprovalModeSpy.mockReturnValue(ApprovalMode.YOLO); - const mockTool = new MockTool({ - name: 'test-tool-yolo', - displayName: 'Test Tool YOLO', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Tool executed successfully.', - returnDisplay: 'Tool executed successfully.', - }), + const mockTool = new MockTool( + 'test-tool-yolo', + 'Test Tool YOLO', + false, + false, + ); + mockTool.execute.mockResolvedValue({ + llmContent: 'Tool executed successfully.', + returnDisplay: 'Tool executed successfully.', }); getToolRegistrySpy.mockReturnValue({ diff --git a/packages/a2a-server/src/testing_utils.ts b/packages/a2a-server/src/testing_utils.ts index b9c1de17..bd7ddaaa 100644 --- a/packages/a2a-server/src/testing_utils.ts +++ b/packages/a2a-server/src/testing_utils.ts @@ -9,7 +9,90 @@ import type { TaskStatusUpdateEvent, SendStreamingMessageSuccessResponse, } from '@a2a-js/sdk'; -import { expect } from 'vitest'; +import { + BaseDeclarativeTool, + BaseToolInvocation, + Kind, +} from '@google/gemini-cli-core'; +import type { + ToolCallConfirmationDetails, + ToolResult, + ToolInvocation, +} from '@google/gemini-cli-core'; +import { expect, vi } from 'vitest'; + +export const mockOnUserConfirmForToolConfirmation = vi.fn(); + +export class MockToolInvocation extends BaseToolInvocation { + constructor( + private readonly tool: MockTool, + params: object, + ) { + super(params); + } + + getDescription(): string { + return JSON.stringify(this.params); + } + + override shouldConfirmExecute( + abortSignal: AbortSignal, + ): Promise { + return this.tool.shouldConfirmExecute(this.params, abortSignal); + } + + execute( + signal: AbortSignal, + updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, + ): Promise { + return this.tool.execute( + this.params, + signal, + updateOutput, + terminalColumns, + terminalRows, + ); + } +} + +// TODO: dedup with gemini-cli, add shouldConfirmExecute() support in core +export class MockTool extends BaseDeclarativeTool { + constructor( + name: string, + displayName: string, + canUpdateOutput = false, + isOutputMarkdown = false, + shouldConfirmExecute?: () => Promise, + ) { + super( + name, + displayName, + 'A mock tool for testing', + Kind.Other, + {}, + isOutputMarkdown, + canUpdateOutput, + ); + + if (shouldConfirmExecute) { + this.shouldConfirmExecute.mockImplementation(shouldConfirmExecute); + } else { + // Default to no confirmation needed + this.shouldConfirmExecute.mockResolvedValue(false); + } + } + + execute = vi.fn(); + shouldConfirmExecute = vi.fn(); + + protected createInvocation( + params: object, + ): ToolInvocation { + return new MockToolInvocation(this, params); + } +} export function createStreamMessageRequest( text: string, diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 4584cfa3..0100b551 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -22,13 +22,16 @@ import type { ToolCallResponseInfo, ToolCall, // Import from core Status as ToolCallStatusType, + ToolInvocation, AnyDeclarativeTool, AnyToolInvocation, } from '@google/gemini-cli-core'; import { ToolConfirmationOutcome, ApprovalMode, - MockTool, + Kind, + BaseDeclarativeTool, + BaseToolInvocation, } from '@google/gemini-cli-core'; import type { HistoryItemWithoutId, HistoryItemToolGroup } from '../types.js'; import { ToolCallStatus } from '../types.js'; @@ -61,20 +64,96 @@ const mockConfig = { }), } as unknown as Config; -const mockTool = new MockTool({ name: 'mockTool', displayName: 'Mock Tool' }); -const mockToolWithLiveOutput = new MockTool({ - name: 'mockToolWithLiveOutput', - displayName: 'Mock Tool With Live Output', - description: 'A mock tool for testing', - params: {}, - isOutputMarkdown: true, - canUpdateOutput: true, -}); +class MockToolInvocation extends BaseToolInvocation { + constructor( + private readonly tool: MockTool, + params: object, + ) { + super(params); + } + + getDescription(): string { + return JSON.stringify(this.params); + } + + override shouldConfirmExecute( + abortSignal: AbortSignal, + ): Promise { + return this.tool.shouldConfirmExecute(this.params, abortSignal); + } + + execute( + signal: AbortSignal, + updateOutput?: (output: string) => void, + terminalColumns?: number, + terminalRows?: number, + ): Promise { + return this.tool.execute( + this.params, + signal, + updateOutput, + terminalColumns, + terminalRows, + ); + } +} + +class MockTool extends BaseDeclarativeTool { + constructor( + name: string, + displayName: string, + canUpdateOutput = false, + shouldConfirm = false, + isOutputMarkdown = false, + ) { + super( + name, + displayName, + 'A mock tool for testing', + Kind.Other, + {}, + isOutputMarkdown, + canUpdateOutput, + ); + if (shouldConfirm) { + this.shouldConfirmExecute.mockImplementation( + async (): Promise => ({ + type: 'edit', + title: 'Mock Tool Requires Confirmation', + onConfirm: mockOnUserConfirmForToolConfirmation, + filePath: 'mock', + fileName: 'mockToolRequiresConfirmation.ts', + fileDiff: 'Mock tool requires confirmation', + originalContent: 'Original content', + newContent: 'New content', + }), + ); + } + } + + execute = vi.fn(); + shouldConfirmExecute = vi.fn(); + + protected createInvocation( + params: object, + ): ToolInvocation { + return new MockToolInvocation(this, params); + } +} + +const mockTool = new MockTool('mockTool', 'Mock Tool'); +const mockToolWithLiveOutput = new MockTool( + 'mockToolWithLiveOutput', + 'Mock Tool With Live Output', + true, +); let mockOnUserConfirmForToolConfirmation: Mock; -const mockToolRequiresConfirmation = new MockTool({ - name: 'mockToolRequiresConfirmation', - displayName: 'Mock Tool Requires Confirmation', -}); +const mockToolRequiresConfirmation = new MockTool( + 'mockToolRequiresConfirmation', + 'Mock Tool Requires Confirmation', + false, + true, +); describe('useReactToolScheduler in YOLO Mode', () => { let onComplete: Mock; @@ -638,23 +717,19 @@ describe('useReactToolScheduler', () => { }); it('should schedule and execute multiple tool calls', async () => { - const tool1 = new MockTool({ - name: 'tool1', - displayName: 'Tool 1', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Output 1', - returnDisplay: 'Display 1', - } as ToolResult), - }); + const tool1 = new MockTool('tool1', 'Tool 1'); + tool1.execute.mockResolvedValue({ + llmContent: 'Output 1', + returnDisplay: 'Display 1', + } as ToolResult); + tool1.shouldConfirmExecute.mockResolvedValue(null); - const tool2 = new MockTool({ - name: 'tool2', - displayName: 'Tool 2', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Output 2', - returnDisplay: 'Display 2', - } as ToolResult), - }); + const tool2 = new MockTool('tool2', 'Tool 2'); + tool2.execute.mockResolvedValue({ + llmContent: 'Output 2', + returnDisplay: 'Display 2', + } as ToolResult); + tool2.shouldConfirmExecute.mockResolvedValue(null); mockToolRegistry.getTool.mockImplementation((name) => { if (name === 'tool1') return tool1; @@ -795,12 +870,7 @@ describe('mapToDisplay', () => { args: { foo: 'bar' }, } as any; - const baseTool = new MockTool({ - name: 'testTool', - displayName: 'Test Tool Display', - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), - }); + const baseTool = new MockTool('testTool', 'Test Tool Display'); const baseResponse: ToolCallResponseInfo = { callId: 'testCallId', @@ -958,7 +1028,7 @@ describe('mapToDisplay', () => { expectedStatus: ToolCallStatus.Error, expectedResultDisplay: 'Execution failed display', expectedName: baseTool.displayName, // Changed from baseTool.name - expectedDescription: JSON.stringify(baseRequest.args), + expectedDescription: baseInvocation.getDescription(), }, { name: 'cancelled', @@ -1029,13 +1099,13 @@ describe('mapToDisplay', () => { invocation: baseTool.build(baseRequest.args), response: { ...baseResponse, callId: 'call1' }, } as ToolCall; - const toolForCall2 = new MockTool({ - name: baseTool.name, - displayName: baseTool.displayName, - isOutputMarkdown: true, - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), - }); + const toolForCall2 = new MockTool( + baseTool.name, + baseTool.displayName, + false, + false, + true, + ); const toolCall2: ToolCall = { request: { ...baseRequest, callId: 'call2' }, status: 'executing', diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index be437ec1..8ef7a7de 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -108,6 +108,3 @@ export * from './telemetry/index.js'; export { sessionId } from './utils/session.js'; export * from './utils/browser.js'; export { Storage } from './config/storage.js'; - -// Export test utils -export * from './test-utils/index.js'; diff --git a/packages/core/src/test-utils/index.ts b/packages/core/src/test-utils/index.ts deleted file mode 100644 index 6146d39d..00000000 --- a/packages/core/src/test-utils/index.ts +++ /dev/null @@ -1,7 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -export * from './mock-tool.js'; diff --git a/packages/core/src/test-utils/mock-tool.ts b/packages/core/src/test-utils/mock-tool.ts deleted file mode 100644 index 5fb6a6d2..00000000 --- a/packages/core/src/test-utils/mock-tool.ts +++ /dev/null @@ -1,110 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { vi } from 'vitest'; -import type { - ToolCallConfirmationDetails, - ToolInvocation, - ToolResult, -} from '../tools/tools.js'; -import { - BaseDeclarativeTool, - BaseToolInvocation, - Kind, -} from '../tools/tools.js'; - -type MockToolOptions = { - name: string; - displayName?: string; - description?: string; - canUpdateOutput?: boolean; - isOutputMarkdown?: boolean; - shouldConfirmExecute?: ( - ...args: unknown[] - ) => Promise; - execute?: (...args: unknown[]) => Promise; - params?: object; -}; - -class MockToolInvocation extends BaseToolInvocation< - { [key: string]: unknown }, - ToolResult -> { - constructor( - private readonly tool: MockTool, - params: { [key: string]: unknown }, - ) { - super(params); - } - - execute( - signal: AbortSignal, - updateOutput?: (output: string) => void, - terminalColumns?: number, - terminalRows?: number, - ): Promise { - return this.tool.execute( - this.params, - signal, - updateOutput, - terminalColumns, - terminalRows, - ); - } - - override shouldConfirmExecute( - abortSignal: AbortSignal, - ): Promise { - return this.tool.shouldConfirmExecute(this.params, abortSignal); - } - - getDescription(): string { - return `A mock tool invocation for ${this.tool.name}`; - } -} - -/** - * A highly configurable mock tool for testing purposes. - */ -export class MockTool extends BaseDeclarativeTool< - { [key: string]: unknown }, - ToolResult -> { - execute: (...args: unknown[]) => Promise; - shouldConfirmExecute: ( - ...args: unknown[] - ) => Promise; - - constructor(options: MockToolOptions) { - super( - options.name, - options.displayName ?? options.name, - options.description ?? options.name, - Kind.Other, - options.params, - options.isOutputMarkdown ?? false, - options.canUpdateOutput ?? false, - ); - - if (options.shouldConfirmExecute) { - this.shouldConfirmExecute = options.shouldConfirmExecute; - } else { - this.shouldConfirmExecute = vi.fn().mockResolvedValue(false); - } - - if (options.execute) { - this.execute = options.execute; - } else { - this.execute = vi.fn(); - } - } - - protected createInvocation(params: { - [key: string]: unknown; - }): ToolInvocation<{ [key: string]: unknown }, ToolResult> { - return new MockToolInvocation(this, params); - } -}