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

This commit is contained in:
Gal Zahavi
2025-09-17 15:37:13 -07:00
committed by GitHub
parent d54cdd8802
commit 6756a8b8a9
13 changed files with 499 additions and 85 deletions
+1 -1
View File
@@ -17,7 +17,7 @@
"sandboxImageUri": "us-docker.pkg.dev/gemini-code-dev/gemini-cli/sandbox:0.7.0-nightly.20250917.0b10ba2c"
},
"scripts": {
"start": "node scripts/start.js",
"start": "cross-env node scripts/start.js",
"start:a2a-server": "CODER_AGENT_PORT=41242 npm run start --workspace @google/gemini-cli-a2a-server",
"debug": "cross-env DEBUG=1 node --inspect-brk scripts/start.js",
"auth:npm": "npx google-artifactregistry-auth",
+12 -2
View File
@@ -4,6 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import React from 'react';
import { render } from 'ink';
import { AppContainer } from './ui/AppContainer.js';
import { loadCliConfig, parseArguments } from './config/config.js';
@@ -212,10 +213,19 @@ export async function startInteractiveUI(
);
};
const instance = render(<AppWrapper />, {
const instance = render(
process.env['DEBUG'] ? (
<React.StrictMode>
<AppWrapper />
</React.StrictMode>
) : (
<AppWrapper />
),
{
exitOnCtrlC: false,
isScreenReaderEnabled: config.getScreenReader(),
});
},
);
checkForUpdates()
.then((info) => {
+3
View File
@@ -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 || [],
}}
>
<FocusContext.Provider value={isFocused}>
<App />
</FocusContext.Provider>
</AppContext.Provider>
</ConfigContext.Provider>
</UIActionsContext.Provider>
+3 -1
View File
@@ -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={
@@ -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', () => {
@@ -1984,7 +1984,8 @@ export function useTextBuffer({
dispatch({ type: 'move_to_offset', payload: { offset } });
}, []);
const returnValue: TextBuffer = {
const returnValue: TextBuffer = useMemo(
() => ({
lines,
text,
cursor: [cursorRow, cursorCol],
@@ -2048,7 +2049,70 @@ export function useTextBuffer({
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;
}
@@ -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);
@@ -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(() => {});
+131 -4
View File
@@ -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) => ({
setStats((prevState) => {
if (
prevState.lastPromptTokenCount === lastPromptTokenCount &&
areMetricsEqual(prevState.metrics, metrics)
) {
return prevState;
}
return {
...prevState,
metrics,
lastPromptTokenCount,
}));
};
});
};
uiTelemetryService.on('update', handleUpdate);
@@ -78,7 +78,6 @@ export interface UIState {
ctrlCPressedOnce: boolean;
ctrlDPressedOnce: boolean;
showEscapePrompt: boolean;
isFocused: boolean;
elapsedTime: number;
currentLoadingPhrase: string;
historyRemountKey: number;
@@ -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(
+8 -5
View File
@@ -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 () => {
+2 -2
View File
@@ -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;
};