diff options
| -rw-r--r-- | packages/cli/src/ui/App.tsx | 49 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/atCommandProcessor.ts | 3 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 272 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useHistoryManager.ts | 8 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useStateAndRef.ts | 36 | ||||
| -rw-r--r-- | packages/cli/src/ui/types.ts | 8 |
6 files changed, 183 insertions, 193 deletions
diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 1c1ec424..9fe651b5 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -36,7 +36,7 @@ interface 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 [showHelp, setShowHelp] = useState<boolean>(false); const { @@ -57,9 +57,9 @@ export const App = ({ config, settings, cliVersion }: AppProps) => { initError, debugMessage, slashCommands, + pendingHistoryItem, } = useGeminiStream( addItem, - updateItem, clearItems, refreshStatic, setShowHelp, @@ -124,9 +124,6 @@ export const App = ({ config, settings, cliVersion }: AppProps) => { // --- Render Logic --- - const { staticallyRenderedHistoryItems, updatableHistoryItems } = - getHistoryRenderSlices(history); - // Get terminal width const { stdout } = useStdout(); 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" * down a notch. Without Static the area which gets erased and redrawn continuously grows. */} - <Static - key={'static-key-' + staticKey} - items={['header', ...staticallyRenderedHistoryItems]} - > + <Static key={'static-key-' + staticKey} items={['header', ...history]}> {(item, index) => { if (item === 'header') { return ( @@ -170,19 +164,14 @@ export const App = ({ config, settings, cliVersion }: AppProps) => { ); }} </Static> - - {updatableHistoryItems.length > 0 && ( - <Box flexDirection="column" alignItems="flex-start"> - {updatableHistoryItems.map((historyItem) => ( - <HistoryItemDisplay - key={'history-' + historyItem.id} - item={historyItem} - onSubmit={submitQuery} - /> - ))} - </Box> + {pendingHistoryItem && ( + <HistoryItemDisplay + // TODO(taehykim): It seems like references to ids aren't necessary in + // HistoryItemDisplay. Refactor later. Use a fake id for now. + item={{ ...pendingHistoryItem, id: 0 }} + onSubmit={submitQuery} + /> )} - {showHelp && <Help commands={slashCommands} />} {startupWarnings.length > 0 && ( @@ -295,21 +284,3 @@ export const App = ({ config, settings, cliVersion }: AppProps) => { </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 }; -} diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index 50e81589..5ffa5383 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -24,7 +24,6 @@ interface HandleAtCommandParams { query: string; config: Config; addItem: UseHistoryManagerReturn['addItem']; - updateItem: UseHistoryManagerReturn['updateItem']; setDebugMessage: React.Dispatch<React.SetStateAction<string>>; messageId: number; } @@ -88,7 +87,7 @@ function parseAtCommand( export async function handleAtCommand({ query, config, - addItem: addItem, + addItem, setDebugMessage, messageId: userMessageTimestamp, }: HandleAtCommandParams): Promise<HandleAtCommandResult> { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 40c688a3..1c860501 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -23,15 +23,16 @@ import { import { type Chat, type PartListUnion, type Part } from '@google/genai'; import { StreamingState, - HistoryItem, IndividualToolCallDisplay, ToolCallStatus, + HistoryItemWithoutId, } from '../types.js'; import { isAtCommand } from '../utils/commandUtils.js'; import { useSlashCommandProcessor } from './slashCommandProcessor.js'; import { useShellCommandProcessor } from './shellCommandProcessor.js'; import { handleAtCommand } from './atCommandProcessor.js'; import { findSafeSplitPoint } from '../utils/markdownUtilities.js'; +import { useStateAndRef } from './useStateAndRef.js'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; /** @@ -40,7 +41,6 @@ import { UseHistoryManagerReturn } from './useHistoryManager.js'; */ export const useGeminiStream = ( addItem: UseHistoryManagerReturn['addItem'], - updateItem: UseHistoryManagerReturn['updateItem'], clearItems: UseHistoryManagerReturn['clearItems'], refreshStatic: () => void, setShowHelp: React.Dispatch<React.SetStateAction<boolean>>, @@ -56,7 +56,8 @@ export const useGeminiStream = ( const abortControllerRef = useRef<AbortController | null>(null); const chatSessionRef = useRef<Chat | 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( addItem, @@ -93,13 +94,6 @@ export const useGeminiStream = ( } }); - const updateGeminiMessage = useCallback( - (messageId: number, newContent: string) => { - updateItem(messageId, { text: newContent }); - }, - [updateItem], - ); - const submitQuery = useCallback( async (query: PartListUnion) => { if (streamingState === StreamingState.Responding) return; @@ -124,7 +118,6 @@ export const useGeminiStream = ( query: trimmedQuery, config, addItem, - updateItem, setDebugMessage, messageId: userMessageTimestamp, }); @@ -170,7 +163,6 @@ export const useGeminiStream = ( setStreamingState(StreamingState.Responding); setInitError(null); const chat = chatSessionRef.current; - let currentToolGroupMessageId: number | null = null; try { abortControllerRef.current = new AbortController(); @@ -183,55 +175,56 @@ export const useGeminiStream = ( ); let currentGeminiText = ''; - let hasInitialGeminiResponse = false; for await (const event of stream) { if (signal.aborted) break; if (event.type === ServerGeminiEventType.Content) { currentGeminiText += event.value; - currentToolGroupMessageId = null; // Reset group on new text content - if (!hasInitialGeminiResponse) { - hasInitialGeminiResponse = true; - const eventId = addItem( - { type: 'gemini', text: currentGeminiText }, + if (pendingHistoryItemRef.current?.type !== 'gemini') { + // Flush out existing pending history item. + if (pendingHistoryItemRef.current) { + addItem(pendingHistoryItemRef.current, userMessageTimestamp); + } + setPendingHistoryItem({ + type: 'gemini', + text: currentGeminiText, + }); + } + + // Split large messages for better rendering performance + const splitPoint = findSafeSplitPoint(currentGeminiText); + if (splitPoint === currentGeminiText.length) { + // Update the existing message with accumulated content + setPendingHistoryItem((pending) => ({ + // There might be a more typesafe way to do this. + ...pending!, + text: currentGeminiText, + })); + } else { + // This indicates that we need to split up this Gemini Message. + // Splitting a message is primarily a performance consideration. There is a + // <Static> component at the root of App.tsx which takes care of rendering + // content statically or dynamically. Everything but the last message is + // treated as static in order to prevent re-rendering an entire message history + // 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 + // broken up so that there are more "statically" rendered. + const beforeText = currentGeminiText.substring(0, splitPoint); + const afterText = currentGeminiText.substring(splitPoint); + currentGeminiText = afterText; // Continue accumulating from split point + addItem( + { type: 'gemini_content', text: beforeText }, userMessageTimestamp, ); - currentGeminiMessageIdRef.current = eventId; - } else if (currentGeminiMessageIdRef.current !== null) { - // Split large messages for better rendering performance - const splitPoint = findSafeSplitPoint(currentGeminiText); - if (splitPoint === currentGeminiText.length) { - updateGeminiMessage( - currentGeminiMessageIdRef.current, - currentGeminiText, - ); - } else { - // This indicates that we need to split up this Gemini Message. - // Splitting a message is primarily a performance consideration. There is a - // <Static> component at the root of App.tsx which takes care of rendering - // content statically or dynamically. Everything but the last message is - // treated as static in order to prevent re-rendering an entire message history - // 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 - // broken up so that there are more "statically" rendered. - const originalMessageRef = currentGeminiMessageIdRef.current; - const beforeText = currentGeminiText.substring(0, splitPoint); - const afterText = currentGeminiText.substring(splitPoint); - currentGeminiText = afterText; // Continue accumulating from split point - updateItem(originalMessageRef, { text: beforeText }); - const nextId = addItem( - { type: 'gemini_content', text: afterText }, - userMessageTimestamp, - ); - currentGeminiMessageIdRef.current = nextId; - } + setPendingHistoryItem({ + type: 'gemini_content', + text: afterText, + }); } } else if (event.type === ServerGeminiEventType.ToolCallRequest) { currentGeminiText = ''; - hasInitialGeminiResponse = false; - currentGeminiMessageIdRef.current = null; const { callId, name, args } = event.value; const cliTool = toolRegistry.getTool(name); @@ -240,12 +233,15 @@ export const useGeminiStream = ( continue; } - // Create a new tool group if needed - if (currentToolGroupMessageId === null) { - currentToolGroupMessageId = addItem( - { type: 'tool_group', tools: [] } as Omit<HistoryItem, 'id'>, - userMessageTimestamp, - ); + if (pendingHistoryItemRef.current?.type !== 'tool_group') { + // Flush out existing pending history item. + if (pendingHistoryItemRef.current) { + addItem(pendingHistoryItemRef.current, userMessageTimestamp); + } + setPendingHistoryItem({ + type: 'tool_group', + tools: [], + }); } let description: string; @@ -264,27 +260,16 @@ export const useGeminiStream = ( confirmationDetails: undefined, }; - // Add the pending tool call to the current group - if (currentToolGroupMessageId !== null) { - updateItem( - currentToolGroupMessageId, - ( - currentItem: HistoryItem, - ): Partial<Omit<HistoryItem, 'id'>> => { - 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'>>; + // Add pending tool call to the UI history group + setPendingHistoryItem((pending) => + // Should always be true. + pending?.type === 'tool_group' + ? { + ...pending, + tools: [...pending.tools, toolCallDisplay], } - const currentTools = currentItem.tools; - return { - ...currentItem, - tools: [...currentTools, toolCallDisplay], - } as Partial<Omit<HistoryItem, 'id'>>; - }, - ); - } + : null, + ); } else if (event.type === ServerGeminiEventType.ToolCallResponse) { const status = event.value.error ? ToolCallStatus.Error @@ -303,6 +288,12 @@ export const useGeminiStream = ( } } // 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); } catch (error: unknown) { if (!isNodeError(error) || error.name !== 'AbortError') { @@ -326,29 +317,22 @@ export const useGeminiStream = ( callId: string, confirmationDetails: ToolCallConfirmationDetails | undefined, ) { - if (currentToolGroupMessageId === null) return; - updateItem( - currentToolGroupMessageId, - (currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => { - if (currentItem?.type !== 'tool_group') { - console.error( - `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, - status: ToolCallStatus.Confirming, - confirmationDetails, - } - : tool, - ), - } as Partial<Omit<HistoryItem, 'id'>>; - }, + if (pendingHistoryItemRef.current?.type !== 'tool_group') return; + setPendingHistoryItem((item) => + item?.type === 'tool_group' + ? { + ...item, + tools: item.tools.map((tool) => + tool.callId === callId + ? { + ...tool, + status: ToolCallStatus.Confirming, + confirmationDetails, + } + : tool, + ), + } + : null, ); } @@ -356,31 +340,23 @@ export const useGeminiStream = ( toolResponse: ToolCallResponseInfo, status: ToolCallStatus, ) { - if (currentToolGroupMessageId === null) return; - updateItem( - currentToolGroupMessageId, - (currentItem: HistoryItem): Partial<Omit<HistoryItem, 'id'>> => { - if (currentItem?.type !== 'tool_group') { - 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) { - return { - ...tool, - status, - resultDisplay: toolResponse.resultDisplay, - }; - } else { - return tool; - } - }), - } as Partial<Omit<HistoryItem, 'id'>>; - }, + setPendingHistoryItem((item) => + item?.type === 'tool_group' + ? { + ...item, + tools: item.tools.map((tool) => { + if (tool.callId === toolResponse.callId) { + return { + ...tool, + status, + resultDisplay: toolResponse.resultDisplay, + }; + } else { + return tool; + } + }), + } + : null, ); } @@ -397,25 +373,22 @@ export const useGeminiStream = ( originalConfirmationDetails.onConfirm(outcome); // Ensure UI updates before potentially long-running operations - if (currentToolGroupMessageId !== null) { - updateItem( - currentToolGroupMessageId, - (currentItem: HistoryItem) => { - if (currentItem?.type !== 'tool_group') - return currentItem as Partial<Omit<HistoryItem, 'id'>>; - return { - ...currentItem, - tools: (currentItem.tools || []).map((tool) => - tool.callId === request.callId - ? { - ...tool, - confirmationDetails: undefined, - status: ToolCallStatus.Executing, - } - : tool, - ), - } as Partial<Omit<HistoryItem, 'id'>>; - }, + if (pendingHistoryItemRef?.current?.type === 'tool_group') { + setPendingHistoryItem((item) => + item?.type === 'tool_group' + ? { + ...item, + tools: item.tools.map((tool) => + tool.callId === request.callId + ? { + ...tool, + confirmationDetails: undefined, + status: ToolCallStatus.Executing, + } + : tool, + ), + } + : item, ); refreshStatic(); } @@ -485,17 +458,14 @@ export const useGeminiStream = ( }, [ streamingState, - config, - updateGeminiMessage, + setShowHelp, handleSlashCommand, handleShellCommand, - setDebugMessage, - setStreamingState, + config, addItem, - updateItem, - setShowHelp, + pendingHistoryItemRef, + setPendingHistoryItem, toolRegistry, - setInitError, refreshStatic, ], ); @@ -506,5 +476,9 @@ export const useGeminiStream = ( initError, debugMessage, 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, }; }; diff --git a/packages/cli/src/ui/hooks/useHistoryManager.ts b/packages/cli/src/ui/hooks/useHistoryManager.ts index 52dcfd4e..424f1fb6 100644 --- a/packages/cli/src/ui/hooks/useHistoryManager.ts +++ b/packages/cli/src/ui/hooks/useHistoryManager.ts @@ -49,7 +49,13 @@ export function useHistory(): UseHistoryManagerReturn { [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( ( id: number, diff --git a/packages/cli/src/ui/hooks/useStateAndRef.ts b/packages/cli/src/ui/hooks/useStateAndRef.ts new file mode 100644 index 00000000..d073a1dc --- /dev/null +++ b/packages/cli/src/ui/hooks/useStateAndRef.ts @@ -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; +}; diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index d6279598..62869dbe 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -52,11 +52,13 @@ export interface IndividualToolCallDisplay { } export interface HistoryItemBase { - id: number; 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: 'gemini'; text: string } @@ -65,3 +67,5 @@ export type HistoryItem = HistoryItemBase & | { type: 'error'; text: string } | { type: 'tool_group'; tools: IndividualToolCallDisplay[] } ); + +export type HistoryItem = HistoryItemWithoutId & { id: number }; |