fix: InputPrompt wrapped lines maintain highlighting, increase responsiveness in narrow cases (#7656)

Co-authored-by: Jacob Richman <jacob314@gmail.com>
Co-authored-by: Arya Gummadi <aryagummadi@google.com>

packages/cli/src/ui/AppContainer.tsx

@@ -56,6 +56,7 @@ import { useSlashCommandProcessor } from './hooks/slashCommandProcessor.js';
 import { useVimMode } from './contexts/VimModeContext.js';
 import { useConsoleMessages } from './hooks/useConsoleMessages.js';
 import { useTerminalSize } from './hooks/useTerminalSize.js';
+import { calculatePromptWidths } from './components/InputPrompt.js';
 import { useStdin, useStdout } from 'ink';
 import ansiEscapes from 'ansi-escapes';
 import * as fs from 'node:fs';
@@ -227,12 +228,12 @@ export const AppContainer = (props: AppContainerProps) => {
     registerCleanup(consolePatcher.cleanup);
   }, [handleNewMessage, config]);
 
-  const widthFraction = 0.9;
+  // Derive widths for InputPrompt using shared helper
-  const inputWidth = Math.max(
+  const { inputWidth, suggestionsWidth } = useMemo(() => {
-    20,
+    const { inputWidth, suggestionsWidth } =
-    Math.floor(terminalWidth * widthFraction) - 3,
+      calculatePromptWidths(terminalWidth);
-  );
+    return { inputWidth, suggestionsWidth };
-  const suggestionsWidth = Math.max(20, Math.floor(terminalWidth * 1.0));
+  }, [terminalWidth]);
   const mainAreaWidth = Math.floor(terminalWidth * 0.9);
   const staticAreaMaxItemHeight = Math.max(terminalHeight * 4, 100);
 

packages/cli/src/ui/components/Composer.test.tsx

@@ -50,6 +50,11 @@ vi.mock('./DetailedMessagesDisplay.js', () => ({
 
 vi.mock('./InputPrompt.js', () => ({
   InputPrompt: () => <Text>InputPrompt</Text>,
+  calculatePromptWidths: vi.fn(() => ({
+    inputWidth: 80,
+    suggestionsWidth: 40,
+    containerWidth: 84,
+  })),
 }));
 
 vi.mock('./Footer.js', () => ({

packages/cli/src/ui/components/Composer.tsx

@@ -5,12 +5,13 @@
  */
 
 import { Box, Text } from 'ink';
+import { useMemo } from 'react';
 import { LoadingIndicator } from './LoadingIndicator.js';
 import { ContextSummaryDisplay } from './ContextSummaryDisplay.js';
 import { AutoAcceptIndicator } from './AutoAcceptIndicator.js';
 import { ShellModeIndicator } from './ShellModeIndicator.js';
 import { DetailedMessagesDisplay } from './DetailedMessagesDisplay.js';
-import { InputPrompt } from './InputPrompt.js';
+import { InputPrompt, calculatePromptWidths } from './InputPrompt.js';
 import { Footer, type FooterProps } from './Footer.js';
 import { ShowMoreLines } from './ShowMoreLines.js';
 import { OverflowProvider } from '../contexts/OverflowContext.js';
@@ -39,6 +40,12 @@ export const Composer = () => {
 
   const { contextFileNames, showAutoAcceptIndicator } = uiState;
 
+  // Use the container width of InputPrompt for width of DetailedMessagesDisplay
+  const { containerWidth } = useMemo(
+    () => calculatePromptWidths(uiState.terminalWidth),
+    [uiState.terminalWidth],
+  );
+
   // Build footer props from context values
   const footerProps: Omit<FooterProps, 'vimMode'> = {
     model: config.getModel(),
@@ -163,7 +170,7 @@ export const Composer = () => {
               maxHeight={
                 uiState.constrainHeight ? debugConsoleMaxHeight : undefined
               }
-              width={uiState.inputWidth}
+              width={containerWidth}
             />
             <ShowMoreLines constrainHeight={uiState.constrainHeight} />
           </Box>

packages/cli/src/ui/components/InputPrompt.test.tsx

@@ -113,6 +113,7 @@ describe('InputPrompt', () => {
         mockBuffer.cursor = [0, newText.length];
         mockBuffer.viewportVisualLines = [newText];
         mockBuffer.allVisualLines = [newText];
+        mockBuffer.visualToLogicalMap = [[0, 0]];
       }),
       replaceRangeByOffset: vi.fn(),
       viewportVisualLines: [''],
@@ -138,6 +139,7 @@ describe('InputPrompt', () => {
       replaceRange: vi.fn(),
       deleteWordLeft: vi.fn(),
       deleteWordRight: vi.fn(),
+      visualToLogicalMap: [[0, 0]],
     } as unknown as TextBuffer;
 
     mockShellHistory = {
@@ -1344,6 +1346,12 @@ describe('InputPrompt', () => {
       mockBuffer.viewportVisualLines = text.split('\n');
       mockBuffer.allVisualLines = text.split('\n');
       mockBuffer.visualCursor = [2, 5]; // cursor at the end of "world"
+      // Provide a visual-to-logical mapping for each visual line
+      mockBuffer.visualToLogicalMap = [
+        [0, 0], // 'hello' starts at col 0 of logical line 0
+        [1, 0], // '' (blank) is logical line 1, col 0
+        [2, 0], // 'world' is logical line 2, col 0
+      ];
 
       const { stdout, unmount } = renderWithProviders(
         <InputPrompt {...props} />,

packages/cli/src/ui/components/InputPrompt.tsx

@@ -24,7 +24,10 @@ import { keyMatchers, Command } from '../keyMatchers.js';
 import type { CommandContext, SlashCommand } from '../commands/types.js';
 import type { Config } from '@google/gemini-cli-core';
 import { ApprovalMode } from '@google/gemini-cli-core';
-import { parseInputForHighlighting } from '../utils/highlight.js';
+import {
+  parseInputForHighlighting,
+  buildSegmentsForVisualSlice,
+} from '../utils/highlight.js';
 import {
   clipboardHasImage,
   saveClipboardImage,
@@ -52,6 +55,31 @@ export interface InputPromptProps {
   isShellFocused?: boolean;
 }
 
+// The input content, input container, and input suggestions list may have different widths
+export const calculatePromptWidths = (terminalWidth: number) => {
+  const widthFraction = 0.9;
+  const FRAME_PADDING_AND_BORDER = 4; // Border (2) + padding (2)
+  const PROMPT_PREFIX_WIDTH = 2; // '> ' or '! '
+  const MIN_CONTENT_WIDTH = 2;
+
+  const innerContentWidth =
+    Math.floor(terminalWidth * widthFraction) -
+    FRAME_PADDING_AND_BORDER -
+    PROMPT_PREFIX_WIDTH;
+
+  const inputWidth = Math.max(MIN_CONTENT_WIDTH, innerContentWidth);
+  const FRAME_OVERHEAD = FRAME_PADDING_AND_BORDER + PROMPT_PREFIX_WIDTH;
+  const containerWidth = inputWidth + FRAME_OVERHEAD;
+  const suggestionsWidth = Math.max(20, Math.floor(terminalWidth * 1.0));
+
+  return {
+    inputWidth,
+    containerWidth,
+    suggestionsWidth,
+    frameOverhead: FRAME_OVERHEAD,
+  } as const;
+};
+
 export const InputPrompt: React.FC<InputPromptProps> = ({
   buffer,
   onSubmit,
@@ -854,64 +882,79 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
           ) : (
             linesToRender
               .map((lineText, visualIdxInRenderedSet) => {
-                const tokens = parseInputForHighlighting(
+                const absoluteVisualIdx =
-                  lineText,
+                  scrollVisualRow + visualIdxInRenderedSet;
-                  visualIdxInRenderedSet,
+                const mapEntry = buffer.visualToLogicalMap[absoluteVisualIdx];
-                );
                 const cursorVisualRow =
                   cursorVisualRowAbsolute - scrollVisualRow;
                 const isOnCursorLine =
                   focus && visualIdxInRenderedSet === cursorVisualRow;
 
                 const renderedLine: React.ReactNode[] = [];
-                let charCount = 0;
 
-                tokens.forEach((token, tokenIdx) => {
+                const [logicalLineIdx, logicalStartCol] = mapEntry;
-                  let display = token.text;
+                const logicalLine = buffer.lines[logicalLineIdx] || '';
+                const tokens = parseInputForHighlighting(
+                  logicalLine,
+                  logicalLineIdx,
+                );
+
+                const visualStart = logicalStartCol;
+                const visualEnd = logicalStartCol + cpLen(lineText);
+                const segments = buildSegmentsForVisualSlice(
+                  tokens,
+                  visualStart,
+                  visualEnd,
+                );
+
+                let charCount = 0;
+                segments.forEach((seg, segIdx) => {
+                  const segLen = cpLen(seg.text);
+                  let display = seg.text;
+
                   if (isOnCursorLine) {
                     const relativeVisualColForHighlight =
                       cursorVisualColAbsolute;
-                    const tokenStart = charCount;
+                    const segStart = charCount;
-                    const tokenEnd = tokenStart + cpLen(token.text);
+                    const segEnd = segStart + segLen;
-
                     if (
-                      relativeVisualColForHighlight >= tokenStart &&
+                      relativeVisualColForHighlight >= segStart &&
-                      relativeVisualColForHighlight < tokenEnd
+                      relativeVisualColForHighlight < segEnd
                     ) {
                       const charToHighlight = cpSlice(
-                        token.text,
+                        seg.text,
-                        relativeVisualColForHighlight - tokenStart,
+                        relativeVisualColForHighlight - segStart,
-                        relativeVisualColForHighlight - tokenStart + 1,
+                        relativeVisualColForHighlight - segStart + 1,
                       );
                       const highlighted = chalk.inverse(charToHighlight);
                       display =
                         cpSlice(
-                          token.text,
+                          seg.text,
                           0,
-                          relativeVisualColForHighlight - tokenStart,
+                          relativeVisualColForHighlight - segStart,
                         ) +
                         highlighted +
                         cpSlice(
-                          token.text,
+                          seg.text,
-                          relativeVisualColForHighlight - tokenStart + 1,
+                          relativeVisualColForHighlight - segStart + 1,
                         );
                     }
-                    charCount = tokenEnd;
+                    charCount = segEnd;
                   }
 
                   const color =
-                    token.type === 'command' || token.type === 'file'
+                    seg.type === 'command' || seg.type === 'file'
                       ? theme.text.accent
                       : theme.text.primary;
 
                   renderedLine.push(
-                    <Text key={`token-${tokenIdx}`} color={color}>
+                    <Text key={`token-${segIdx}`} color={color}>
                       {display}
                     </Text>,
                   );
                 });
-                const currentLineGhost = isOnCursorLine ? inlineGhost : '';
 
+                const currentLineGhost = isOnCursorLine ? inlineGhost : '';
                 if (
                   isOnCursorLine &&
                   cursorVisualColAbsolute === cpLen(lineText)

packages/cli/src/ui/components/shared/text-buffer.ts

@@ -1568,7 +1568,7 @@ export function useTextBuffer({
     [visualLayout, cursorRow, cursorCol],
   );
 
-  const { visualLines } = visualLayout;
+  const { visualLines, visualToLogicalMap } = visualLayout;
 
   const [visualScrollRow, setVisualScrollRow] = useState<number>(0);
 
@@ -1588,6 +1588,8 @@ export function useTextBuffer({
   // Update visual scroll (vertical)
   useEffect(() => {
     const { height } = viewport;
+    const totalVisualLines = visualLines.length;
+    const maxScrollStart = Math.max(0, totalVisualLines - height);
     let newVisualScrollRow = visualScrollRow;
 
     if (visualCursor[0] < visualScrollRow) {
@@ -1595,10 +1597,15 @@ export function useTextBuffer({
     } else if (visualCursor[0] >= visualScrollRow + height) {
       newVisualScrollRow = visualCursor[0] - height + 1;
     }
+
+    // When the number of visual lines shrinks (e.g., after widening the viewport),
+    // ensure scroll never starts beyond the last valid start so we can render a full window.
+    newVisualScrollRow = clamp(newVisualScrollRow, 0, maxScrollStart);
+
     if (newVisualScrollRow !== visualScrollRow) {
       setVisualScrollRow(newVisualScrollRow);
     }
-  }, [visualCursor, visualScrollRow, viewport]);
+  }, [visualCursor, visualScrollRow, viewport, visualLines.length]);
 
   const insert = useCallback(
     (ch: string, { paste = false }: { paste?: boolean } = {}): void => {
@@ -1988,6 +1995,7 @@ export function useTextBuffer({
     viewportVisualLines: renderedVisualLines,
     visualCursor,
     visualScrollRow,
+    visualToLogicalMap,
 
     setText,
     insert,
@@ -2063,6 +2071,12 @@ export interface TextBuffer {
   viewportVisualLines: string[]; // The subset of visual lines to be rendered based on visualScrollRow and viewport.height
   visualCursor: [number, number]; // Visual cursor [row, col] relative to the start of all visualLines
   visualScrollRow: number; // Scroll position for visual lines (index of the first visible visual line)
+  /**
+   * For each visual line (by absolute index in allVisualLines) provides a tuple
+   * [logicalLineIndex, startColInLogical] that maps where that visual line
+   * begins within the logical buffer. Indices are code-point based.
+   */
+  visualToLogicalMap: Array<[number, number]>;
 
   // Actions
 

packages/cli/src/ui/utils/highlight.ts

@@ -4,6 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { cpLen, cpSlice } from './textUtils.js';
+
 export type HighlightToken = {
   text: string;
   type: 'default' | 'command' | 'file';
@@ -63,3 +65,39 @@ export function parseInputForHighlighting(
 
   return tokens;
 }
+
+export function buildSegmentsForVisualSlice(
+  tokens: readonly HighlightToken[],
+  sliceStart: number,
+  sliceEnd: number,
+): readonly HighlightToken[] {
+  if (sliceStart >= sliceEnd) return [];
+
+  const segments: HighlightToken[] = [];
+  let tokenCpStart = 0;
+
+  for (const token of tokens) {
+    const tokenLen = cpLen(token.text);
+    const tokenStart = tokenCpStart;
+    const tokenEnd = tokenStart + tokenLen;
+
+    const overlapStart = Math.max(tokenStart, sliceStart);
+    const overlapEnd = Math.min(tokenEnd, sliceEnd);
+    if (overlapStart < overlapEnd) {
+      const sliceStartInToken = overlapStart - tokenStart;
+      const sliceEndInToken = overlapEnd - tokenStart;
+      const rawSlice = cpSlice(token.text, sliceStartInToken, sliceEndInToken);
+
+      const last = segments[segments.length - 1];
+      if (last && last.type === token.type) {
+        last.text += rawSlice;
+      } else {
+        segments.push({ type: token.type, text: rawSlice });
+      }
+    }
+
+    tokenCpStart += tokenLen;
+  }
+
+  return segments;
+}