Fix issues where escape codes could end up on startup in the input prompt (google-gemini#7267)

Author: jacob314

packages/cli/src/gemini.test.tsx

@@ -4,8 +4,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import {
+  describe,
+  it,
+  expect,
+  vi,
+  beforeEach,
+  afterEach,
+  type MockInstance,
+} from 'vitest';
+import {
+  main,
   setupUnhandledRejectionHandler,
   validateDnsResolutionOrder,
   startInteractiveUI,
@@ -33,14 +42,10 @@ vi.mock('./config/settings.js', async (importOriginal) => {
 
 vi.mock('./config/config.js', () => ({
   loadCliConfig: vi.fn().mockResolvedValue({
-    config: {
-      getSandbox: vi.fn(() => false),
-      getQuestion: vi.fn(() => ''),
-    },
-    modelWasSwitched: false,
-    originalModelBeforeSwitch: null,
-    finalModel: 'test-model',
-  }),
+    getSandbox: vi.fn(() => false),
+    getQuestion: vi.fn(() => ''),
+  } as unknown as Config),
+  parseArguments: vi.fn().mockResolvedValue({ sessionSummary: null }),
 }));
 
 vi.mock('read-package-up', () => ({
@@ -157,6 +162,87 @@ describe('gemini.tsx main function', () => {
   });
 });
 
+describe('gemini.tsx main function kitty protocol', () => {
+  let setRawModeSpy: MockInstance<
+    (mode: boolean) => NodeJS.ReadStream & { fd: 0 }
+  >;
+
+  beforeEach(() => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    if (!(process.stdin as any).setRawMode) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (process.stdin as any).setRawMode = vi.fn();
+    }
+    setRawModeSpy = vi.spyOn(process.stdin, 'setRawMode');
+  });
+
+  it('should call setRawMode and detectAndEnableKittyProtocol when isInteractive is true', async () => {
+    const { detectAndEnableKittyProtocol } = await import(
+      './ui/utils/kittyProtocolDetector.js'
+    );
+    const { loadCliConfig, parseArguments } = await import(
+      './config/config.js'
+    );
+    const { loadSettings } = await import('./config/settings.js');
+    vi.mocked(loadCliConfig).mockResolvedValue({
+      isInteractive: () => true,
+      getQuestion: () => '',
+      getSandbox: () => false,
+      getDebugMode: () => false,
+      getListExtensions: () => false,
+      getMcpServers: () => ({}),
+      initialize: vi.fn(),
+      getIdeMode: () => false,
+      getExperimentalZedIntegration: () => false,
+      getScreenReader: () => false,
+    } as unknown as Config);
+    vi.mocked(loadSettings).mockReturnValue({
+      errors: [],
+      merged: {
+        advanced: {},
+        security: { auth: {} },
+        ui: {},
+      },
+      setValue: vi.fn(),
+    } as never);
+    vi.mocked(parseArguments).mockResolvedValue({
+      model: undefined,
+      sandbox: undefined,
+      sandboxImage: undefined,
+      debug: undefined,
+      prompt: undefined,
+      promptInteractive: undefined,
+      allFiles: undefined,
+      showMemoryUsage: undefined,
+      yolo: undefined,
+      approvalMode: undefined,
+      telemetry: undefined,
+      checkpointing: undefined,
+      telemetryTarget: undefined,
+      telemetryOtlpEndpoint: undefined,
+      telemetryOtlpProtocol: undefined,
+      telemetryLogPrompts: undefined,
+      telemetryOutfile: undefined,
+      allowedMcpServerNames: undefined,
+      allowedTools: undefined,
+      experimentalAcp: undefined,
+      extensions: undefined,
+      listExtensions: undefined,
+      proxy: undefined,
+      includeDirectories: undefined,
+      screenReader: undefined,
+      useSmartEdit: undefined,
+      sessionSummary: undefined,
+      promptWords: undefined,
+    });
+
+    await main();
+
+    expect(setRawModeSpy).toHaveBeenCalledWith(true);
+    expect(detectAndEnableKittyProtocol).toHaveBeenCalledTimes(1);
+  });
+});
+
 describe('validateDnsResolutionOrder', () => {
   let consoleWarnSpy: ReturnType<typeof vi.spyOn>;
 
@@ -213,7 +299,7 @@ describe('startInteractiveUI', () => {
   }));
 
   vi.mock('./ui/utils/kittyProtocolDetector.js', () => ({
-    detectAndEnableKittyProtocol: vi.fn(() => Promise.resolve()),
+    detectAndEnableKittyProtocol: vi.fn(() => Promise.resolve(true)),
   }));
 
   vi.mock('./ui/utils/updateCheck.js', () => ({
@@ -260,9 +346,6 @@ describe('startInteractiveUI', () => {
 
   it('should perform all startup tasks in correct order', async () => {
     const { getCliVersion } = await import('./utils/version.js');
-    const { detectAndEnableKittyProtocol } = await import(
-      './ui/utils/kittyProtocolDetector.js'
-    );
     const { checkForUpdates } = await import('./ui/utils/updateCheck.js');
     const { registerCleanup } = await import('./utils/cleanup.js');
 
@@ -275,7 +358,6 @@ describe('startInteractiveUI', () => {
 
     // Verify all startup tasks were called
     expect(getCliVersion).toHaveBeenCalledTimes(1);
-    expect(detectAndEnableKittyProtocol).toHaveBeenCalledTimes(1);
     expect(registerCleanup).toHaveBeenCalledTimes(1);
 
     // Verify cleanup handler is registered with unmount function

packages/cli/src/gemini.tsx

@@ -172,8 +172,6 @@ export async function startInteractiveUI(
   workspaceRoot: string = process.cwd(),
 ) {
   const version = await getCliVersion();
-  // Detect and enable Kitty keyboard protocol once at startup
-  await detectAndEnableKittyProtocol();
   setWindowTitle(basename(workspaceRoot), settings);
   const instance = render(
     <React.StrictMode>
@@ -218,6 +216,24 @@ export async function main() {
     argv,
   );
 
+  const wasRaw = process.stdin.isRaw;
+  let kittyProtocolDetectionComplete: Promise<boolean> | undefined;
+  if (config.isInteractive() && !wasRaw) {
+    // Set this as early as possible to avoid spurious characters from
+    // input showing up in the output.
+    process.stdin.setRawMode(true);
+
+    // This cleanup isn't strictly needed but may help in certain situations.
+    process.on('SIGTERM', () => {
+      process.stdin.setRawMode(wasRaw);
+    });
+    process.on('SIGINT', () => {
+      process.stdin.setRawMode(wasRaw);
+    });
+
+    // Detect and enable Kitty keyboard protocol once at startup.
+    kittyProtocolDetectionComplete = detectAndEnableKittyProtocol();
+  }
   if (argv.sessionSummary) {
     registerCleanup(() => {
       const metrics = uiTelemetryService.getMetrics();
@@ -385,6 +401,8 @@ export async function main() {
 
   // Render UI, passing necessary config values. Check that there is no command line question.
   if (config.isInteractive()) {
+    // Need kitty detection to be complete before we can start the interactive UI.
+    await kittyProtocolDetectionComplete;
     await startInteractiveUI(config, settings, startupWarnings);
     return;
   }

packages/cli/src/ui/contexts/KeypressContext.tsx

@@ -115,7 +115,10 @@ export function KeypressProvider({
       }
     };
 
-    setRawMode(true);
+    const wasRaw = stdin.isRaw;
+    if (wasRaw === false) {
+      setRawMode(true);
+    }
 
     const keypressStream = new PassThrough();
     let usePassthrough = false;
@@ -677,7 +680,9 @@ export function KeypressProvider({
       rl.close();
 
       // Restore the terminal to its original state.
-      setRawMode(false);
+      if (wasRaw === false) {
+        setRawMode(false);
+      }
 
       if (backslashTimeout) {
         clearTimeout(backslashTimeout);

packages/cli/src/ui/hooks/useKeypress.test.ts

@@ -54,6 +54,7 @@ vi.mock('readline', () => {
 
 class MockStdin extends EventEmitter {
   isTTY = true;
+  isRaw = false;
   setRawMode = vi.fn();
   on = this.addListener;
   removeListener = this.removeListener;

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

@@ -32,40 +32,58 @@ export async function detectAndEnableKittyProtocol(): Promise<boolean> {
 
     let responseBuffer = '';
     let progressiveEnhancementReceived = false;
-    let checkFinished = false;
+    let timeoutId: NodeJS.Timeout | undefined;
+
+    const onTimeout = () => {
+      timeoutId = undefined;
+      process.stdin.removeListener('data', handleData);
+      if (!originalRawMode) {
+        process.stdin.setRawMode(false);
+      }
+      detectionComplete = true;
+      resolve(false);
+    };
 
     const handleData = (data: Buffer) => {
+      if (timeoutId === undefined) {
+        // Race condition. We have already timed out.
+        return;
+      }
       responseBuffer += data.toString();
 
       // Check for progressive enhancement response (CSI ? <flags> u)
       if (responseBuffer.includes('\x1b[?') && responseBuffer.includes('u')) {
         progressiveEnhancementReceived = true;
+        // Give more time to get the full set of kitty responses if we have an
+        // indication the terminal probably supports kitty and we just need to
+        // wait a bit longer for a response.
+        clearTimeout(timeoutId);
+        timeoutId = setTimeout(onTimeout, 1000);
       }
 
       // Check for device attributes response (CSI ? <attrs> c)
       if (responseBuffer.includes('\x1b[?') && responseBuffer.includes('c')) {
-        if (!checkFinished) {
-          checkFinished = true;
-          process.stdin.removeListener('data', handleData);
-
-          if (!originalRawMode) {
-            process.stdin.setRawMode(false);
-          }
-
-          if (progressiveEnhancementReceived) {
-            // Enable the protocol
-            process.stdout.write('\x1b[>1u');
-            protocolSupported = true;
-            protocolEnabled = true;
-
-            // Set up cleanup on exit
-            process.on('exit', disableProtocol);
-            process.on('SIGTERM', disableProtocol);
-          }
-
-          detectionComplete = true;
-          resolve(protocolSupported);
+        clearTimeout(timeoutId);
+        timeoutId = undefined;
+        process.stdin.removeListener('data', handleData);
+
+        if (!originalRawMode) {
+          process.stdin.setRawMode(false);
         }
+
+        if (progressiveEnhancementReceived) {
+          // Enable the protocol
+          process.stdout.write('\x1b[>1u');
+          protocolSupported = true;
+          protocolEnabled = true;
+
+          // Set up cleanup on exit
+          process.on('exit', disableProtocol);
+          process.on('SIGTERM', disableProtocol);
+        }
+
+        detectionComplete = true;
+        resolve(protocolSupported);
       }
     };
 
@@ -75,17 +93,10 @@ export async function detectAndEnableKittyProtocol(): Promise<boolean> {
     process.stdout.write('\x1b[?u'); // Query progressive enhancement
     process.stdout.write('\x1b[c'); // Query device attributes
 
-    // Timeout after 50ms
-    setTimeout(() => {
-      if (!checkFinished) {
-        process.stdin.removeListener('data', handleData);
-        if (!originalRawMode) {
-          process.stdin.setRawMode(false);
-        }
-        detectionComplete = true;
-        resolve(false);
-      }
-    }, 50);
+    // Timeout after 200ms
+    // When a iterm2 terminal does not have focus this can take over 90s on a
+    // fast macbook so we need a somewhat longer threshold than would be ideal.
+    timeoutId = setTimeout(onTimeout, 200);
   });
 }