diff options
| author | Taylor Mullen <[email protected]> | 2025-04-25 17:11:08 -0700 |
|---|---|---|
| committer | N. Taylor Mullen <[email protected]> | 2025-04-26 16:08:05 -0700 |
| commit | 5be89befeff9c4d4f3ab9f508f030bc153fdd06b (patch) | |
| tree | 9d4f679a0e7292132cab04fdd3c24062fcd66ce8 /packages/cli/src/ui/hooks/useGeminiStream.ts | |
| parent | aa65a4a1fc3f51589c7633217f9d3c8bd0141abb (diff) | |
feat: Fix flickering in iTerm + scrolling + performance issues.
- Refactors history display using Ink's <Static> component to prevent flickering and improve performance by rendering completed items statically.
- Introduces ConsolePatcher component to capture and display console.log, console.warn, and console.error output within the Ink UI, addressing native handling issues.
- Introduce a new content splitting mechanism to work better for static items. Basically when content gets too long we will now split content into multiple blocks for Gemini messages to ensure that we can statically cache larger pieces of history.
Fixes:
- https://b.corp.google.com/issues/411450097
- https://b.corp.google.com/issues/412716309
Diffstat (limited to 'packages/cli/src/ui/hooks/useGeminiStream.ts')
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 67 |
1 files changed, 56 insertions, 11 deletions
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d8254cd3..5a9de54d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -28,6 +28,7 @@ import { IndividualToolCallDisplay, ToolCallStatus, } from '../types.js'; +import { findSafeSplitPoint } from '../utils/markdownUtilities.js'; const addHistoryItem = ( setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>, @@ -169,6 +170,28 @@ export const useGeminiStream = ( return false; // Not handled by a manual command. }; + // Helper function to update Gemini message content + const updateAndAddGeminiMessageContent = useCallback( + ( + messageId: number, + previousContent: string, + nextId: number, + nextContent: string, + ) => { + setHistory((prevHistory) => { + const beforeNextHistory = prevHistory.map((item) => + item.id === messageId ? { ...item, text: previousContent } : item, + ); + + return [ + ...beforeNextHistory, + { id: nextId, type: 'gemini_content', text: nextContent }, + ]; + }); + }, + [setHistory], + ); + // Improved submit query function const submitQuery = useCallback( async (query: PartListUnion) => { @@ -250,11 +273,37 @@ export const useGeminiStream = ( eventTimestamp, ); } else if (currentGeminiMessageIdRef.current !== null) { - // Update the existing message with accumulated content - updateGeminiMessage( - currentGeminiMessageIdRef.current, - currentGeminiText, - ); + const splitPoint = findSafeSplitPoint(currentGeminiText); + + if (splitPoint === currentGeminiText.length) { + // Update the existing message with accumulated content + 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); + + currentGeminiMessageIdRef.current = + getNextMessageId(userMessageTimestamp); + const afterText = currentGeminiText.substring(splitPoint); + currentGeminiText = afterText; + updateAndAddGeminiMessageContent( + originalMessageRef, + beforeText, + currentGeminiMessageIdRef.current, + afterText, + ); + } } } else if (event.type === ServerGeminiEventType.ToolCallRequest) { // Reset the Gemini message tracking for the next response @@ -414,9 +463,6 @@ export const useGeminiStream = ( ) => { originalConfirmationDetails.onConfirm(outcome); - // Reset streaming state since confirmation has been chosen. - setStreamingState(StreamingState.Idle); - if (outcome === ToolConfirmationOutcome.Cancel) { let resultDisplay: ToolResultDisplay | undefined; if ('fileDiff' in originalConfirmationDetails) { @@ -444,8 +490,7 @@ export const useGeminiStream = ( }; updateFunctionResponseUI(responseInfo, ToolCallStatus.Error); - - await submitQuery(functionResponse); + setStreamingState(StreamingState.Idle); } else { const tool = toolRegistry.getTool(request.name); if (!tool) { @@ -469,7 +514,7 @@ export const useGeminiStream = ( error: undefined, }; updateFunctionResponseUI(responseInfo, ToolCallStatus.Success); - + setStreamingState(StreamingState.Idle); await submitQuery(functionResponse); } }; |
