diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 1c860501..7a7484df 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -31,7 +31,7 @@ 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 { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; import { useStateAndRef } from './useStateAndRef.js'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -174,14 +174,12 @@ export const useGeminiStream = ( signal, ); - let currentGeminiText = ''; + let geminiMessageBuffer = ''; for await (const event of stream) { if (signal.aborted) break; if (event.type === ServerGeminiEventType.Content) { - currentGeminiText += event.value; - if (pendingHistoryItemRef.current?.type !== 'gemini') { // Flush out existing pending history item. if (pendingHistoryItemRef.current) { @@ -189,19 +187,22 @@ export const useGeminiStream = ( } setPendingHistoryItem({ type: 'gemini', - text: currentGeminiText, + text: '', }); + geminiMessageBuffer = ''; } - // Split large messages for better rendering performance - const splitPoint = findSafeSplitPoint(currentGeminiText); - if (splitPoint === currentGeminiText.length) { + geminiMessageBuffer += event.value; + + // Split large messages for better rendering performance. Ideally, + // we should maximize the amount of output sent to . + const splitPoint = findLastSafeSplitPoint(geminiMessageBuffer); + if (splitPoint === geminiMessageBuffer.length) { // Update the existing message with accumulated content - setPendingHistoryItem((pending) => ({ - // There might be a more typesafe way to do this. - ...pending!, - text: currentGeminiText, - })); + setPendingHistoryItem({ + type: 'gemini', + text: geminiMessageBuffer, + }); } else { // This indicates that we need to split up this Gemini Message. // Splitting a message is primarily a performance consideration. There is a @@ -211,21 +212,19 @@ export const useGeminiStream = ( // 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 + const beforeText = geminiMessageBuffer.substring(0, splitPoint); + const afterText = geminiMessageBuffer.substring(splitPoint); + geminiMessageBuffer = afterText; // Continue accumulating from split point addItem( - { type: 'gemini_content', text: beforeText }, + { type: 'gemini', text: beforeText }, userMessageTimestamp, ); setPendingHistoryItem({ - type: 'gemini_content', + type: 'gemini', text: afterText, }); } } else if (event.type === ServerGeminiEventType.ToolCallRequest) { - currentGeminiText = ''; - const { callId, name, args } = event.value; const cliTool = toolRegistry.getTool(name); if (!cliTool) { diff --git a/packages/cli/src/ui/utils/markdownUtilities.ts b/packages/cli/src/ui/utils/markdownUtilities.ts index 94492b8c..96c2ba39 100644 --- a/packages/cli/src/ui/utils/markdownUtilities.ts +++ b/packages/cli/src/ui/utils/markdownUtilities.ts @@ -184,3 +184,35 @@ export const findSafeSplitPoint = ( // to keep the entire content as one piece. return content.length; }; + +export const findLastSafeSplitPoint = (content: string) => { + const enclosingBlockStart = findEnclosingCodeBlockStart( + content, + content.length, + ); + if (enclosingBlockStart !== -1) { + // The end of the content is contained in a code block. Split right before. + return enclosingBlockStart; + } + + // Search for the last double newline (\n\n) not in a code block. + let searchStartIndex = content.length; + while (searchStartIndex >= 0) { + const dnlIndex = content.lastIndexOf('\n\n', searchStartIndex); + if (dnlIndex === -1) { + // No more double newlines found after idealMaxLength + break; + } + + const potentialSplitPoint = dnlIndex + 2; + if (!isIndexInsideCodeBlock(content, potentialSplitPoint)) { + return potentialSplitPoint; + } + + searchStartIndex = potentialSplitPoint; // Continue search after the found \n\n + } + + // If no safe double newline found after idealMaxLength, return content.length + // to keep the entire content as one piece. + return content.length; +};