diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index bb002adb71..90cea8bbc0 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/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; @@ -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 diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 603c1cbd76..b6ffda16d1 100644 --- a/packages/cli/src/gemini.tsx +++ b/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( @@ -218,6 +216,24 @@ export async function main() { argv, ); + const wasRaw = process.stdin.isRaw; + let kittyProtocolDetectionComplete: Promise | 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; } diff --git a/packages/cli/src/ui/contexts/KeypressContext.tsx b/packages/cli/src/ui/contexts/KeypressContext.tsx index c0d1a600fe..4df48d8d6b 100644 --- a/packages/cli/src/ui/contexts/KeypressContext.tsx +++ b/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); diff --git a/packages/cli/src/ui/hooks/useKeypress.test.ts b/packages/cli/src/ui/hooks/useKeypress.test.ts index 07129c4f87..22d729e4bc 100644 --- a/packages/cli/src/ui/hooks/useKeypress.test.ts +++ b/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; diff --git a/packages/cli/src/ui/utils/kittyProtocolDetector.ts b/packages/cli/src/ui/utils/kittyProtocolDetector.ts index 5d77943afa..3355330a62 100644 --- a/packages/cli/src/ui/utils/kittyProtocolDetector.ts +++ b/packages/cli/src/ui/utils/kittyProtocolDetector.ts @@ -32,40 +32,58 @@ export async function detectAndEnableKittyProtocol(): Promise { 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 ? 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 ? c) if (responseBuffer.includes('\x1b[?') && responseBuffer.includes('c')) { - if (!checkFinished) { - checkFinished = true; - process.stdin.removeListener('data', handleData); + 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); + 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 { 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); }); }