From 8ab74ef1bb4bd1c475596088d9d3b52e0a9c5ca7 Mon Sep 17 00:00:00 2001 From: "N. Taylor Mullen" Date: Mon, 2 Jun 2025 22:30:52 -0700 Subject: [PATCH] Refactor: Use config.getGeminiClient() for GeminiClient instantiation (#715) --- packages/cli/src/nonInteractiveCli.test.ts | 1 + packages/cli/src/nonInteractiveCli.ts | 3 +-- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 27 ++++++++++++------- packages/cli/src/ui/hooks/useGeminiStream.ts | 2 +- packages/core/src/core/client.ts | 4 +-- packages/core/src/tools/edit.test.ts | 8 ++++++ packages/core/src/tools/edit.ts | 2 +- packages/core/src/tools/write-file.test.ts | 6 +++++ packages/core/src/tools/write-file.ts | 2 +- 9 files changed, 39 insertions(+), 16 deletions(-) diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index c071dedc..97e4416a 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -51,6 +51,7 @@ describe('runNonInteractive', () => { mockConfig = { getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry), + getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient), } as unknown as Config; mockProcessStdoutWrite = vi.fn().mockImplementation(() => true); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 7505c736..45c9ce2d 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -6,7 +6,6 @@ import { Config, - GeminiClient, ToolCallRequestInfo, executeToolCall, ToolRegistry, @@ -39,7 +38,7 @@ export async function runNonInteractive( config: Config, input: string, ): Promise { - const geminiClient = new GeminiClient(config); + const geminiClient = config.getGeminiClient(); const toolRegistry: ToolRegistry = await config.getToolRegistry(); const chat = await geminiClient.getChat(); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index d46fab9e..4dfea5e0 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -25,20 +25,20 @@ const mockSendMessageStream = vi .mockReturnValue((async function* () {})()); const mockStartChat = vi.fn(); -vi.mock('@gemini-code/core', async (importOriginal) => { - const actualCoreModule = (await importOriginal()) as any; - const MockedGeminiClientClass = vi.fn().mockImplementation(function ( - this: any, - _config: any, - ) { +const MockedGeminiClientClass = vi.hoisted(() => + vi.fn().mockImplementation(function (this: any, _config: any) { // _config this.startChat = mockStartChat; this.sendMessageStream = mockSendMessageStream; - }); + }), +); + +vi.mock('@gemini-code/core', async (importOriginal) => { + const actualCoreModule = (await importOriginal()) as any; return { ...(actualCoreModule || {}), - GeminiClient: MockedGeminiClientClass, - // GeminiChat will be from actualCoreModule if it exists, otherwise undefined + GeminiClient: MockedGeminiClientClass, // Export the class for type checking or other direct uses + Config: actualCoreModule.Config, // Ensure Config is passed through }; }); @@ -235,6 +235,14 @@ describe('useGeminiStream', () => { mockAddItem = vi.fn(); mockSetShowHelp = vi.fn(); + // Define the mock for getGeminiClient + const mockGetGeminiClient = vi.fn().mockImplementation(() => { + // MockedGeminiClientClass is defined in the module scope by the previous change. + // It will use the mockStartChat and mockSendMessageStream that are managed within beforeEach. + const clientInstance = new MockedGeminiClientClass(mockConfig); + return clientInstance; + }); + mockConfig = { apiKey: 'test-api-key', model: 'gemini-pro', @@ -258,6 +266,7 @@ describe('useGeminiStream', () => { getToolRegistry: vi.fn( () => ({ getToolSchemaList: vi.fn(() => []) }) as any, ), + getGeminiClient: mockGetGeminiClient, } as unknown as Config; mockOnDebugMessage = vi.fn(); mockHandleSlashCommand = vi.fn().mockReturnValue(false); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 284709cf..c50d4c43 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -145,7 +145,7 @@ export const useGeminiStream = ( setInitError(null); if (!geminiClientRef.current) { try { - geminiClientRef.current = new GeminiClient(config); + geminiClientRef.current = config.getGeminiClient(); } catch (error: unknown) { const errorMsg = `Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`; setInitError(errorMsg); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index fcad1ef0..311f53bf 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -25,6 +25,7 @@ import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js'; import { reportError } from '../utils/errorReporting.js'; import { GeminiChat } from './geminiChat.js'; import { retryWithBackoff } from '../utils/retry.js'; +import { getErrorMessage } from '../utils/errors.js'; export class GeminiClient { private chat: Promise; @@ -158,8 +159,7 @@ export class GeminiClient { history, 'startChat', ); - const message = error instanceof Error ? error.message : 'Unknown error.'; - throw new Error(`Failed to initialize chat: ${message}`); + throw new Error(`Failed to initialize chat: ${getErrorMessage(error)}`); } } diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 3b93708a..87ae3b6f 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -39,7 +39,15 @@ describe('EditTool', () => { rootDir = path.join(tempDir, 'root'); fs.mkdirSync(rootDir); + // The client instance that EditTool will use + const mockClientInstanceWithGenerateJson = { + generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted + }; + mockConfig = { + getGeminiClient: vi + .fn() + .mockReturnValue(mockClientInstanceWithGenerateJson), getTargetDir: () => rootDir, getApprovalMode: vi.fn(() => false), setApprovalMode: vi.fn(), diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index b2f648f8..4a337faa 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -114,7 +114,7 @@ Expectation for required parameters: ); this.config = config; this.rootDirectory = path.resolve(this.config.getTargetDir()); - this.client = new GeminiClient(this.config); + this.client = config.getGeminiClient(); } /** diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index c94edfd1..4646f30a 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -53,6 +53,7 @@ const mockConfigInternal = { getTargetDir: () => rootDir, getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), setApprovalMode: vi.fn(), + getGeminiClient: vi.fn(), // Initialize as a plain mock function getApiKey: () => 'test-key', getModel: () => 'test-model', getSandbox: () => false, @@ -97,6 +98,11 @@ describe('WriteFileTool', () => { ) as Mocked; vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance); + // Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient + mockConfigInternal.getGeminiClient.mockReturnValue( + mockGeminiClientInstance, + ); + tool = new WriteFileTool(mockConfig); // Reset mocks before each test diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 2e04a10a..57bfe8ed 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -77,7 +77,7 @@ export class WriteFileTool extends BaseTool { }, ); - this.client = new GeminiClient(this.config); + this.client = this.config.getGeminiClient(); } private isWithinRoot(pathToCheck: string): boolean {