Fix multiple bugs with auth flow including using the implemented but unused restart support. (#13565)

Author: jacob314

packages/cli/index.ts

@@ -8,23 +8,26 @@
 
 import './src/gemini.js';
 import { main } from './src/gemini.js';
-import { debugLogger, FatalError } from '@google/gemini-cli-core';
+import { FatalError, writeToStderr } from '@google/gemini-cli-core';
+import { runExitCleanup } from './src/utils/cleanup.js';
 
 // --- Global Entry Point ---
-main().catch((error) => {
+main().catch(async (error) => {
+  await runExitCleanup();
+
   if (error instanceof FatalError) {
     let errorMessage = error.message;
     if (!process.env['NO_COLOR']) {
       errorMessage = `\x1b[31m${errorMessage}\x1b[0m`;
     }
-    debugLogger.error(errorMessage);
+    writeToStderr(errorMessage + '\n');
     process.exit(error.exitCode);
   }
-  debugLogger.error('An unexpected critical error occurred:');
+  writeToStderr('An unexpected critical error occurred:');
   if (error instanceof Error) {
-    debugLogger.error(error.stack);
+    writeToStderr(error.stack + '\n');
   } else {
-    debugLogger.error(String(error));
+    writeToStderr(String(error) + '\n');
   }
   process.exit(1);
 });

packages/cli/src/gemini.tsx

@@ -33,13 +33,11 @@ import {
   runExitCleanup,
 } from './utils/cleanup.js';
 import { getCliVersion } from './utils/version.js';
-import type {
-  Config,
-  ResumedSessionData,
-  OutputPayload,
-  ConsoleLogPayload,
-} from '@google/gemini-cli-core';
 import {
+  type Config,
+  type ResumedSessionData,
+  type OutputPayload,
+  type ConsoleLogPayload,
   sessionId,
   logUserPrompt,
   AuthType,
@@ -53,6 +51,11 @@ import {
   patchStdio,
   writeToStdout,
   writeToStderr,
+  disableMouseEvents,
+  enableMouseEvents,
+  enterAlternateScreen,
+  disableLineWrapping,
+  shouldEnterAlternateScreen,
 } from '@google/gemini-cli-core';
 import {
   initializeApp,
@@ -85,9 +88,7 @@ import { deleteSession, listSessions } from './utils/sessions.js';
 import { ExtensionManager } from './config/extension-manager.js';
 import { createPolicyUpdater } from './config/policy.js';
 import { requestConsentNonInteractive } from './config/extensions/consent.js';
-import { disableMouseEvents, enableMouseEvents } from './ui/utils/mouse.js';
 import { ScrollProvider } from './ui/contexts/ScrollProvider.js';
-import ansiEscapes from 'ansi-escapes';
 import { isAlternateBufferEnabled } from './ui/hooks/useAlternateBuffer.js';
 
 import { profiler } from './ui/components/DebugProfiler.js';
@@ -176,8 +177,10 @@ export async function startInteractiveUI(
   // as there is no benefit of alternate buffer mode when using a screen reader
   // and the Ink alternate buffer mode requires line wrapping harmful to
   // screen readers.
-  const useAlternateBuffer =
-    isAlternateBufferEnabled(settings) && !config.getScreenReader();
+  const useAlternateBuffer = shouldEnterAlternateScreen(
+    isAlternateBufferEnabled(settings),
+    config.getScreenReader(),
+  );
   const mouseEventsEnabled = useAlternateBuffer;
   if (mouseEventsEnabled) {
     enableMouseEvents();
@@ -481,8 +484,14 @@ export async function main() {
       // input showing up in the output.
       process.stdin.setRawMode(true);
 
-      if (isAlternateBufferEnabled(settings)) {
-        writeToStdout(ansiEscapes.enterAlternativeScreen);
+      if (
+        shouldEnterAlternateScreen(
+          isAlternateBufferEnabled(settings),
+          config.getScreenReader(),
+        )
+      ) {
+        enterAlternateScreen();
+        disableLineWrapping();
 
         // Ink will cleanup so there is no need for us to manually cleanup.
       }

packages/cli/src/ui/AppContainer.test.tsx

@@ -69,6 +69,8 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
       stdout: process.stdout,
       stderr: process.stderr,
     })),
+    enableMouseEvents: vi.fn(),
+    disableMouseEvents: vi.fn(),
   };
 });
 import type { LoadedSettings } from '../config/settings.js';
@@ -137,10 +139,6 @@ vi.mock('../utils/events.js');
 vi.mock('../utils/handleAutoUpdate.js');
 vi.mock('./utils/ConsolePatcher.js');
 vi.mock('../utils/cleanup.js');
-vi.mock('./utils/mouse.js', () => ({
-  enableMouseEvents: vi.fn(),
-  disableMouseEvents: vi.fn(),
-}));
 
 import { useHistory } from './hooks/useHistoryManager.js';
 import { useThemeCommand } from './hooks/useThemeCommand.js';
@@ -165,9 +163,13 @@ import { useLoadingIndicator } from './hooks/useLoadingIndicator.js';
 import { useKeypress, type Key } from './hooks/useKeypress.js';
 import { measureElement } from 'ink';
 import { useTerminalSize } from './hooks/useTerminalSize.js';
-import { ShellExecutionService, writeToStdout } from '@google/gemini-cli-core';
+import {
+  ShellExecutionService,
+  writeToStdout,
+  enableMouseEvents,
+  disableMouseEvents,
+} from '@google/gemini-cli-core';
 import { type ExtensionManager } from '../config/extension-manager.js';
-import { enableMouseEvents, disableMouseEvents } from './utils/mouse.js';
 
 describe('AppContainer State Management', () => {
   let mockConfig: Config;

packages/cli/src/ui/AppContainer.tsx

@@ -52,6 +52,12 @@ import {
   refreshServerHierarchicalMemory,
   type ModelChangedPayload,
   type MemoryChangedPayload,
+  writeToStdout,
+  disableMouseEvents,
+  enterAlternateScreen,
+  enableMouseEvents,
+  disableLineWrapping,
+  shouldEnterAlternateScreen,
 } from '@google/gemini-cli-core';
 import { validateAuthMethod } from '../config/auth.js';
 import process from 'node:process';
@@ -92,6 +98,7 @@ import { appEvents, AppEvent } from '../utils/events.js';
 import { type UpdateObject } from './utils/updateCheck.js';
 import { setUpdateHandler } from '../utils/handleAutoUpdate.js';
 import { registerCleanup, runExitCleanup } from '../utils/cleanup.js';
+import { RELAUNCH_EXIT_CODE } from '../utils/processUtils.js';
 import { useMessageQueue } from './hooks/useMessageQueue.js';
 import { useAutoAcceptIndicator } from './hooks/useAutoAcceptIndicator.js';
 import { useSessionStats } from './contexts/SessionContext.js';
@@ -106,11 +113,9 @@ import { type ExtensionManager } from '../config/extension-manager.js';
 import { requestConsentInteractive } from '../config/extensions/consent.js';
 import { useIncludeDirsTrust } from './hooks/useIncludeDirsTrust.js';
 import { isWorkspaceTrusted } from '../config/trustedFolders.js';
-import { disableMouseEvents, enableMouseEvents } from './utils/mouse.js';
 import { useAlternateBuffer } from './hooks/useAlternateBuffer.js';
 import { useSettings } from './contexts/SettingsContext.js';
 import { enableSupportedProtocol } from './utils/kittyProtocolDetector.js';
-import { writeToStdout } from '@google/gemini-cli-core';
 
 const WARNING_PROMPT_DURATION_MS = 1000;
 const QUEUE_ERROR_DISPLAY_DURATION_MS = 3000;
@@ -372,16 +377,19 @@ export const AppContainer = (props: AppContainerProps) => {
     setHistoryRemountKey((prev) => prev + 1);
   }, [setHistoryRemountKey, isAlternateBuffer, stdout]);
   const handleEditorClose = useCallback(() => {
-    if (isAlternateBuffer) {
+    if (
+      shouldEnterAlternateScreen(isAlternateBuffer, config.getScreenReader())
+    ) {
       // The editor may have exited alternate buffer mode so we need to
       // enter it again to be safe.
-      writeToStdout(ansiEscapes.enterAlternativeScreen);
+      enterAlternateScreen();
       enableMouseEvents();
+      disableLineWrapping();
       app.rerender();
     }
     enableSupportedProtocol();
     refreshStatic();
-  }, [refreshStatic, isAlternateBuffer, app]);
+  }, [refreshStatic, isAlternateBuffer, app, config]);
 
   useEffect(() => {
     coreEvents.on(CoreEvent.ExternalEditorClosed, handleEditorClose);
@@ -458,12 +466,12 @@ export const AppContainer = (props: AppContainerProps) => {
           config.isBrowserLaunchSuppressed()
         ) {
           await runExitCleanup();
-          debugLogger.log(`
+          writeToStdout(`
 ----------------------------------------------------------------
-Logging in with Google... Please restart Gemini CLI to continue.
+Logging in with Google... Restarting Gemini CLI to continue.
 ----------------------------------------------------------------
           `);
-          process.exit(0);
+          process.exit(RELAUNCH_EXIT_CODE);
         }
       }
       setAuthState(AuthState.Authenticated);

packages/cli/src/ui/auth/AuthDialog.test.tsx

@@ -25,6 +25,7 @@ import { validateAuthMethodWithSettings } from './useAuth.js';
 import { runExitCleanup } from '../../utils/cleanup.js';
 import { clearCachedCredentialFile } from '@google/gemini-cli-core';
 import { Text } from 'ink';
+import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
 
 // Mocks
 vi.mock('@google/gemini-cli-core', async (importOriginal) => {
@@ -229,6 +230,7 @@ describe('AuthDialog', () => {
     });
 
     it('exits process for Login with Google when browser is suppressed', async () => {
+      vi.useFakeTimers();
       const exitSpy = vi
         .spyOn(process, 'exit')
         .mockImplementation(() => undefined as never);
@@ -241,14 +243,14 @@ describe('AuthDialog', () => {
         mockedRadioButtonSelect.mock.calls[0][0];
       await handleAuthSelect(AuthType.LOGIN_WITH_GOOGLE);
 
+      await vi.runAllTimersAsync();
+
       expect(mockedRunExitCleanup).toHaveBeenCalled();
-      expect(logSpy).toHaveBeenCalledWith(
-        expect.stringContaining('Please restart Gemini CLI'),
-      );
-      expect(exitSpy).toHaveBeenCalledWith(0);
+      expect(exitSpy).toHaveBeenCalledWith(RELAUNCH_EXIT_CODE);
 
       exitSpy.mockRestore();
       logSpy.mockRestore();
+      vi.useRealTimers();
     });
   });
 

packages/cli/src/ui/auth/AuthDialog.tsx

@@ -5,7 +5,7 @@
  */
 
 import type React from 'react';
-import { useCallback } from 'react';
+import { useCallback, useState } from 'react';
 import { Box, Text } from 'ink';
 import { theme } from '../semantic-colors.js';
 import { RadioButtonSelect } from '../components/shared/RadioButtonSelect.js';
@@ -17,13 +17,13 @@ import { SettingScope } from '../../config/settings.js';
 import {
   AuthType,
   clearCachedCredentialFile,
-  debugLogger,
   type Config,
 } from '@google/gemini-cli-core';
 import { useKeypress } from '../hooks/useKeypress.js';
 import { AuthState } from '../types.js';
 import { runExitCleanup } from '../../utils/cleanup.js';
 import { validateAuthMethodWithSettings } from './useAuth.js';
+import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
 
 interface AuthDialogProps {
   config: Config;
@@ -40,6 +40,7 @@ export function AuthDialog({
   authError,
   onAuthError,
 }: AuthDialogProps): React.JSX.Element {
+  const [exiting, setExiting] = useState(false);
   let items = [
     {
       label: 'Login with Google',
@@ -111,6 +112,9 @@ export function AuthDialog({
 
   const onSelect = useCallback(
     async (authType: AuthType | undefined, scope: LoadableSettingScope) => {
+      if (exiting) {
+        return;
+      }
       if (authType) {
         await clearCachedCredentialFile();
 
@@ -119,15 +123,12 @@ export function AuthDialog({
           authType === AuthType.LOGIN_WITH_GOOGLE &&
           config.isBrowserLaunchSuppressed()
         ) {
-          runExitCleanup();
-          debugLogger.log(
-            `
-----------------------------------------------------------------
-Logging in with Google... Please restart Gemini CLI to continue.
-----------------------------------------------------------------
-            `,
-          );
-          process.exit(0);
+          setExiting(true);
+          setTimeout(async () => {
+            await runExitCleanup();
+            process.exit(RELAUNCH_EXIT_CODE);
+          }, 100);
+          return;
         }
       }
       if (authType === AuthType.USE_GEMINI) {
@@ -136,7 +137,7 @@ Logging in with Google... Please restart Gemini CLI to continue.
       }
       setAuthState(AuthState.Unauthenticated);
     },
-    [settings, config, setAuthState],
+    [settings, config, setAuthState, exiting],
   );
 
   const handleAuthSelect = (authMethod: AuthType) => {
@@ -169,6 +170,23 @@ Logging in with Google... Please restart Gemini CLI to continue.
     { isActive: true },
   );
 
+  if (exiting) {
+    return (
+      <Box
+        borderStyle="round"
+        borderColor={theme.border.focused}
+        flexDirection="row"
+        padding={1}
+        width="100%"
+        alignItems="flex-start"
+      >
+        <Text color={theme.text.primary}>
+          Logging in with Google... Restarting Gemini CLI to continue.
+        </Text>
+      </Box>
+    );
+  }
+
   return (
     <Box
       borderStyle="round"

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

@@ -18,6 +18,8 @@ import { ApiAuthDialog } from '../auth/ApiAuthDialog.js';
 import { EditorSettingsDialog } from './EditorSettingsDialog.js';
 import { PrivacyNotice } from '../privacy/PrivacyNotice.js';
 import { ProQuotaDialog } from './ProQuotaDialog.js';
+import { runExitCleanup } from '../../utils/cleanup.js';
+import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js';
 import { PermissionsModifyTrustDialog } from './PermissionsModifyTrustDialog.js';
 import { ModelDialog } from './ModelDialog.js';
 import { theme } from '../semantic-colors.js';
@@ -137,7 +139,10 @@ export const DialogManager = ({
         <SettingsDialog
           settings={settings}
           onSelect={() => uiActions.closeSettingsDialog()}
-          onRestartRequest={() => process.exit(0)}
+          onRestartRequest={async () => {
+            await runExitCleanup();
+            process.exit(RELAUNCH_EXIT_CODE);
+          }}
           availableTerminalHeight={terminalHeight - staticExtraHeight}
           config={config}
         />

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

@@ -97,15 +97,19 @@ export async function detectAndEnableKittyProtocol(): Promise<void> {
   });
 }
 
+import {
+  enableKittyKeyboardProtocol,
+  disableKittyKeyboardProtocol,
+} from '@google/gemini-cli-core';
+
 export function isKittyProtocolEnabled(): boolean {
   return kittyEnabled;
 }
 
 function disableAllProtocols() {
   try {
     if (kittyEnabled) {
-      // use writeSync to avoid race conditions
-      fs.writeSync(process.stdout.fd, '\x1b[<u');
+      disableKittyKeyboardProtocol();
       kittyEnabled = false;
     }
   } catch {
@@ -120,8 +124,7 @@ function disableAllProtocols() {
 export function enableSupportedProtocol(): void {
   try {
     if (kittySupported) {
-      // use writeSync to avoid race conditions
-      fs.writeSync(process.stdout.fd, '\x1b[>1u');
+      enableKittyKeyboardProtocol();
       kittyEnabled = true;
     }
   } catch {