refactor: remove unnecessary useRefs (#780)

5 files changed, 46 insertions, 46 deletions

diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx
index 0dc4cb79..01acd2d3 100644
--- a/packages/cli/src/ui/App.test.tsx
+++ b/packages/cli/src/ui/App.test.tsx
@@ -117,6 +117,7 @@ vi.mock('@gemini-code/core', async (importOriginal) => {
         getVertexAI: vi.fn(() => opts.vertexai),
         getShowMemoryUsage: vi.fn(() => opts.showMemoryUsage ?? false),
         getAccessibility: vi.fn(() => opts.accessibility ?? {}),
+        getGeminiClient: vi.fn(() => ({})),
       };
     });
   return {
diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx
index 8dd037a5..ea235be9 100644
--- a/packages/cli/src/ui/App.tsx
+++ b/packages/cli/src/ui/App.tsx
@@ -201,6 +201,7 @@ export const App = ({
 
   const { streamingState, submitQuery, initError, pendingHistoryItems } =
     useGeminiStream(
+      config.getGeminiClient(),
       addItem,
       setShowHelp,
       config,
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index 4dfea5e0..ca489fab 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -304,6 +304,7 @@ describe('useGeminiStream', () => {
 
     const { result, rerender } = renderHook(() =>
       useGeminiStream(
+        mockConfig.getGeminiClient(),
         mockAddItem as unknown as UseHistoryManagerReturn['addItem'],
         mockSetShowHelp,
         mockConfig,
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 0b42c161..3889fbe6 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -65,6 +65,7 @@ enum StreamProcessingStatus {
  * API interaction, and tool call lifecycle.
  */
 export const useGeminiStream = (
+  geminiClient: GeminiClient | null,
   addItem: UseHistoryManagerReturn['addItem'],
   setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
   config: Config,
@@ -76,7 +77,6 @@ export const useGeminiStream = (
 ) => {
   const [initError, setInitError] = useState<string | null>(null);
   const abortControllerRef = useRef<AbortController | null>(null);
-  const geminiClientRef = useRef<GeminiClient | null>(null);
   const [isResponding, setIsResponding] = useState<boolean>(false);
   const [pendingHistoryItemRef, setPendingHistoryItem] =
     useStateAndRef<HistoryItemWithoutId | null>(null);
@@ -142,19 +142,6 @@ export const useGeminiStream = (
     return StreamingState.Idle;
   }, [isResponding, toolCalls]);
 
-  useEffect(() => {
-    setInitError(null);
-    if (!geminiClientRef.current) {
-      try {
-        geminiClientRef.current = config.getGeminiClient();
-      } catch (error: unknown) {
-        const errorMsg = `Failed to initialize client: ${getErrorMessage(error) || 'Unknown error'}`;
-        setInitError(errorMsg);
-        addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
-      }
-    }
-  }, [config, addItem]);
-
   useInput((_input, key) => {
     if (streamingState !== StreamingState.Idle && key.escape) {
       abortControllerRef.current?.abort();
@@ -450,9 +437,7 @@ export const useGeminiStream = (
         return;
       }
 
-      const client = geminiClientRef.current;
-
-      if (!client) {
+      if (!geminiClient) {
         const errorMsg = 'Gemini client is not available.';
         setInitError(errorMsg);
         addItem({ type: MessageType.ERROR, text: errorMsg }, Date.now());
@@ -463,7 +448,7 @@ export const useGeminiStream = (
       setInitError(null);
 
       try {
-        const stream = client.sendMessageStream(queryToSend, abortSignal);
+        const stream = geminiClient.sendMessageStream(queryToSend, abortSignal);
         const processingStatus = await processGeminiStreamEvents(
           stream,
           userMessageTimestamp,
@@ -501,6 +486,7 @@ export const useGeminiStream = (
       addItem,
       setPendingHistoryItem,
       setInitError,
+      geminiClient,
     ],
   );
 
diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
index 8e027ade..f9d2be5f 100644
--- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts
+++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts
@@ -22,7 +22,7 @@ import {
   Status as CoreStatus,
   logToolCall,
 } from '@gemini-code/core';
-import { useCallback, useEffect, useState, useRef } from 'react';
+import { useCallback, useState, useMemo } from 'react';
 import {
   HistoryItemToolGroup,
   IndividualToolCallDisplay,
@@ -73,13 +73,9 @@ export function useReactToolScheduler(
   const [toolCallsForDisplay, setToolCallsForDisplay] = useState<
     TrackedToolCall[]
   >([]);
-  const schedulerRef = useRef<CoreToolScheduler | null>(null);
 
-  useEffect(() => {
-    const outputUpdateHandler: OutputUpdateHandler = (
-      toolCallId,
-      outputChunk,
-    ) => {
+  const outputUpdateHandler: OutputUpdateHandler = useCallback(
+    (toolCallId, outputChunk) => {
       setPendingHistoryItem((prevItem) => {
         if (prevItem?.type === 'tool_group') {
           return {
@@ -104,11 +100,12 @@ export function useReactToolScheduler(
           return tc;
         }),
       );
-    };
+    },
+    [setPendingHistoryItem],
+  );
 
-    const allToolCallsCompleteHandler: AllToolCallsCompleteHandler = (
-      completedToolCalls,
-    ) => {
+  const allToolCallsCompleteHandler: AllToolCallsCompleteHandler = useCallback(
+    (completedToolCalls) => {
       completedToolCalls.forEach((call) => {
         let success = false;
         let errorMessage: string | undefined;
@@ -134,11 +131,12 @@ export function useReactToolScheduler(
         });
       });
       onComplete(completedToolCalls);
-    };
+    },
+    [onComplete],
+  );
 
-    const toolCallsUpdateHandler: ToolCallsUpdateHandler = (
-      updatedCoreToolCalls: ToolCall[],
-    ) => {
+  const toolCallsUpdateHandler: ToolCallsUpdateHandler = useCallback(
+    (updatedCoreToolCalls: ToolCall[]) => {
       setToolCallsForDisplay((prevTrackedCalls) =>
         updatedCoreToolCalls.map((coreTc) => {
           const existingTrackedCall = prevTrackedCalls.find(
@@ -152,27 +150,40 @@ export function useReactToolScheduler(
           return newTrackedCall;
         }),
       );
-    };
+    },
+    [],
+  );
 
-    schedulerRef.current = new CoreToolScheduler({
-      toolRegistry: config.getToolRegistry(),
+  const scheduler = useMemo(
+    () =>
+      new CoreToolScheduler({
+        toolRegistry: config.getToolRegistry(),
+        outputUpdateHandler,
+        onAllToolCallsComplete: allToolCallsCompleteHandler,
+        onToolCallsUpdate: toolCallsUpdateHandler,
+        approvalMode: config.getApprovalMode(),
+      }),
+    [
+      config,
       outputUpdateHandler,
-      onAllToolCallsComplete: allToolCallsCompleteHandler,
-      onToolCallsUpdate: toolCallsUpdateHandler,
-      approvalMode: config.getApprovalMode(),
-    });
-  }, [config, onComplete, setPendingHistoryItem]);
+      allToolCallsCompleteHandler,
+      toolCallsUpdateHandler,
+    ],
+  );
 
   const schedule: ScheduleFn = useCallback(
     async (request: ToolCallRequestInfo | ToolCallRequestInfo[]) => {
-      schedulerRef.current?.schedule(request);
+      scheduler.schedule(request);
     },
-    [],
+    [scheduler],
   );
 
-  const cancel: CancelFn = useCallback((reason: string = 'unspecified') => {
-    schedulerRef.current?.cancelAll(reason);
-  }, []);
+  const cancel: CancelFn = useCallback(
+    (reason: string = 'unspecified') => {
+      scheduler.cancelAll(reason);
+    },
+    [scheduler],
+  );
 
   const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
     (callIdsToMark: string[]) => {