feat(auth): improve API key authentication flow (#11760)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

packages/cli/src/ui/components/shared/TextInput.test.tsx

@@ -0,0 +1,311 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { render } from 'ink-testing-library';
+import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
+import { TextInput } from './TextInput.js';
+import { useKeypress } from '../../hooks/useKeypress.js';
+import { useTextBuffer, type TextBuffer } from './text-buffer.js';
+
+// Mocks
+vi.mock('../../hooks/useKeypress.js', () => ({
+  useKeypress: vi.fn(),
+}));
+
+vi.mock('./text-buffer.js', () => {
+  const mockTextBuffer = {
+    text: '',
+    lines: [''],
+    cursor: [0, 0],
+    visualCursor: [0, 0],
+    viewportVisualLines: [''],
+    handleInput: vi.fn((key) => {
+      // Simulate basic input for testing
+      if (key.sequence) {
+        mockTextBuffer.text += key.sequence;
+        mockTextBuffer.viewportVisualLines = [mockTextBuffer.text];
+        mockTextBuffer.visualCursor[1] = mockTextBuffer.text.length;
+      } else if (key.name === 'backspace') {
+        mockTextBuffer.text = mockTextBuffer.text.slice(0, -1);
+        mockTextBuffer.viewportVisualLines = [mockTextBuffer.text];
+        mockTextBuffer.visualCursor[1] = mockTextBuffer.text.length;
+      } else if (key.name === 'left') {
+        mockTextBuffer.visualCursor[1] = Math.max(
+          0,
+          mockTextBuffer.visualCursor[1] - 1,
+        );
+      } else if (key.name === 'right') {
+        mockTextBuffer.visualCursor[1] = Math.min(
+          mockTextBuffer.text.length,
+          mockTextBuffer.visualCursor[1] + 1,
+        );
+      }
+    }),
+    setText: vi.fn((newText) => {
+      mockTextBuffer.text = newText;
+      mockTextBuffer.viewportVisualLines = [newText];
+      mockTextBuffer.visualCursor[1] = newText.length;
+    }),
+  };
+
+  return {
+    useTextBuffer: vi.fn(() => mockTextBuffer as unknown as TextBuffer),
+    TextBuffer: vi.fn(() => mockTextBuffer as unknown as TextBuffer),
+  };
+});
+
+const mockedUseKeypress = useKeypress as Mock;
+const mockedUseTextBuffer = useTextBuffer as Mock;
+
+describe('TextInput', () => {
+  const onCancel = vi.fn();
+  const onSubmit = vi.fn();
+  let mockBuffer: TextBuffer;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+    // Reset the internal state of the mock buffer for each test
+    const buffer = {
+      text: '',
+      lines: [''],
+      cursor: [0, 0],
+      visualCursor: [0, 0],
+      viewportVisualLines: [''],
+      handleInput: vi.fn((key) => {
+        if (key.sequence) {
+          buffer.text += key.sequence;
+          buffer.viewportVisualLines = [buffer.text];
+          buffer.visualCursor[1] = buffer.text.length;
+        } else if (key.name === 'backspace') {
+          buffer.text = buffer.text.slice(0, -1);
+          buffer.viewportVisualLines = [buffer.text];
+          buffer.visualCursor[1] = buffer.text.length;
+        } else if (key.name === 'left') {
+          buffer.visualCursor[1] = Math.max(0, buffer.visualCursor[1] - 1);
+        } else if (key.name === 'right') {
+          buffer.visualCursor[1] = Math.min(
+            buffer.text.length,
+            buffer.visualCursor[1] + 1,
+          );
+        }
+      }),
+      setText: vi.fn((newText) => {
+        buffer.text = newText;
+        buffer.viewportVisualLines = [newText];
+        buffer.visualCursor[1] = newText.length;
+      }),
+    };
+    mockBuffer = buffer as unknown as TextBuffer;
+    mockedUseTextBuffer.mockReturnValue(mockBuffer);
+  });
+
+  it('renders with an initial value', () => {
+    const buffer = {
+      text: 'test',
+      lines: ['test'],
+      cursor: [0, 4],
+      visualCursor: [0, 4],
+      viewportVisualLines: ['test'],
+      handleInput: vi.fn(),
+      setText: vi.fn(),
+    };
+    const { lastFrame } = render(
+      <TextInput
+        buffer={buffer as unknown as TextBuffer}
+        onSubmit={onSubmit}
+        onCancel={onCancel}
+      />,
+    );
+    expect(lastFrame()).toContain('test');
+  });
+
+  it('renders a placeholder', () => {
+    const buffer = {
+      text: '',
+      lines: [''],
+      cursor: [0, 0],
+      visualCursor: [0, 0],
+      viewportVisualLines: [''],
+      handleInput: vi.fn(),
+      setText: vi.fn(),
+    };
+    const { lastFrame } = render(
+      <TextInput
+        buffer={buffer as unknown as TextBuffer}
+        placeholder="testing"
+        onSubmit={onSubmit}
+        onCancel={onCancel}
+      />,
+    );
+    expect(lastFrame()).toContain('testing');
+  });
+
+  it('handles character input', () => {
+    render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'a',
+      sequence: 'a',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(mockBuffer.handleInput).toHaveBeenCalledWith({
+      name: 'a',
+      sequence: 'a',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+    expect(mockBuffer.text).toBe('a');
+  });
+
+  it('handles backspace', () => {
+    mockBuffer.setText('test');
+    render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'backspace',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(mockBuffer.handleInput).toHaveBeenCalledWith({
+      name: 'backspace',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+    expect(mockBuffer.text).toBe('tes');
+  });
+
+  it('handles left arrow', () => {
+    mockBuffer.setText('test');
+    render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'left',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    // Cursor moves from end to before 't'
+    expect(mockBuffer.visualCursor[1]).toBe(3);
+  });
+
+  it('handles right arrow', () => {
+    mockBuffer.setText('test');
+    mockBuffer.visualCursor[1] = 2; // Set initial cursor for right arrow test
+    render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'right',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(mockBuffer.visualCursor[1]).toBe(3);
+  });
+
+  it('calls onSubmit on return', () => {
+    mockBuffer.setText('test');
+    render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'return',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+
+    expect(onSubmit).toHaveBeenCalledWith('test');
+  });
+
+  it('calls onCancel on escape', async () => {
+    vi.useFakeTimers();
+    render(
+      <TextInput buffer={mockBuffer} onCancel={onCancel} onSubmit={onSubmit} />,
+    );
+    const keypressHandler = mockedUseKeypress.mock.calls[0][0];
+
+    keypressHandler({
+      name: 'escape',
+      sequence: '',
+      ctrl: false,
+      meta: false,
+      shift: false,
+      paste: false,
+    });
+    await vi.runAllTimersAsync();
+
+    expect(onCancel).toHaveBeenCalled();
+    vi.useRealTimers();
+  });
+
+  it('renders the input value', () => {
+    mockBuffer.setText('secret');
+    const { lastFrame } = render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+    expect(lastFrame()).toContain('secret');
+  });
+
+  it('does not show cursor when not focused', () => {
+    mockBuffer.setText('test');
+    const { lastFrame } = render(
+      <TextInput
+        buffer={mockBuffer}
+        focus={false}
+        onSubmit={onSubmit}
+        onCancel={onCancel}
+      />,
+    );
+    expect(lastFrame()).not.toContain('\u001b[7m'); // Inverse video chalk
+  });
+
+  it('renders multiple lines when text wraps', () => {
+    mockBuffer.text = 'line1\nline2';
+    mockBuffer.viewportVisualLines = ['line1', 'line2'];
+
+    const { lastFrame } = render(
+      <TextInput buffer={mockBuffer} onSubmit={onSubmit} onCancel={onCancel} />,
+    );
+
+    expect(lastFrame()).toContain('line1');
+    expect(lastFrame()).toContain('line2');
+  });
+});

packages/cli/src/ui/components/shared/TextInput.tsx

@@ -0,0 +1,104 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type React from 'react';
+import { useCallback } from 'react';
+import type { Key } from '../../hooks/useKeypress.js';
+import { Text, Box } from 'ink';
+import { useKeypress } from '../../hooks/useKeypress.js';
+import chalk from 'chalk';
+import { theme } from '../../semantic-colors.js';
+import type { TextBuffer } from './text-buffer.js';
+import { cpSlice } from '../../utils/textUtils.js';
+
+export interface TextInputProps {
+  buffer: TextBuffer;
+  placeholder?: string;
+  onSubmit?: (value: string) => void;
+  onCancel?: () => void;
+  focus?: boolean;
+}
+
+export function TextInput({
+  buffer,
+  placeholder = '',
+  onSubmit,
+  onCancel,
+  focus = true,
+}: TextInputProps): React.JSX.Element {
+  const {
+    text,
+    handleInput,
+    visualCursor,
+    viewportVisualLines,
+    visualScrollRow,
+  } = buffer;
+  const [cursorVisualRowAbsolute, cursorVisualColAbsolute] = visualCursor;
+
+  const handleKeyPress = useCallback(
+    (key: Key) => {
+      if (key.name === 'escape') {
+        onCancel?.();
+        return;
+      }
+
+      if (key.name === 'return') {
+        onSubmit?.(text);
+        return;
+      }
+
+      handleInput(key);
+    },
+    [handleInput, onCancel, onSubmit, text],
+  );
+
+  useKeypress(handleKeyPress, { isActive: focus });
+
+  const showPlaceholder = text.length === 0 && placeholder;
+
+  if (showPlaceholder) {
+    return (
+      <Box>
+        {focus ? (
+          <Text>
+            {chalk.inverse(placeholder[0] || ' ')}
+            <Text color={theme.text.secondary}>{placeholder.slice(1)}</Text>
+          </Text>
+        ) : (
+          <Text color={theme.text.secondary}>{placeholder}</Text>
+        )}
+      </Box>
+    );
+  }
+
+  return (
+    <Box flexDirection="column">
+      {viewportVisualLines.map((lineText, idx) => {
+        const currentVisualRow = visualScrollRow + idx;
+        const isCursorLine =
+          focus && currentVisualRow === cursorVisualRowAbsolute;
+
+        const lineDisplay = isCursorLine
+          ? cpSlice(lineText, 0, cursorVisualColAbsolute) +
+            chalk.inverse(
+              cpSlice(
+                lineText,
+                cursorVisualColAbsolute,
+                cursorVisualColAbsolute + 1,
+              ) || ' ',
+            ) +
+            cpSlice(lineText, cursorVisualColAbsolute + 1)
+          : lineText;
+
+        return (
+          <Box key={idx} height={1}>
+            <Text>{lineDisplay}</Text>
+          </Box>
+        );
+      })}
+    </Box>
+  );
+}

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

@@ -14,6 +14,7 @@ import type {
   TextBufferState,
   TextBufferAction,
   VisualLayout,
+  TextBufferOptions,
 } from './text-buffer.js';
 import {
   useTextBuffer,
@@ -101,6 +102,43 @@ describe('textBufferReducer', () => {
     });
   });
 
+  describe('insert action with options', () => {
+    it('should filter input using inputFilter option', () => {
+      const action: TextBufferAction = { type: 'insert', payload: 'a1b2c3' };
+      const options: TextBufferOptions = {
+        inputFilter: (text) => text.replace(/[0-9]/g, ''),
+      };
+      const state = textBufferReducer(initialState, action, options);
+      expect(state.lines).toEqual(['abc']);
+      expect(state.cursorCol).toBe(3);
+    });
+
+    it('should strip newlines when singleLine option is true', () => {
+      const action: TextBufferAction = {
+        type: 'insert',
+        payload: 'hello\nworld',
+      };
+      const options: TextBufferOptions = { singleLine: true };
+      const state = textBufferReducer(initialState, action, options);
+      expect(state.lines).toEqual(['helloworld']);
+      expect(state.cursorCol).toBe(10);
+    });
+
+    it('should apply both inputFilter and singleLine options', () => {
+      const action: TextBufferAction = {
+        type: 'insert',
+        payload: 'h\ne\nl\nl\no\n1\n2\n3',
+      };
+      const options: TextBufferOptions = {
+        singleLine: true,
+        inputFilter: (text) => text.replace(/[0-9]/g, ''),
+      };
+      const state = textBufferReducer(initialState, action, options);
+      expect(state.lines).toEqual(['hello']);
+      expect(state.cursorCol).toBe(5);
+    });
+  });
+
   describe('backspace action', () => {
     it('should remove a character', () => {
       const stateWithText: TextBufferState = {
@@ -1520,6 +1558,75 @@ Contrary to popular belief, Lorem Ipsum is not simply random text. It has roots
     });
   });
 
+  describe('inputFilter', () => {
+    it('should filter input based on the provided filter function', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          inputFilter: (text) => text.replace(/[^0-9]/g, ''),
+        }),
+      );
+
+      act(() => result.current.insert('a1b2c3'));
+      expect(getBufferState(result).text).toBe('123');
+    });
+
+    it('should handle empty result from filter', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          inputFilter: (text) => text.replace(/[^0-9]/g, ''),
+        }),
+      );
+
+      act(() => result.current.insert('abc'));
+      expect(getBufferState(result).text).toBe('');
+    });
+
+    it('should filter pasted text', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          inputFilter: (text) => text.toUpperCase(),
+        }),
+      );
+
+      act(() => result.current.insert('hello', { paste: true }));
+      expect(getBufferState(result).text).toBe('HELLO');
+    });
+
+    it('should not filter newlines if they are allowed by the filter', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          inputFilter: (text) => text, // Allow everything including newlines
+        }),
+      );
+
+      act(() => result.current.insert('a\nb'));
+      // The insert function splits by newline and inserts separately if it detects them.
+      // If the filter allows them, they should be handled correctly by the subsequent logic in insert.
+      expect(getBufferState(result).text).toBe('a\nb');
+    });
+
+    it('should filter before newline check in insert', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          inputFilter: (text) => text.replace(/\n/g, ''), // Filter out newlines
+        }),
+      );
+
+      act(() => result.current.insert('a\nb'));
+      expect(getBufferState(result).text).toBe('ab');
+    });
+  });
+
   describe('stripAnsi', () => {
     it('should correctly strip ANSI escape codes', () => {
       const textWithAnsi = '\x1B[31mHello\x1B[0m World';
@@ -1587,6 +1694,74 @@ Contrary to popular belief, Lorem Ipsum is not simply random text. It has roots
       expect(getBufferState(result).text).toBe('hello world');
     });
   });
+
+  describe('singleLine mode', () => {
+    it('should not insert a newline character when singleLine is true', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          singleLine: true,
+        }),
+      );
+      act(() => result.current.insert('\n'));
+      const state = getBufferState(result);
+      expect(state.text).toBe('');
+      expect(state.lines).toEqual(['']);
+    });
+
+    it('should not create a new line when newline() is called and singleLine is true', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          initialText: 'ab',
+          viewport,
+          isValidPath: () => false,
+          singleLine: true,
+        }),
+      );
+      act(() => result.current.move('end')); // cursor at [0,2]
+      act(() => result.current.newline());
+      const state = getBufferState(result);
+      expect(state.text).toBe('ab');
+      expect(state.lines).toEqual(['ab']);
+      expect(state.cursor).toEqual([0, 2]);
+    });
+
+    it('should not handle "Enter" key as newline when singleLine is true', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          singleLine: true,
+        }),
+      );
+      act(() =>
+        result.current.handleInput({
+          name: 'return',
+          ctrl: false,
+          meta: false,
+          shift: false,
+          paste: false,
+          sequence: '\r',
+        }),
+      );
+      expect(getBufferState(result).lines).toEqual(['']);
+    });
+
+    it('should strip newlines from pasted text when singleLine is true', () => {
+      const { result } = renderHook(() =>
+        useTextBuffer({
+          viewport,
+          isValidPath: () => false,
+          singleLine: true,
+        }),
+      );
+      act(() => result.current.insert('hello\nworld', { paste: true }));
+      const state = getBufferState(result);
+      expect(state.text).toBe('helloworld');
+      expect(state.lines).toEqual(['helloworld']);
+    });
+  });
 });
 
 describe('offsetToLogicalPos', () => {

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

@@ -518,6 +518,8 @@ interface UseTextBufferProps {
   onChange?: (text: string) => void; // Callback for when text changes
   isValidPath: (path: string) => boolean;
   shellModeActive?: boolean; // Whether the text buffer is in shell mode
+  inputFilter?: (text: string) => string; // Optional filter for input text
+  singleLine?: boolean;
 }
 
 interface UndoHistoryEntry {
@@ -949,9 +951,15 @@ export type TextBufferAction =
   | { type: 'vim_move_to_line'; payload: { lineNumber: number } }
   | { type: 'vim_escape_insert_mode' };
 
+export interface TextBufferOptions {
+  inputFilter?: (text: string) => string;
+  singleLine?: boolean;
+}
+
 function textBufferReducerLogic(
   state: TextBufferState,
   action: TextBufferAction,
+  options: TextBufferOptions = {},
 ): TextBufferState {
   const pushUndoLocal = pushUndo;
 
@@ -986,8 +994,20 @@ function textBufferReducerLogic(
 
       const currentLine = (r: number) => newLines[r] ?? '';
 
+      let payload = action.payload;
+      if (options.singleLine) {
+        payload = payload.replace(/[\r\n]/g, '');
+      }
+      if (options.inputFilter) {
+        payload = options.inputFilter(payload);
+      }
+
+      if (payload.length === 0) {
+        return state;
+      }
+
       const str = stripUnsafeCharacters(
-        action.payload.replace(/\r\n/g, '\n').replace(/\r/g, '\n'),
+        payload.replace(/\r\n/g, '\n').replace(/\r/g, '\n'),
       );
       const parts = str.split('\n');
       const lineContent = currentLine(newCursorRow);
@@ -1498,8 +1518,9 @@ function textBufferReducerLogic(
 export function textBufferReducer(
   state: TextBufferState,
   action: TextBufferAction,
+  options: TextBufferOptions = {},
 ): TextBufferState {
-  const newState = textBufferReducerLogic(state, action);
+  const newState = textBufferReducerLogic(state, action, options);
 
   if (
     newState.lines !== state.lines ||
@@ -1525,6 +1546,8 @@ export function useTextBuffer({
   onChange,
   isValidPath,
   shellModeActive = false,
+  inputFilter,
+  singleLine = false,
 }: UseTextBufferProps): TextBuffer {
   const initialState = useMemo((): TextBufferState => {
     const lines = initialText.split('\n');
@@ -1551,7 +1574,11 @@ export function useTextBuffer({
     };
   }, [initialText, initialCursorOffset, viewport.width, viewport.height]);
 
-  const [state, dispatch] = useReducer(textBufferReducer, initialState);
+  const [state, dispatch] = useReducer(
+    (s: TextBufferState, a: TextBufferAction) =>
+      textBufferReducer(s, a, { inputFilter, singleLine }),
+    initialState,
+  );
   const {
     lines,
     cursorRow,
@@ -1609,7 +1636,7 @@ export function useTextBuffer({
 
   const insert = useCallback(
     (ch: string, { paste = false }: { paste?: boolean } = {}): void => {
-      if (/[\n\r]/.test(ch)) {
+      if (!singleLine && /[\n\r]/.test(ch)) {
         dispatch({ type: 'insert', payload: ch });
         return;
       }
@@ -1648,12 +1675,15 @@ export function useTextBuffer({
         dispatch({ type: 'insert', payload: currentText });
       }
     },
-    [isValidPath, shellModeActive],
+    [isValidPath, shellModeActive, singleLine],
   );
 
   const newline = useCallback((): void => {
+    if (singleLine) {
+      return;
+    }
     dispatch({ type: 'insert', payload: '\n' });
-  }, []);
+  }, [singleLine]);
 
   const backspace = useCallback((): void => {
     dispatch({ type: 'backspace' });
@@ -1895,10 +1925,11 @@ export function useTextBuffer({
       }
 
       if (
-        key.name === 'return' ||
-        input === '\r' ||
-        input === '\n' ||
-        input === '\\\r' // VSCode terminal represents shift + enter this way
+        !singleLine &&
+        (key.name === 'return' ||
+          input === '\r' ||
+          input === '\n' ||
+          input === '\\r') // VSCode terminal represents shift + enter this way
       )
         newline();
       else if (key.name === 'left' && !key.meta && !key.ctrl) move('left');
@@ -1947,6 +1978,7 @@ export function useTextBuffer({
       insert,
       undo,
       redo,
+      singleLine,
     ],
   );