refactor(ui): Optimize rendering performance (#8239)

packages/cli/src/ui/AppContainer.tsx

@@ -85,6 +85,7 @@ import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js';
 import { useWorkspaceMigration } from './hooks/useWorkspaceMigration.js';
 import { useSessionStats } from './contexts/SessionContext.js';
 import { useGitBranchName } from './hooks/useGitBranchName.js';
+import { FocusContext } from './contexts/FocusContext.js';
 import type { ExtensionUpdateState } from './state/extensions.js';
 import { checkForAllExtensionUpdates } from '../config/extension.js';
 
@@ -1210,7 +1211,9 @@ Logging in with Google... Please restart Gemini CLI to continue.
               startupWarnings: props.startupWarnings || [],
             }}
           >
-            <App />
+            <FocusContext.Provider value={isFocused}>
+              <App />
+            </FocusContext.Provider>
           </AppContext.Provider>
         </ConfigContext.Provider>
       </UIActionsContext.Provider>

packages/cli/src/ui/components/Composer.tsx

@@ -18,6 +18,7 @@ import { OverflowProvider } from '../contexts/OverflowContext.js';
 import { theme } from '../semantic-colors.js';
 import { isNarrowWidth } from '../utils/isNarrowWidth.js';
 import { useUIState } from '../contexts/UIStateContext.js';
+import { useFocusState } from '../contexts/FocusContext.js';
 import { useUIActions } from '../contexts/UIActionsContext.js';
 import { useVimMode } from '../contexts/VimModeContext.js';
 import { useConfig } from '../contexts/ConfigContext.js';
@@ -32,6 +33,7 @@ export const Composer = () => {
   const config = useConfig();
   const settings = useSettings();
   const uiState = useUIState();
+  const isFocused = useFocusState();
   const uiActions = useUIActions();
   const { vimEnabled, vimMode } = useVimMode();
   const terminalWidth = process.stdout.columns;
@@ -192,7 +194,7 @@ export const Composer = () => {
           setShellModeActive={uiActions.setShellModeActive}
           approvalMode={showAutoAcceptIndicator}
           onEscapePromptChange={uiActions.onEscapePromptChange}
-          focus={uiState.isFocused}
+          focus={isFocused}
           vimHandleInput={uiActions.vimHandleInput}
           isShellFocused={uiState.shellFocused}
           placeholder={

packages/cli/src/ui/components/shared/text-buffer.test.ts

@@ -1495,6 +1495,53 @@ Contrary to popular belief, Lorem Ipsum is not simply random text. It has roots
       expect(stripAnsi('')).toBe('');
     });
   });
+
+  describe('Memoization', () => {
+    it('should keep action references stable across re-renders', () => {
+      // We pass a stable `isValidPath` so that callbacks that depend on it
+      // are not recreated on every render.
+      const isValidPath = () => false;
+      const { result, rerender } = renderHook(() =>
+        useTextBuffer({ viewport, isValidPath }),
+      );
+
+      const initialInsert = result.current.insert;
+      const initialBackspace = result.current.backspace;
+      const initialMove = result.current.move;
+      const initialHandleInput = result.current.handleInput;
+
+      rerender();
+
+      expect(result.current.insert).toBe(initialInsert);
+      expect(result.current.backspace).toBe(initialBackspace);
+      expect(result.current.move).toBe(initialMove);
+      expect(result.current.handleInput).toBe(initialHandleInput);
+    });
+
+    it('should have memoized actions that operate on the latest state', () => {
+      const isValidPath = () => false;
+      const { result } = renderHook(() =>
+        useTextBuffer({ viewport, isValidPath }),
+      );
+
+      // Store a reference to the memoized insert function.
+      const memoizedInsert = result.current.insert;
+
+      // Update the buffer state.
+      act(() => {
+        result.current.insert('hello');
+      });
+      expect(getBufferState(result).text).toBe('hello');
+
+      // Now, call the original memoized function reference.
+      act(() => {
+        memoizedInsert(' world');
+      });
+
+      // It should have operated on the updated state.
+      expect(getBufferState(result).text).toBe('hello world');
+    });
+  });
 });
 
 describe('offsetToLogicalPos', () => {

packages/cli/src/ui/components/shared/text-buffer.ts

@@ -1984,71 +1984,135 @@ export function useTextBuffer({
     dispatch({ type: 'move_to_offset', payload: { offset } });
   }, []);
 
-  const returnValue: TextBuffer = {
-    lines,
-    text,
-    cursor: [cursorRow, cursorCol],
-    preferredCol,
-    selectionAnchor,
+  const returnValue: TextBuffer = useMemo(
+    () => ({
+      lines,
+      text,
+      cursor: [cursorRow, cursorCol],
+      preferredCol,
+      selectionAnchor,
 
-    allVisualLines: visualLines,
-    viewportVisualLines: renderedVisualLines,
-    visualCursor,
-    visualScrollRow,
-    visualToLogicalMap,
+      allVisualLines: visualLines,
+      viewportVisualLines: renderedVisualLines,
+      visualCursor,
+      visualScrollRow,
+      visualToLogicalMap,
 
-    setText,
-    insert,
-    newline,
-    backspace,
-    del,
-    move,
-    undo,
-    redo,
-    replaceRange,
-    replaceRangeByOffset,
-    moveToOffset,
-    deleteWordLeft,
-    deleteWordRight,
+      setText,
+      insert,
+      newline,
+      backspace,
+      del,
+      move,
+      undo,
+      redo,
+      replaceRange,
+      replaceRangeByOffset,
+      moveToOffset,
+      deleteWordLeft,
+      deleteWordRight,
 
-    killLineRight,
-    killLineLeft,
-    handleInput,
-    openInExternalEditor,
-    // Vim-specific operations
-    vimDeleteWordForward,
-    vimDeleteWordBackward,
-    vimDeleteWordEnd,
-    vimChangeWordForward,
-    vimChangeWordBackward,
-    vimChangeWordEnd,
-    vimDeleteLine,
-    vimChangeLine,
-    vimDeleteToEndOfLine,
-    vimChangeToEndOfLine,
-    vimChangeMovement,
-    vimMoveLeft,
-    vimMoveRight,
-    vimMoveUp,
-    vimMoveDown,
-    vimMoveWordForward,
-    vimMoveWordBackward,
-    vimMoveWordEnd,
-    vimDeleteChar,
-    vimInsertAtCursor,
-    vimAppendAtCursor,
-    vimOpenLineBelow,
-    vimOpenLineAbove,
-    vimAppendAtLineEnd,
-    vimInsertAtLineStart,
-    vimMoveToLineStart,
-    vimMoveToLineEnd,
-    vimMoveToFirstNonWhitespace,
-    vimMoveToFirstLine,
-    vimMoveToLastLine,
-    vimMoveToLine,
-    vimEscapeInsertMode,
-  };
+      killLineRight,
+      killLineLeft,
+      handleInput,
+      openInExternalEditor,
+      // Vim-specific operations
+      vimDeleteWordForward,
+      vimDeleteWordBackward,
+      vimDeleteWordEnd,
+      vimChangeWordForward,
+      vimChangeWordBackward,
+      vimChangeWordEnd,
+      vimDeleteLine,
+      vimChangeLine,
+      vimDeleteToEndOfLine,
+      vimChangeToEndOfLine,
+      vimChangeMovement,
+      vimMoveLeft,
+      vimMoveRight,
+      vimMoveUp,
+      vimMoveDown,
+      vimMoveWordForward,
+      vimMoveWordBackward,
+      vimMoveWordEnd,
+      vimDeleteChar,
+      vimInsertAtCursor,
+      vimAppendAtCursor,
+      vimOpenLineBelow,
+      vimOpenLineAbove,
+      vimAppendAtLineEnd,
+      vimInsertAtLineStart,
+      vimMoveToLineStart,
+      vimMoveToLineEnd,
+      vimMoveToFirstNonWhitespace,
+      vimMoveToFirstLine,
+      vimMoveToLastLine,
+      vimMoveToLine,
+      vimEscapeInsertMode,
+    }),
+    [
+      lines,
+      text,
+      cursorRow,
+      cursorCol,
+      preferredCol,
+      selectionAnchor,
+      visualLines,
+      renderedVisualLines,
+      visualCursor,
+      visualScrollRow,
+      setText,
+      insert,
+      newline,
+      backspace,
+      del,
+      move,
+      undo,
+      redo,
+      replaceRange,
+      replaceRangeByOffset,
+      moveToOffset,
+      deleteWordLeft,
+      deleteWordRight,
+      killLineRight,
+      killLineLeft,
+      handleInput,
+      openInExternalEditor,
+      vimDeleteWordForward,
+      vimDeleteWordBackward,
+      vimDeleteWordEnd,
+      vimChangeWordForward,
+      vimChangeWordBackward,
+      vimChangeWordEnd,
+      vimDeleteLine,
+      vimChangeLine,
+      vimDeleteToEndOfLine,
+      vimChangeToEndOfLine,
+      vimChangeMovement,
+      vimMoveLeft,
+      vimMoveRight,
+      vimMoveUp,
+      vimMoveDown,
+      vimMoveWordForward,
+      vimMoveWordBackward,
+      vimMoveWordEnd,
+      vimDeleteChar,
+      vimInsertAtCursor,
+      vimAppendAtCursor,
+      vimOpenLineBelow,
+      vimOpenLineAbove,
+      vimAppendAtLineEnd,
+      vimInsertAtLineStart,
+      vimMoveToLineStart,
+      vimMoveToLineEnd,
+      vimMoveToFirstNonWhitespace,
+      vimMoveToFirstLine,
+      vimMoveToLastLine,
+      vimMoveToLine,
+      vimEscapeInsertMode,
+      visualToLogicalMap,
+    ],
+  );
   return returnValue;
 }
 

packages/cli/src/ui/contexts/FocusContext.tsx

@@ -0,0 +1,11 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { createContext, useContext } from 'react';
+
+export const FocusContext = createContext<boolean>(true);
+
+export const useFocusState = () => useContext(FocusContext);

packages/cli/src/ui/contexts/SessionContext.test.tsx

@@ -113,6 +113,88 @@ describe('SessionStatsContext', () => {
     expect(stats?.lastPromptTokenCount).toBe(100);
   });
 
+  it('should not update metrics if the data is the same', () => {
+    const contextRef: MutableRefObject<
+      ReturnType<typeof useSessionStats> | undefined
+    > = { current: undefined };
+
+    let renderCount = 0;
+    const CountingTestHarness = () => {
+      contextRef.current = useSessionStats();
+      renderCount++;
+      return null;
+    };
+
+    render(
+      <SessionStatsProvider>
+        <CountingTestHarness />
+      </SessionStatsProvider>,
+    );
+
+    expect(renderCount).toBe(1);
+
+    const metrics: SessionMetrics = {
+      models: {
+        'gemini-pro': {
+          api: { totalRequests: 1, totalErrors: 0, totalLatencyMs: 100 },
+          tokens: {
+            prompt: 10,
+            candidates: 20,
+            total: 30,
+            cached: 0,
+            thoughts: 0,
+            tool: 0,
+          },
+        },
+      },
+      tools: {
+        totalCalls: 0,
+        totalSuccess: 0,
+        totalFail: 0,
+        totalDurationMs: 0,
+        totalDecisions: { accept: 0, reject: 0, modify: 0 },
+        byName: {},
+      },
+    };
+
+    act(() => {
+      uiTelemetryService.emit('update', { metrics, lastPromptTokenCount: 10 });
+    });
+
+    expect(renderCount).toBe(2);
+
+    act(() => {
+      uiTelemetryService.emit('update', { metrics, lastPromptTokenCount: 10 });
+    });
+
+    expect(renderCount).toBe(2);
+
+    const newMetrics = {
+      ...metrics,
+      models: {
+        'gemini-pro': {
+          api: { totalRequests: 2, totalErrors: 0, totalLatencyMs: 200 },
+          tokens: {
+            prompt: 20,
+            candidates: 40,
+            total: 60,
+            cached: 0,
+            thoughts: 0,
+            tool: 0,
+          },
+        },
+      },
+    };
+    act(() => {
+      uiTelemetryService.emit('update', {
+        metrics: newMetrics,
+        lastPromptTokenCount: 20,
+      });
+    });
+
+    expect(renderCount).toBe(3);
+  });
+
   it('should throw an error when useSessionStats is used outside of a provider', () => {
     // Suppress console.error for this test since we expect an error
     const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});

packages/cli/src/ui/contexts/SessionContext.tsx

@@ -14,10 +14,129 @@ import {
   useEffect,
 } from 'react';
 
-import type { SessionMetrics, ModelMetrics } from '@google/gemini-cli-core';
+import type {
+  SessionMetrics,
+  ModelMetrics,
+  ToolCallStats,
+} from '@google/gemini-cli-core';
 import { uiTelemetryService, sessionId } from '@google/gemini-cli-core';
 
-// --- Interface Definitions ---
+export enum ToolCallDecision {
+  ACCEPT = 'accept',
+  REJECT = 'reject',
+  MODIFY = 'modify',
+  AUTO_ACCEPT = 'auto_accept',
+}
+
+function areModelMetricsEqual(a: ModelMetrics, b: ModelMetrics): boolean {
+  if (
+    a.api.totalRequests !== b.api.totalRequests ||
+    a.api.totalErrors !== b.api.totalErrors ||
+    a.api.totalLatencyMs !== b.api.totalLatencyMs
+  ) {
+    return false;
+  }
+  if (
+    a.tokens.prompt !== b.tokens.prompt ||
+    a.tokens.candidates !== b.tokens.candidates ||
+    a.tokens.total !== b.tokens.total ||
+    a.tokens.cached !== b.tokens.cached ||
+    a.tokens.thoughts !== b.tokens.thoughts ||
+    a.tokens.tool !== b.tokens.tool
+  ) {
+    return false;
+  }
+  return true;
+}
+
+function areToolCallStatsEqual(a: ToolCallStats, b: ToolCallStats): boolean {
+  if (
+    a.count !== b.count ||
+    a.success !== b.success ||
+    a.fail !== b.fail ||
+    a.durationMs !== b.durationMs
+  ) {
+    return false;
+  }
+  if (
+    a.decisions[ToolCallDecision.ACCEPT] !==
+      b.decisions[ToolCallDecision.ACCEPT] ||
+    a.decisions[ToolCallDecision.REJECT] !==
+      b.decisions[ToolCallDecision.REJECT] ||
+    a.decisions[ToolCallDecision.MODIFY] !==
+      b.decisions[ToolCallDecision.MODIFY] ||
+    a.decisions[ToolCallDecision.AUTO_ACCEPT] !==
+      b.decisions[ToolCallDecision.AUTO_ACCEPT]
+  ) {
+    return false;
+  }
+  return true;
+}
+
+function areMetricsEqual(a: SessionMetrics, b: SessionMetrics): boolean {
+  if (a === b) return true;
+  if (!a || !b) return false;
+
+  // Compare files
+  if (
+    a.files.totalLinesAdded !== b.files.totalLinesAdded ||
+    a.files.totalLinesRemoved !== b.files.totalLinesRemoved
+  ) {
+    return false;
+  }
+
+  // Compare tools
+  const toolsA = a.tools;
+  const toolsB = b.tools;
+  if (
+    toolsA.totalCalls !== toolsB.totalCalls ||
+    toolsA.totalSuccess !== toolsB.totalSuccess ||
+    toolsA.totalFail !== toolsB.totalFail ||
+    toolsA.totalDurationMs !== toolsB.totalDurationMs
+  ) {
+    return false;
+  }
+
+  // Compare tool decisions
+  if (
+    toolsA.totalDecisions[ToolCallDecision.ACCEPT] !==
+      toolsB.totalDecisions[ToolCallDecision.ACCEPT] ||
+    toolsA.totalDecisions[ToolCallDecision.REJECT] !==
+      toolsB.totalDecisions[ToolCallDecision.REJECT] ||
+    toolsA.totalDecisions[ToolCallDecision.MODIFY] !==
+      toolsB.totalDecisions[ToolCallDecision.MODIFY] ||
+    toolsA.totalDecisions[ToolCallDecision.AUTO_ACCEPT] !==
+      toolsB.totalDecisions[ToolCallDecision.AUTO_ACCEPT]
+  ) {
+    return false;
+  }
+
+  // Compare tools.byName
+  const toolsByNameAKeys = Object.keys(toolsA.byName);
+  const toolsByNameBKeys = Object.keys(toolsB.byName);
+  if (toolsByNameAKeys.length !== toolsByNameBKeys.length) return false;
+
+  for (const key of toolsByNameAKeys) {
+    const toolA = toolsA.byName[key];
+    const toolB = toolsB.byName[key];
+    if (!toolB || !areToolCallStatsEqual(toolA, toolB)) {
+      return false;
+    }
+  }
+
+  // Compare models
+  const modelsAKeys = Object.keys(a.models);
+  const modelsBKeys = Object.keys(b.models);
+  if (modelsAKeys.length !== modelsBKeys.length) return false;
+
+  for (const key of modelsAKeys) {
+    if (!b.models[key] || !areModelMetricsEqual(a.models[key], b.models[key])) {
+      return false;
+    }
+  }
+
+  return true;
+}
 
 export type { SessionMetrics, ModelMetrics };
 
@@ -80,11 +199,19 @@ export const SessionStatsProvider: React.FC<{ children: React.ReactNode }> = ({
       metrics: SessionMetrics;
       lastPromptTokenCount: number;
     }) => {
-      setStats((prevState) => ({
-        ...prevState,
-        metrics,
-        lastPromptTokenCount,
-      }));
+      setStats((prevState) => {
+        if (
+          prevState.lastPromptTokenCount === lastPromptTokenCount &&
+          areMetricsEqual(prevState.metrics, metrics)
+        ) {
+          return prevState;
+        }
+        return {
+          ...prevState,
+          metrics,
+          lastPromptTokenCount,
+        };
+      });
     };
 
     uiTelemetryService.on('update', handleUpdate);

packages/cli/src/ui/contexts/UIStateContext.tsx

@@ -78,7 +78,6 @@ export interface UIState {
   ctrlCPressedOnce: boolean;
   ctrlDPressedOnce: boolean;
   showEscapePrompt: boolean;
-  isFocused: boolean;
   elapsedTime: number;
   currentLoadingPhrase: string;
   historyRemountKey: number;

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -114,7 +114,7 @@ vi.mock('./useStateAndRef.js', () => ({
       }
       ref.current = val;
     });
-    return [ref, setVal];
+    return [val, ref, setVal];
   }),
 }));
 
@@ -2216,6 +2216,72 @@ describe('useGeminiStream', () => {
       });
     });
 
+    it('should memoize pendingHistoryItems', () => {
+      mockUseReactToolScheduler.mockReturnValue([
+        [],
+        mockScheduleToolCalls,
+        mockCancelAllToolCalls,
+        mockMarkToolsAsSubmitted,
+      ]);
+
+      const { result, rerender } = renderHook(() =>
+        useGeminiStream(
+          mockConfig.getGeminiClient(),
+          [],
+          mockAddItem,
+          mockConfig,
+          mockLoadedSettings,
+          mockOnDebugMessage,
+          mockHandleSlashCommand,
+          false,
+          () => 'vscode' as EditorType,
+          () => {},
+          () => Promise.resolve(),
+          false,
+          () => {},
+          () => {},
+          () => {},
+          () => {},
+          80,
+          24,
+        ),
+      );
+
+      const firstResult = result.current.pendingHistoryItems;
+      rerender();
+      const secondResult = result.current.pendingHistoryItems;
+
+      expect(firstResult).toStrictEqual(secondResult);
+
+      const newToolCalls: TrackedToolCall[] = [
+        {
+          request: { callId: 'call1', name: 'tool1', args: {} },
+          status: 'executing',
+          tool: {
+            name: 'tool1',
+            displayName: 'tool1',
+            description: 'desc1',
+            build: vi.fn(),
+          },
+          invocation: {
+            getDescription: () => 'Mock description',
+          },
+        } as unknown as TrackedExecutingToolCall,
+      ];
+
+      mockUseReactToolScheduler.mockReturnValue([
+        newToolCalls,
+        mockScheduleToolCalls,
+        mockCancelAllToolCalls,
+        mockMarkToolsAsSubmitted,
+      ]);
+
+      rerender();
+      const thirdResult = result.current.pendingHistoryItems;
+
+      expect(thirdResult).not.toStrictEqual(secondResult);
+    });
+
     it('should reset thought to null when user cancels', async () => {
       // Mock a stream that yields a thought then gets cancelled
       mockSendMessageStream.mockReturnValue(

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -111,7 +111,7 @@ export const useGeminiStream = (
   const turnCancelledRef = useRef(false);
   const [isResponding, setIsResponding] = useState<boolean>(false);
   const [thought, setThought] = useState<ThoughtSummary | null>(null);
-  const [pendingHistoryItemRef, setPendingHistoryItem] =
+  const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] =
     useStateAndRef<HistoryItemWithoutId | null>(null);
   const processedMemoryToolsRef = useRef<Set<string>>(new Set());
   const { startNewPrompt, getPromptCount } = useSessionStats();
@@ -1015,10 +1015,13 @@ export const useGeminiStream = (
     ],
   );
 
-  const pendingHistoryItems = [
-    pendingHistoryItemRef.current,
-    pendingToolCallGroupDisplay,
-  ].filter((i) => i !== undefined && i !== null);
+  const pendingHistoryItems = useMemo(
+    () =>
+      [pendingHistoryItem, pendingToolCallGroupDisplay].filter(
+        (i) => i !== undefined && i !== null,
+      ),
+    [pendingHistoryItem, pendingToolCallGroupDisplay],
+  );
 
   useEffect(() => {
     const saveRestorableToolCalls = async () => {

packages/cli/src/ui/hooks/useStateAndRef.ts

@@ -15,7 +15,7 @@ export const useStateAndRef = <
 >(
   initialValue: T,
 ) => {
-  const [_, setState] = React.useState<T>(initialValue);
+  const [state, setState] = React.useState<T>(initialValue);
   const ref = React.useRef<T>(initialValue);
 
   const setStateInternal = React.useCallback<typeof setState>(
@@ -32,5 +32,5 @@ export const useStateAndRef = <
     [],
   );
 
-  return [ref, setStateInternal] as const;
+  return [state, ref, setStateInternal] as const;
 };