diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 135516db6f..28ebc47d71 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -582,6 +582,15 @@ Logging in with Google... Please restart Gemini CLI to continue. const cancelHandlerRef = useRef<() => void>(() => {}); + const getPreferredEditor = useCallback( + () => settings.merged.general?.preferredEditor as EditorType, + [settings.merged.general?.preferredEditor], + ); + + const onCancelSubmit = useCallback(() => { + cancelHandlerRef.current(); + }, []); + const { streamingState, submitQuery, @@ -601,13 +610,13 @@ Logging in with Google... Please restart Gemini CLI to continue. setDebugMessage, handleSlashCommand, shellModeActive, - () => settings.merged.general?.preferredEditor as EditorType, + getPreferredEditor, onAuthError, performMemoryRefresh, modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError, refreshStatic, - () => cancelHandlerRef.current(), + onCancelSubmit, setEmbeddedShellFocused, terminalWidth, terminalHeight, diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts new file mode 100644 index 0000000000..b3fcfad8b7 --- /dev/null +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts @@ -0,0 +1,85 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { CoreToolScheduler } from '@google/gemini-cli-core'; +import type { Config } from '@google/gemini-cli-core'; +import { renderHook } from '@testing-library/react'; +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { useReactToolScheduler } from './useReactToolScheduler.js'; + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + CoreToolScheduler: vi.fn(), + }; +}); + +const mockCoreToolScheduler = vi.mocked(CoreToolScheduler); + +describe('useReactToolScheduler', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('only creates one instance of CoreToolScheduler even if props change', () => { + const onComplete = vi.fn(); + const getPreferredEditor = vi.fn(); + const onEditorClose = vi.fn(); + const config = {} as Config; + + const { rerender } = renderHook( + (props) => + useReactToolScheduler( + props.onComplete, + props.config, + props.getPreferredEditor, + props.onEditorClose, + ), + { + initialProps: { + onComplete, + config, + getPreferredEditor, + onEditorClose, + }, + }, + ); + + expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1); + + // Rerender with a new onComplete function + const newOnComplete = vi.fn(); + rerender({ + onComplete: newOnComplete, + config, + getPreferredEditor, + onEditorClose, + }); + expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1); + + // Rerender with a new getPreferredEditor function + const newGetPreferredEditor = vi.fn(); + rerender({ + onComplete: newOnComplete, + config, + getPreferredEditor: newGetPreferredEditor, + onEditorClose, + }); + expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1); + + // Rerender with a new onEditorClose function + const newOnEditorClose = vi.fn(); + rerender({ + onComplete: newOnComplete, + config, + getPreferredEditor: newGetPreferredEditor, + onEditorClose: newOnEditorClose, + }); + expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index d78b66e142..883690d79a 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -21,7 +21,7 @@ import type { EditorType, } from '@google/gemini-cli-core'; import { CoreToolScheduler, debugLogger } from '@google/gemini-cli-core'; -import { useCallback, useState, useMemo } from 'react'; +import { useCallback, useState, useMemo, useEffect, useRef } from 'react'; import type { HistoryItemToolGroup, IndividualToolCallDisplay, @@ -72,6 +72,23 @@ export function useReactToolScheduler( TrackedToolCall[] >([]); + // Store callbacks in refs to keep them up-to-date without causing re-renders. + const onCompleteRef = useRef(onComplete); + const getPreferredEditorRef = useRef(getPreferredEditor); + const onEditorCloseRef = useRef(onEditorClose); + + useEffect(() => { + onCompleteRef.current = onComplete; + }, [onComplete]); + + useEffect(() => { + getPreferredEditorRef.current = getPreferredEditor; + }, [getPreferredEditor]); + + useEffect(() => { + onEditorCloseRef.current = onEditorClose; + }, [onEditorClose]); + const outputUpdateHandler: OutputUpdateHandler = useCallback( (toolCallId, outputChunk) => { setToolCallsForDisplay((prevCalls) => @@ -89,9 +106,9 @@ export function useReactToolScheduler( const allToolCallsCompleteHandler: AllToolCallsCompleteHandler = useCallback( async (completedToolCalls) => { - await onComplete(completedToolCalls); + await onCompleteRef.current(completedToolCalls); }, - [onComplete], + [], ); const toolCallsUpdateHandler: ToolCallsUpdateHandler = useCallback( @@ -130,24 +147,29 @@ export function useReactToolScheduler( [setToolCallsForDisplay], ); + const stableGetPreferredEditor = useCallback( + () => getPreferredEditorRef.current(), + [], + ); + const stableOnEditorClose = useCallback(() => onEditorCloseRef.current(), []); + const scheduler = useMemo( () => new CoreToolScheduler({ outputUpdateHandler, onAllToolCallsComplete: allToolCallsCompleteHandler, onToolCallsUpdate: toolCallsUpdateHandler, - getPreferredEditor, + getPreferredEditor: stableGetPreferredEditor, config, - onEditorClose, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any), + onEditorClose: stableOnEditorClose, + }), [ config, outputUpdateHandler, allToolCallsCompleteHandler, toolCallsUpdateHandler, - getPreferredEditor, - onEditorClose, + stableGetPreferredEditor, + stableOnEditorClose, ], ); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 31265690dc..6c76f4aa5c 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -10,7 +10,6 @@ import type { ToolCallConfirmationDetails, ToolResult, ToolResultDisplay, - ToolRegistry, EditorType, Config, ToolConfirmationPayload, @@ -332,7 +331,6 @@ interface CoreToolSchedulerOptions { } export class CoreToolScheduler { - private toolRegistry: ToolRegistry; private toolCalls: ToolCall[] = []; private outputUpdateHandler?: OutputUpdateHandler; private onAllToolCallsComplete?: AllToolCallsCompleteHandler; @@ -351,7 +349,6 @@ export class CoreToolScheduler { constructor(options: CoreToolSchedulerOptions) { this.config = options.config; - this.toolRegistry = options.config.getToolRegistry(); this.outputUpdateHandler = options.outputUpdateHandler; this.onAllToolCallsComplete = options.onAllToolCallsComplete; this.onToolCallsUpdate = options.onToolCallsUpdate; @@ -603,7 +600,7 @@ export class CoreToolScheduler { * @returns A suggestion string like " Did you mean 'tool'?" or " Did you mean one of: 'tool1', 'tool2'?", or an empty string if no suggestions are found. */ private getToolSuggestion(unknownToolName: string, topN = 3): string { - const allToolNames = this.toolRegistry.getAllToolNames(); + const allToolNames = this.config.getToolRegistry().getAllToolNames(); const matches = allToolNames.map((toolName) => ({ name: toolName, @@ -680,7 +677,9 @@ export class CoreToolScheduler { const newToolCalls: ToolCall[] = requestsToProcess.map( (reqInfo): ToolCall => { - const toolInstance = this.toolRegistry.getTool(reqInfo.name); + const toolInstance = this.config + .getToolRegistry() + .getTool(reqInfo.name); if (!toolInstance) { const suggestion = this.getToolSuggestion(reqInfo.name); const errorMessage = `Tool "${reqInfo.name}" not found in registry. Tools must use the exact names that are registered.${suggestion}`;