From eec5d5ebf8393a7d2da59656c9196a55470df327 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Sun, 4 Jan 2026 14:59:35 -0500 Subject: [PATCH] feat(core): restore MessageBus optionality for soft migration (Phase 1) (#15774) --- packages/a2a-server/src/http/app.test.ts | 3 +- .../cli/src/ui/hooks/useToolScheduler.test.ts | 2 +- packages/core/src/index.ts | 3 -- packages/core/src/test-utils/mock-tool.ts | 39 ++++++++++---- packages/core/src/tools/edit.test.ts | 18 +++++-- packages/core/src/tools/glob.test.ts | 3 +- packages/core/src/tools/grep.test.ts | 18 +++++-- packages/core/src/tools/ls.test.ts | 3 +- .../src/tools/message-bus-integration.test.ts | 51 +++---------------- packages/core/src/tools/read-file.test.ts | 5 +- .../core/src/tools/read-many-files.test.ts | 5 +- packages/core/src/tools/tool-registry.ts | 18 ++++--- packages/core/src/tools/write-file.test.ts | 8 ++- packages/core/src/tools/write-file.ts | 3 +- packages/core/src/utils/editCorrector.test.ts | 5 +- packages/core/src/utils/tool-utils.test.ts | 3 +- 16 files changed, 105 insertions(+), 82 deletions(-) diff --git a/packages/a2a-server/src/http/app.test.ts b/packages/a2a-server/src/http/app.test.ts index 6ff4b37867..f427bdfe63 100644 --- a/packages/a2a-server/src/http/app.test.ts +++ b/packages/a2a-server/src/http/app.test.ts @@ -35,7 +35,8 @@ import { createStreamMessageRequest, createMockConfig, } from '../utils/testing_utils.js'; -import { MockTool } from '@google/gemini-cli-core'; +// Import MockTool from specific path to avoid vitest dependency in main core bundle +import { MockTool } from '@google/gemini-cli-core/src/test-utils/mock-tool.js'; import type { Command, CommandContext } from '../commands/types.js'; const mockToolConfirmationFn = async () => diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 69767198c6..7015e6afea 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -31,10 +31,10 @@ import { DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, ToolConfirmationOutcome, ApprovalMode, - MockTool, HookSystem, PREVIEW_GEMINI_MODEL, } from '@google/gemini-cli-core'; +import { MockTool } from '@google/gemini-cli-core/src/test-utils/mock-tool.js'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; import { ToolCallStatus } from '../types.js'; diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 246d5604e6..4cf9df113a 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -157,9 +157,6 @@ export { Storage } from './config/storage.js'; // Export hooks system export * from './hooks/index.js'; -// Export test utils -export * from './test-utils/index.js'; - // Export hook types export * from './hooks/types.js'; diff --git a/packages/core/src/test-utils/mock-tool.ts b/packages/core/src/test-utils/mock-tool.ts index cdfc649d46..8f5d44d7fd 100644 --- a/packages/core/src/test-utils/mock-tool.ts +++ b/packages/core/src/test-utils/mock-tool.ts @@ -18,6 +18,8 @@ import { BaseToolInvocation, Kind, } from '../tools/tools.js'; +import { createMockMessageBus } from './mock-message-bus.js'; +import type { MessageBus } from '../confirmation-bus/message-bus.js'; interface MockToolOptions { name: string; @@ -35,6 +37,7 @@ interface MockToolOptions { updateOutput?: (output: string) => void, ) => Promise; params?: object; + messageBus?: MessageBus; } class MockToolInvocation extends BaseToolInvocation< @@ -44,8 +47,9 @@ class MockToolInvocation extends BaseToolInvocation< constructor( private readonly tool: MockTool, params: { [key: string]: unknown }, + messageBus: MessageBus, ) { - super(params); + super(params, messageBus, tool.name, tool.displayName); } execute( @@ -96,6 +100,7 @@ export class MockTool extends BaseDeclarativeTool< options.params, options.isOutputMarkdown ?? false, options.canUpdateOutput ?? false, + options.messageBus ?? createMockMessageBus(), ); if (options.shouldConfirmExecute) { @@ -115,10 +120,11 @@ export class MockTool extends BaseDeclarativeTool< } } - protected createInvocation(params: { - [key: string]: unknown; - }): ToolInvocation<{ [key: string]: unknown }, ToolResult> { - return new MockToolInvocation(this, params); + protected createInvocation( + params: { [key: string]: unknown }, + messageBus: MessageBus, + ): ToolInvocation<{ [key: string]: unknown }, ToolResult> { + return new MockToolInvocation(this, params, messageBus); } } @@ -138,8 +144,9 @@ export class MockModifiableToolInvocation extends BaseToolInvocation< constructor( private readonly tool: MockModifiableTool, params: Record, + messageBus: MessageBus, ) { - super(params); + super(params, messageBus, tool.name, tool.displayName); } async execute(_abortSignal: AbortSignal): Promise { @@ -189,10 +196,19 @@ export class MockModifiableTool shouldConfirm = true; constructor(name = 'mockModifiableTool') { - super(name, name, 'A mock modifiable tool for testing.', Kind.Other, { - type: 'object', - properties: { param: { type: 'string' } }, - }); + super( + name, + name, + 'A mock modifiable tool for testing.', + Kind.Other, + { + type: 'object', + properties: { param: { type: 'string' } }, + }, + true, + false, + createMockMessageBus(), + ); } getModifyContext( @@ -212,7 +228,8 @@ export class MockModifiableTool protected createInvocation( params: Record, + messageBus: MessageBus, ): ToolInvocation, ToolResult> { - return new MockModifiableToolInvocation(this, params); + return new MockModifiableToolInvocation(this, params, messageBus); } } diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 52e3d8702a..ca1505a2c4 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -59,6 +59,10 @@ import { ApprovalMode } from '../policy/types.js'; import type { Content, Part, SchemaUnion } from '@google/genai'; import { StandardFileSystemService } from '../services/fileSystemService.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; +import { + createMockMessageBus, + getMockMessageBusInstance, +} from '../test-utils/mock-message-bus.js'; describe('EditTool', () => { let tool: EditTool; @@ -177,7 +181,9 @@ describe('EditTool', () => { }, ); - tool = new EditTool(mockConfig); + const bus = createMockMessageBus(); + getMockMessageBusInstance(bus).defaultToolDecision = 'ask_user'; + tool = new EditTool(mockConfig, bus); }); afterEach(() => { @@ -1029,7 +1035,10 @@ describe('EditTool', () => { it('should use windows-style path examples on windows', () => { vi.spyOn(process, 'platform', 'get').mockReturnValue('win32'); - const tool = new EditTool({} as unknown as Config); + const tool = new EditTool( + {} as unknown as Config, + createMockMessageBus(), + ); const schema = tool.schema; expect( (schema.parametersJsonSchema as EditFileParameterSchema).properties @@ -1040,7 +1049,10 @@ describe('EditTool', () => { it('should use unix-style path examples on non-windows platforms', () => { vi.spyOn(process, 'platform', 'get').mockReturnValue('linux'); - const tool = new EditTool({} as unknown as Config); + const tool = new EditTool( + {} as unknown as Config, + createMockMessageBus(), + ); const schema = tool.schema; expect( (schema.parametersJsonSchema as EditFileParameterSchema).properties diff --git a/packages/core/src/tools/glob.test.ts b/packages/core/src/tools/glob.test.ts index 4c65480089..381a47dc12 100644 --- a/packages/core/src/tools/glob.test.ts +++ b/packages/core/src/tools/glob.test.ts @@ -16,6 +16,7 @@ import type { Config } from '../config/config.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { ToolErrorType } from './tool-error.js'; import * as glob from 'glob'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; vi.mock('glob', { spy: true }); @@ -43,7 +44,7 @@ describe('GlobTool', () => { // Create a unique root directory for each test run tempRootDir = await fs.mkdtemp(path.join(os.tmpdir(), 'glob-tool-root-')); await fs.writeFile(path.join(tempRootDir, '.git'), ''); // Fake git repo - globTool = new GlobTool(mockConfig); + globTool = new GlobTool(mockConfig, createMockMessageBus()); // Create some test files and directories within this root // Top-level files diff --git a/packages/core/src/tools/grep.test.ts b/packages/core/src/tools/grep.test.ts index ded3d0d311..131042fe3b 100644 --- a/packages/core/src/tools/grep.test.ts +++ b/packages/core/src/tools/grep.test.ts @@ -14,6 +14,7 @@ import type { Config } from '../config/config.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { ToolErrorType } from './tool-error.js'; import * as glob from 'glob'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; vi.mock('glob', { spy: true }); @@ -47,7 +48,7 @@ describe('GrepTool', () => { beforeEach(async () => { tempRootDir = await fs.mkdtemp(path.join(os.tmpdir(), 'grep-tool-root-')); - grepTool = new GrepTool(mockConfig); + grepTool = new GrepTool(mockConfig, createMockMessageBus()); // Create some test files and directories await fs.writeFile( @@ -270,7 +271,10 @@ describe('GrepTool', () => { }), } as unknown as Config; - const multiDirGrepTool = new GrepTool(multiDirConfig); + const multiDirGrepTool = new GrepTool( + multiDirConfig, + createMockMessageBus(), + ); const params: GrepToolParams = { pattern: 'world' }; const invocation = multiDirGrepTool.build(params); const result = await invocation.execute(abortSignal); @@ -323,7 +327,10 @@ describe('GrepTool', () => { }), } as unknown as Config; - const multiDirGrepTool = new GrepTool(multiDirConfig); + const multiDirGrepTool = new GrepTool( + multiDirConfig, + createMockMessageBus(), + ); // Search only in the 'sub' directory of the first workspace const params: GrepToolParams = { pattern: 'world', dir_path: 'sub' }; @@ -385,7 +392,10 @@ describe('GrepTool', () => { }), } as unknown as Config; - const multiDirGrepTool = new GrepTool(multiDirConfig); + const multiDirGrepTool = new GrepTool( + multiDirConfig, + createMockMessageBus(), + ); const params: GrepToolParams = { pattern: 'testPattern' }; const invocation = multiDirGrepTool.build(params); expect(invocation.getDescription()).toBe( diff --git a/packages/core/src/tools/ls.test.ts b/packages/core/src/tools/ls.test.ts index 63a84aea43..06e8da264e 100644 --- a/packages/core/src/tools/ls.test.ts +++ b/packages/core/src/tools/ls.test.ts @@ -13,6 +13,7 @@ import type { Config } from '../config/config.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { ToolErrorType } from './tool-error.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; describe('LSTool', () => { let lsTool: LSTool; @@ -39,7 +40,7 @@ describe('LSTool', () => { }), } as unknown as Config; - lsTool = new LSTool(mockConfig); + lsTool = new LSTool(mockConfig, createMockMessageBus()); }); afterEach(async () => { diff --git a/packages/core/src/tools/message-bus-integration.test.ts b/packages/core/src/tools/message-bus-integration.test.ts index 2ee38b0d22..3f40468001 100644 --- a/packages/core/src/tools/message-bus-integration.test.ts +++ b/packages/core/src/tools/message-bus-integration.test.ts @@ -56,22 +56,19 @@ class TestToolInvocation extends BaseToolInvocation { override async shouldConfirmExecute( abortSignal: AbortSignal, ): Promise { - // This conditional is here to allow testing of the case where there is no message bus. - if (this.messageBus) { - const decision = await this.getMessageBusDecision(abortSignal); - if (decision === 'ALLOW') { - return false; - } - if (decision === 'DENY') { - throw new Error('Tool execution denied by policy'); - } + const decision = await this.getMessageBusDecision(abortSignal); + if (decision === 'ALLOW') { + return false; + } + if (decision === 'DENY') { + throw new Error('Tool execution denied by policy'); } return false; } } class TestTool extends BaseDeclarativeTool { - constructor(messageBus?: MessageBus) { + constructor(messageBus: MessageBus) { super( 'test-tool', 'Test Tool', @@ -90,7 +87,7 @@ class TestTool extends BaseDeclarativeTool { ); } - protected createInvocation(params: TestParams, messageBus?: MessageBus) { + protected createInvocation(params: TestParams, messageBus: MessageBus) { return new TestToolInvocation(params, messageBus); } } @@ -220,16 +217,6 @@ describe('Message Bus Integration', () => { ); }); - it('should fall back to default behavior when no message bus', async () => { - const tool = new TestTool(); // No message bus - const invocation = tool.build({ testParam: 'test-value' }); - - const result = await invocation.shouldConfirmExecute( - new AbortController().signal, - ); - expect(result).toBe(false); - }); - it('should ignore responses with wrong correlation ID', async () => { vi.useFakeTimers(); @@ -260,28 +247,6 @@ describe('Message Bus Integration', () => { }); }); - describe('Backward Compatibility', () => { - it('should work with existing tools that do not use message bus', async () => { - const tool = new TestTool(); // No message bus - const invocation = tool.build({ testParam: 'test-value' }); - - // Should execute normally - const result = await invocation.execute(new AbortController().signal); - expect(result.testValue).toBe('test-value'); - expect(result.llmContent).toBe('Executed with test-value'); - }); - - it('should work with tools that have message bus but use default confirmation', async () => { - const tool = new TestTool(messageBus); - const invocation = tool.build({ testParam: 'test-value' }); - - // Should execute normally even with message bus available - const result = await invocation.execute(new AbortController().signal); - expect(result.testValue).toBe('test-value'); - expect(result.llmContent).toBe('Executed with test-value'); - }); - }); - describe('Error Handling', () => { it('should handle message bus publish errors gracefully', async () => { const tool = new TestTool(messageBus); diff --git a/packages/core/src/tools/read-file.test.ts b/packages/core/src/tools/read-file.test.ts index a01c4a5261..0194e18288 100644 --- a/packages/core/src/tools/read-file.test.ts +++ b/packages/core/src/tools/read-file.test.ts @@ -17,6 +17,7 @@ import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; vi.mock('../telemetry/loggers.js', () => ({ logFileOperation: vi.fn(), @@ -46,7 +47,7 @@ describe('ReadFileTool', () => { }, isInteractive: () => false, } as unknown as Config; - tool = new ReadFileTool(mockConfigInstance); + tool = new ReadFileTool(mockConfigInstance, createMockMessageBus()); }); afterEach(async () => { @@ -438,7 +439,7 @@ describe('ReadFileTool', () => { getProjectTempDir: () => path.join(tempRootDir, '.temp'), }, } as unknown as Config; - tool = new ReadFileTool(mockConfigInstance); + tool = new ReadFileTool(mockConfigInstance, createMockMessageBus()); }); it('should throw error if path is ignored by a .geminiignore pattern', async () => { diff --git a/packages/core/src/tools/read-many-files.test.ts b/packages/core/src/tools/read-many-files.test.ts index 20a06763c2..e092b6d6a5 100644 --- a/packages/core/src/tools/read-many-files.test.ts +++ b/packages/core/src/tools/read-many-files.test.ts @@ -21,6 +21,7 @@ import { DEFAULT_FILE_EXCLUDES, } from '../utils/ignorePatterns.js'; import * as glob from 'glob'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; vi.mock('glob', { spy: true }); @@ -90,7 +91,7 @@ describe('ReadManyFilesTool', () => { }), isInteractive: () => false, } as Partial as Config; - tool = new ReadManyFilesTool(mockConfig); + tool = new ReadManyFilesTool(mockConfig, createMockMessageBus()); mockReadFileFn = mockControl.mockReadFile; mockReadFileFn.mockReset(); @@ -505,7 +506,7 @@ describe('ReadManyFilesTool', () => { }), isInteractive: () => false, } as Partial as Config; - tool = new ReadManyFilesTool(mockConfig); + tool = new ReadManyFilesTool(mockConfig, createMockMessageBus()); fs.writeFileSync(path.join(tempDir1, 'file1.txt'), 'Content1'); fs.writeFileSync(path.join(tempDir2, 'file2.txt'), 'Content2'); diff --git a/packages/core/src/tools/tool-registry.ts b/packages/core/src/tools/tool-registry.ts index 40b9fc16dc..60e8080fb0 100644 --- a/packages/core/src/tools/tool-registry.ts +++ b/packages/core/src/tools/tool-registry.ts @@ -172,7 +172,7 @@ Signal: Signal number or \`(none)\` if no signal was received. protected createInvocation( params: ToolParams, - _messageBus?: MessageBus, + messageBus?: MessageBus, _toolName?: string, _displayName?: string, ): ToolInvocation { @@ -181,7 +181,7 @@ Signal: Signal number or \`(none)\` if no signal was received. this.originalName, this.name, params, - _messageBus, + messageBus, ); } } @@ -194,11 +194,8 @@ export class ToolRegistry { private config: Config; private messageBus?: MessageBus; - constructor(config: Config) { + constructor(config: Config, messageBus?: MessageBus) { this.config = config; - } - - setMessageBus(messageBus: MessageBus): void { this.messageBus = messageBus; } @@ -206,6 +203,15 @@ export class ToolRegistry { return this.messageBus; } + /** + * @deprecated migration only - will be removed in PR 3 (Enforcement) + * TODO: DELETE ME in PR 3. This is a temporary shim to allow for soft migration + * of tools while the core infrastructure is updated to require a MessageBus at birth. + */ + setMessageBus(messageBus: MessageBus): void { + this.messageBus = messageBus; + } + /** * Registers a tool definition. * diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 9d7e32a0ac..cd5436e7be 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -41,6 +41,10 @@ import { StandardFileSystemService } from '../services/fileSystemService.js'; import type { DiffUpdateResult } from '../ide/ide-client.js'; import { IdeClient } from '../ide/ide-client.js'; import { WorkspaceContext } from '../utils/workspaceContext.js'; +import { + createMockMessageBus, + getMockMessageBusInstance, +} from '../test-utils/mock-message-bus.js'; const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); @@ -150,7 +154,9 @@ describe('WriteFileTool', () => { mockBaseLlmClientInstance, ); - tool = new WriteFileTool(mockConfig); + const bus = createMockMessageBus(); + getMockMessageBusInstance(bus).defaultToolDecision = 'ask_user'; + tool = new WriteFileTool(mockConfig, bus); // Reset mocks before each test mockConfigInternal.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT); diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index bd832ac537..c3d7e39409 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -475,11 +475,12 @@ export class WriteFileTool protected createInvocation( params: WriteFileToolParams, + messageBus?: MessageBus, ): ToolInvocation { return new WriteFileToolInvocation( this.config, params, - this.messageBus, + messageBus ?? this.messageBus, this.name, this.displayName, ); diff --git a/packages/core/src/utils/editCorrector.test.ts b/packages/core/src/utils/editCorrector.test.ts index 913c2d9a63..02b8c1a3fb 100644 --- a/packages/core/src/utils/editCorrector.test.ts +++ b/packages/core/src/utils/editCorrector.test.ts @@ -164,7 +164,10 @@ describe('editCorrector', () => { const abortSignal = new AbortController().signal; beforeEach(() => { - mockToolRegistry = new ToolRegistry({} as Config) as Mocked; + mockToolRegistry = new ToolRegistry( + {} as Config, + {} as any, + ) as Mocked; const configParams = { apiKey: 'test-api-key', model: 'test-model', diff --git a/packages/core/src/utils/tool-utils.test.ts b/packages/core/src/utils/tool-utils.test.ts index 861281b2da..822d0c2999 100644 --- a/packages/core/src/utils/tool-utils.test.ts +++ b/packages/core/src/utils/tool-utils.test.ts @@ -8,6 +8,7 @@ import { expect, describe, it } from 'vitest'; import { doesToolInvocationMatch, getToolSuggestion } from './tool-utils.js'; import type { AnyToolInvocation, Config } from '../index.js'; import { ReadFileTool } from '../tools/read-file.js'; +import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; describe('getToolSuggestion', () => { it('should suggest the top N closest tool names for a typo', () => { @@ -83,7 +84,7 @@ describe('doesToolInvocationMatch', () => { }); describe('for non-shell tools', () => { - const readFileTool = new ReadFileTool({} as Config); + const readFileTool = new ReadFileTool({} as Config, createMockMessageBus()); const invocation = { params: { file: 'test.txt' }, } as AnyToolInvocation;