Fix flicker in iterm2 (#266)

This commit is contained in:
Tae Hyung Kim 2025-05-07 12:57:19 -07:00 committed by GitHub
parent 358281f0fd
commit 0a7f461d39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 183 additions and 193 deletions

View File

@ -36,7 +36,7 @@ interface AppProps {
} }
export const App = ({ config, settings, cliVersion }: AppProps) => { export const App = ({ config, settings, cliVersion }: AppProps) => {
const { history, addItem, updateItem, clearItems } = useHistory(); const { history, addItem, clearItems } = useHistory();
const [startupWarnings, setStartupWarnings] = useState<string[]>([]); const [startupWarnings, setStartupWarnings] = useState<string[]>([]);
const [showHelp, setShowHelp] = useState<boolean>(false); const [showHelp, setShowHelp] = useState<boolean>(false);
const { const {
@ -57,9 +57,9 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
initError, initError,
debugMessage, debugMessage,
slashCommands, slashCommands,
pendingHistoryItem,
} = useGeminiStream( } = useGeminiStream(
addItem, addItem,
updateItem,
clearItems, clearItems,
refreshStatic, refreshStatic,
setShowHelp, setShowHelp,
@ -124,9 +124,6 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
// --- Render Logic --- // --- Render Logic ---
const { staticallyRenderedHistoryItems, updatableHistoryItems } =
getHistoryRenderSlices(history);
// Get terminal width // Get terminal width
const { stdout } = useStdout(); const { stdout } = useStdout();
const terminalWidth = stdout?.columns ?? 80; const terminalWidth = stdout?.columns ?? 80;
@ -146,10 +143,7 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
* content is set it'll flush content to the terminal and move the area which it's "clearing" * content is set it'll flush content to the terminal and move the area which it's "clearing"
* down a notch. Without Static the area which gets erased and redrawn continuously grows. * down a notch. Without Static the area which gets erased and redrawn continuously grows.
*/} */}
<Static <Static key={'static-key-' + staticKey} items={['header', ...history]}>
key={'static-key-' + staticKey}
items={['header', ...staticallyRenderedHistoryItems]}
>
{(item, index) => { {(item, index) => {
if (item === 'header') { if (item === 'header') {
return ( return (
@ -170,19 +164,14 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
); );
}} }}
</Static> </Static>
{pendingHistoryItem && (
{updatableHistoryItems.length > 0 && (
<Box flexDirection="column" alignItems="flex-start">
{updatableHistoryItems.map((historyItem) => (
<HistoryItemDisplay <HistoryItemDisplay
key={'history-' + historyItem.id} // TODO(taehykim): It seems like references to ids aren't necessary in
item={historyItem} // HistoryItemDisplay. Refactor later. Use a fake id for now.
item={{ ...pendingHistoryItem, id: 0 }}
onSubmit={submitQuery} onSubmit={submitQuery}
/> />
))}
</Box>
)} )}
{showHelp && <Help commands={slashCommands} />} {showHelp && <Help commands={slashCommands} />}
{startupWarnings.length > 0 && ( {startupWarnings.length > 0 && (
@ -295,21 +284,3 @@ export const App = ({ config, settings, cliVersion }: AppProps) => {
</Box> </Box>
); );
}; };
function getHistoryRenderSlices(history: HistoryItem[]) {
let staticallyRenderedHistoryItems: HistoryItem[] = [];
let updatableHistoryItems: HistoryItem[] = [];
if (
history.length > 1 &&
history[history.length - 2]?.type === 'tool_group'
) {
// If the second-to-last item is a tool_group, it and the last item are updateable
staticallyRenderedHistoryItems = history.slice(0, -2);
updatableHistoryItems = history.slice(-2);
} else {
// Otherwise, only the last item is updateable
staticallyRenderedHistoryItems = history.slice(0, -1);
updatableHistoryItems = history.slice(-1);
}
return { staticallyRenderedHistoryItems, updatableHistoryItems };
}

View File

@ -24,7 +24,6 @@ interface HandleAtCommandParams {
query: string; query: string;
config: Config; config: Config;
addItem: UseHistoryManagerReturn['addItem']; addItem: UseHistoryManagerReturn['addItem'];
updateItem: UseHistoryManagerReturn['updateItem'];
setDebugMessage: React.Dispatch<React.SetStateAction<string>>; setDebugMessage: React.Dispatch<React.SetStateAction<string>>;
messageId: number; messageId: number;
} }
@ -88,7 +87,7 @@ function parseAtCommand(
export async function handleAtCommand({ export async function handleAtCommand({
query, query,
config, config,
addItem: addItem, addItem,
setDebugMessage, setDebugMessage,
messageId: userMessageTimestamp, messageId: userMessageTimestamp,
}: HandleAtCommandParams): Promise<HandleAtCommandResult> { }: HandleAtCommandParams): Promise<HandleAtCommandResult> {

View File

@ -23,15 +23,16 @@ import {
import { type Chat, type PartListUnion, type Part } from '@google/genai'; import { type Chat, type PartListUnion, type Part } from '@google/genai';
import { import {
StreamingState, StreamingState,
HistoryItem,
IndividualToolCallDisplay, IndividualToolCallDisplay,
ToolCallStatus, ToolCallStatus,
HistoryItemWithoutId,
} from '../types.js'; } from '../types.js';
import { isAtCommand } from '../utils/commandUtils.js'; import { isAtCommand } from '../utils/commandUtils.js';
import { useSlashCommandProcessor } from './slashCommandProcessor.js'; import { useSlashCommandProcessor } from './slashCommandProcessor.js';
import { useShellCommandProcessor } from './shellCommandProcessor.js'; import { useShellCommandProcessor } from './shellCommandProcessor.js';
import { handleAtCommand } from './atCommandProcessor.js'; import { handleAtCommand } from './atCommandProcessor.js';
import { findSafeSplitPoint } from '../utils/markdownUtilities.js'; import { findSafeSplitPoint } from '../utils/markdownUtilities.js';
import { useStateAndRef } from './useStateAndRef.js';
import { UseHistoryManagerReturn } from './useHistoryManager.js'; import { UseHistoryManagerReturn } from './useHistoryManager.js';
/** /**
@ -40,7 +41,6 @@ import { UseHistoryManagerReturn } from './useHistoryManager.js';
*/ */
export const useGeminiStream = ( export const useGeminiStream = (
addItem: UseHistoryManagerReturn['addItem'], addItem: UseHistoryManagerReturn['addItem'],
updateItem: UseHistoryManagerReturn['updateItem'],
clearItems: UseHistoryManagerReturn['clearItems'], clearItems: UseHistoryManagerReturn['clearItems'],
refreshStatic: () => void, refreshStatic: () => void,
setShowHelp: React.Dispatch<React.SetStateAction<boolean>>, setShowHelp: React.Dispatch<React.SetStateAction<boolean>>,
@ -56,7 +56,8 @@ export const useGeminiStream = (
const abortControllerRef = useRef<AbortController | null>(null); const abortControllerRef = useRef<AbortController | null>(null);
const chatSessionRef = useRef<Chat | null>(null); const chatSessionRef = useRef<Chat | null>(null);
const geminiClientRef = useRef<GeminiClient | null>(null); const geminiClientRef = useRef<GeminiClient | null>(null);
const currentGeminiMessageIdRef = useRef<number | null>(null); const [pendingHistoryItemRef, setPendingHistoryItem] =
useStateAndRef<HistoryItemWithoutId | null>(null);
const { handleSlashCommand, slashCommands } = useSlashCommandProcessor( const { handleSlashCommand, slashCommands } = useSlashCommandProcessor(
addItem, addItem,
@ -93,13 +94,6 @@ export const useGeminiStream = (
} }
}); });
const updateGeminiMessage = useCallback(
(messageId: number, newContent: string) => {
updateItem(messageId, { text: newContent });
},
[updateItem],
);
const submitQuery = useCallback( const submitQuery = useCallback(
async (query: PartListUnion) => { async (query: PartListUnion) => {
if (streamingState === StreamingState.Responding) return; if (streamingState === StreamingState.Responding) return;
@ -124,7 +118,6 @@ export const useGeminiStream = (
query: trimmedQuery, query: trimmedQuery,
config, config,
addItem, addItem,
updateItem,
setDebugMessage, setDebugMessage,
messageId: userMessageTimestamp, messageId: userMessageTimestamp,
}); });
@ -170,7 +163,6 @@ export const useGeminiStream = (
setStreamingState(StreamingState.Responding); setStreamingState(StreamingState.Responding);
setInitError(null); setInitError(null);
const chat = chatSessionRef.current; const chat = chatSessionRef.current;
let currentToolGroupMessageId: number | null = null;
try { try {
abortControllerRef.current = new AbortController(); abortControllerRef.current = new AbortController();
@ -183,30 +175,33 @@ export const useGeminiStream = (
); );
let currentGeminiText = ''; let currentGeminiText = '';
let hasInitialGeminiResponse = false;
for await (const event of stream) { for await (const event of stream) {
if (signal.aborted) break; if (signal.aborted) break;
if (event.type === ServerGeminiEventType.Content) { if (event.type === ServerGeminiEventType.Content) {
currentGeminiText += event.value; currentGeminiText += event.value;
currentToolGroupMessageId = null; // Reset group on new text content
if (!hasInitialGeminiResponse) { if (pendingHistoryItemRef.current?.type !== 'gemini') {
hasInitialGeminiResponse = true; // Flush out existing pending history item.
const eventId = addItem( if (pendingHistoryItemRef.current) {
{ type: 'gemini', text: currentGeminiText }, addItem(pendingHistoryItemRef.current, userMessageTimestamp);
userMessageTimestamp, }
); setPendingHistoryItem({
currentGeminiMessageIdRef.current = eventId; type: 'gemini',
} else if (currentGeminiMessageIdRef.current !== null) { text: currentGeminiText,
});
}
// Split large messages for better rendering performance // Split large messages for better rendering performance
const splitPoint = findSafeSplitPoint(currentGeminiText); const splitPoint = findSafeSplitPoint(currentGeminiText);
if (splitPoint === currentGeminiText.length) { if (splitPoint === currentGeminiText.length) {
updateGeminiMessage( // Update the existing message with accumulated content
currentGeminiMessageIdRef.current, setPendingHistoryItem((pending) => ({
currentGeminiText, // There might be a more typesafe way to do this.
); ...pending!,
text: currentGeminiText,
}));
} else { } else {
// This indicates that we need to split up this Gemini Message. // This indicates that we need to split up this Gemini Message.
// Splitting a message is primarily a performance consideration. There is a // Splitting a message is primarily a performance consideration. There is a
@ -216,22 +211,20 @@ export const useGeminiStream = (
// multiple times per-second (as streaming occurs). Prior to this change you'd // multiple times per-second (as streaming occurs). Prior to this change you'd
// see heavy flickering of the terminal. This ensures that larger messages get // see heavy flickering of the terminal. This ensures that larger messages get
// broken up so that there are more "statically" rendered. // broken up so that there are more "statically" rendered.
const originalMessageRef = currentGeminiMessageIdRef.current;
const beforeText = currentGeminiText.substring(0, splitPoint); const beforeText = currentGeminiText.substring(0, splitPoint);
const afterText = currentGeminiText.substring(splitPoint); const afterText = currentGeminiText.substring(splitPoint);
currentGeminiText = afterText; // Continue accumulating from split point currentGeminiText = afterText; // Continue accumulating from split point
updateItem(originalMessageRef, { text: beforeText }); addItem(
const nextId = addItem( { type: 'gemini_content', text: beforeText },
{ type: 'gemini_content', text: afterText },
userMessageTimestamp, userMessageTimestamp,
); );
currentGeminiMessageIdRef.current = nextId; setPendingHistoryItem({
} type: 'gemini_content',
text: afterText,
});
} }
} else if (event.type === ServerGeminiEventType.ToolCallRequest) { } else if (event.type === ServerGeminiEventType.ToolCallRequest) {
currentGeminiText = ''; currentGeminiText = '';
hasInitialGeminiResponse = false;
currentGeminiMessageIdRef.current = null;
const { callId, name, args } = event.value; const { callId, name, args } = event.value;
const cliTool = toolRegistry.getTool(name); const cliTool = toolRegistry.getTool(name);
@ -240,12 +233,15 @@ export const useGeminiStream = (
continue; continue;
} }
// Create a new tool group if needed if (pendingHistoryItemRef.current?.type !== 'tool_group') {
if (currentToolGroupMessageId === null) { // Flush out existing pending history item.
currentToolGroupMessageId = addItem( if (pendingHistoryItemRef.current) {
{ type: 'tool_group', tools: [] } as Omit<HistoryItem, 'id'>, addItem(pendingHistoryItemRef.current, userMessageTimestamp);
userMessageTimestamp, }
); setPendingHistoryItem({
type: 'tool_group',
tools: [],
});
} }
let description: string; let description: string;
@ -264,27 +260,16 @@ export const useGeminiStream = (
confirmationDetails: undefined, confirmationDetails: undefined,
}; };
// Add the pending tool call to the current group // Add pending tool call to the UI history group
if (currentToolGroupMessageId !== null) { setPendingHistoryItem((pending) =>
updateItem( // Should always be true.
currentToolGroupMessageId, pending?.type === 'tool_group'
( ? {
currentItem: HistoryItem, ...pending,
): Partial<Omit<HistoryItem, 'id'>> => { tools: [...pending.tools, toolCallDisplay],
if (currentItem?.type !== 'tool_group') {
console.error(
`Attempted to update non-tool-group item ${currentItem?.id} as tool group.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
} }
const currentTools = currentItem.tools; : null,
return {
...currentItem,
tools: [...currentTools, toolCallDisplay],
} as Partial<Omit<HistoryItem, 'id'>>;
},
); );
}
} else if (event.type === ServerGeminiEventType.ToolCallResponse) { } else if (event.type === ServerGeminiEventType.ToolCallResponse) {
const status = event.value.error const status = event.value.error
? ToolCallStatus.Error ? ToolCallStatus.Error
@ -303,6 +288,12 @@ export const useGeminiStream = (
} }
} // End stream loop } // End stream loop
// We're waiting for user input now so all pending history can be committed.
if (pendingHistoryItemRef.current) {
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
setPendingHistoryItem(null);
}
setStreamingState(StreamingState.Idle); setStreamingState(StreamingState.Idle);
} catch (error: unknown) { } catch (error: unknown) {
if (!isNodeError(error) || error.name !== 'AbortError') { if (!isNodeError(error) || error.name !== 'AbortError') {
@ -326,19 +317,12 @@ export const useGeminiStream = (
callId: string, callId: string,
confirmationDetails: ToolCallConfirmationDetails | undefined, confirmationDetails: ToolCallConfirmationDetails | undefined,
) { ) {
if (currentToolGroupMessageId === null) return; if (pendingHistoryItemRef.current?.type !== 'tool_group') return;
updateItem( setPendingHistoryItem((item) =>
currentToolGroupMessageId, item?.type === 'tool_group'
(currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => { ? {
if (currentItem?.type !== 'tool_group') { ...item,
console.error( tools: item.tools.map((tool) =>
`Attempted to update non-tool-group item ${currentItem?.id} status.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
}
return {
...currentItem,
tools: (currentItem.tools || []).map((tool) =>
tool.callId === callId tool.callId === callId
? { ? {
...tool, ...tool,
@ -347,8 +331,8 @@ export const useGeminiStream = (
} }
: tool, : tool,
), ),
} as Partial<Omit<HistoryItem, 'id'>>; }
}, : null,
); );
} }
@ -356,19 +340,11 @@ export const useGeminiStream = (
toolResponse: ToolCallResponseInfo, toolResponse: ToolCallResponseInfo,
status: ToolCallStatus, status: ToolCallStatus,
) { ) {
if (currentToolGroupMessageId === null) return; setPendingHistoryItem((item) =>
updateItem( item?.type === 'tool_group'
currentToolGroupMessageId, ? {
(currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => { ...item,
if (currentItem?.type !== 'tool_group') { tools: item.tools.map((tool) => {
console.error(
`Attempted to update non-tool-group item ${currentItem?.id} response.`,
);
return currentItem as Partial<Omit<HistoryItem, 'id'>>;
}
return {
...currentItem,
tools: (currentItem.tools || []).map((tool) => {
if (tool.callId === toolResponse.callId) { if (tool.callId === toolResponse.callId) {
return { return {
...tool, ...tool,
@ -379,8 +355,8 @@ export const useGeminiStream = (
return tool; return tool;
} }
}), }),
} as Partial<Omit<HistoryItem, 'id'>>; }
}, : null,
); );
} }
@ -397,15 +373,12 @@ export const useGeminiStream = (
originalConfirmationDetails.onConfirm(outcome); originalConfirmationDetails.onConfirm(outcome);
// Ensure UI updates before potentially long-running operations // Ensure UI updates before potentially long-running operations
if (currentToolGroupMessageId !== null) { if (pendingHistoryItemRef?.current?.type === 'tool_group') {
updateItem( setPendingHistoryItem((item) =>
currentToolGroupMessageId, item?.type === 'tool_group'
(currentItem: HistoryItem) => { ? {
if (currentItem?.type !== 'tool_group') ...item,
return currentItem as Partial<Omit<HistoryItem, 'id'>>; tools: item.tools.map((tool) =>
return {
...currentItem,
tools: (currentItem.tools || []).map((tool) =>
tool.callId === request.callId tool.callId === request.callId
? { ? {
...tool, ...tool,
@ -414,8 +387,8 @@ export const useGeminiStream = (
} }
: tool, : tool,
), ),
} as Partial<Omit<HistoryItem, 'id'>>; }
}, : item,
); );
refreshStatic(); refreshStatic();
} }
@ -485,17 +458,14 @@ export const useGeminiStream = (
}, },
[ [
streamingState, streamingState,
config, setShowHelp,
updateGeminiMessage,
handleSlashCommand, handleSlashCommand,
handleShellCommand, handleShellCommand,
setDebugMessage, config,
setStreamingState,
addItem, addItem,
updateItem, pendingHistoryItemRef,
setShowHelp, setPendingHistoryItem,
toolRegistry, toolRegistry,
setInitError,
refreshStatic, refreshStatic,
], ],
); );
@ -506,5 +476,9 @@ export const useGeminiStream = (
initError, initError,
debugMessage, debugMessage,
slashCommands, slashCommands,
// Normally we would be concerned that the ref would not be up-to-date, but
// this isn't a concern as the ref is updated whenever the corresponding
// state is updated.
pendingHistoryItem: pendingHistoryItemRef.current,
}; };
}; };

View File

@ -49,7 +49,13 @@ export function useHistory(): UseHistoryManagerReturn {
[getNextMessageId], [getNextMessageId],
); );
// Updates an existing history item identified by its ID. /**
* Updates an existing history item identified by its ID.
* @deprecated Prefer not to update history item directly as we are currently
* rendering all history items in <Static /> for performance reasons. Only use
* if ABSOLUTELY NECESSARY
*/
//
const updateItem = useCallback( const updateItem = useCallback(
( (
id: number, id: number,

View File

@ -0,0 +1,36 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import React from 'react';
// Hook to return state, state setter, and ref to most up-to-date value of state.
// We need this in order to setState and reference the updated state multiple
// times in the same function.
export const useStateAndRef = <
// Everything but function.
T extends object | null | undefined | number | string,
>(
initialValue: T,
) => {
const [_, setState] = React.useState<T>(initialValue);
const ref = React.useRef<T>(initialValue);
const setStateInternal = React.useCallback<typeof setState>(
(newStateOrCallback) => {
let newValue: T;
if (typeof newStateOrCallback === 'function') {
newValue = newStateOrCallback(ref.current);
} else {
newValue = newStateOrCallback;
}
setState(newValue);
ref.current = newValue;
},
[],
);
return [ref, setStateInternal] as const;
};

View File

@ -52,11 +52,13 @@ export interface IndividualToolCallDisplay {
} }
export interface HistoryItemBase { export interface HistoryItemBase {
id: number;
text?: string; // Text content for user/gemini/info/error messages text?: string; // Text content for user/gemini/info/error messages
} }
export type HistoryItem = HistoryItemBase & // Using Omit<HistoryItem, 'id'> seems to have some issues with typescript's
// type inference e.g. historyItem.type === 'tool_group' isn't auto-inferring that
// 'tools' in historyItem.
export type HistoryItemWithoutId = HistoryItemBase &
( (
| { type: 'user'; text: string } | { type: 'user'; text: string }
| { type: 'gemini'; text: string } | { type: 'gemini'; text: string }
@ -65,3 +67,5 @@ export type HistoryItem = HistoryItemBase &
| { type: 'error'; text: string } | { type: 'error'; text: string }
| { type: 'tool_group'; tools: IndividualToolCallDisplay[] } | { type: 'tool_group'; tools: IndividualToolCallDisplay[] }
); );
export type HistoryItem = HistoryItemWithoutId & { id: number };