fix(auth): Remove sharp edges from headless auth (#3985)

packages/cli/src/gemini.tsx

@@ -25,7 +25,7 @@ import { getStartupWarnings } from './utils/startupWarnings.js';
 import { getUserStartupWarnings } from './utils/userStartupWarnings.js';
 import { runNonInteractive } from './nonInteractiveCli.js';
 import { loadExtensions, Extension } from './config/extension.js';
-import { cleanupCheckpoints } from './utils/cleanup.js';
+import { cleanupCheckpoints, registerCleanup } from './utils/cleanup.js';
 import { getCliVersion } from './utils/version.js';
 import {
   ApprovalMode,
@@ -202,7 +202,7 @@ export async function main() {
   if (shouldBeInteractive) {
     const version = await getCliVersion();
     setWindowTitle(basename(workspaceRoot), settings);
-    render(
+    const instance = render(
       <React.StrictMode>
         <AppWrapper
           config={config}
@@ -213,6 +213,8 @@ export async function main() {
       </React.StrictMode>,
       { exitOnCtrlC: false },
     );
+
+    registerCleanup(() => instance.unmount());
     return;
   }
   // If not a TTY, read from stdin

packages/cli/src/ui/App.tsx

@@ -41,7 +41,8 @@ import { Help } from './components/Help.js';
 import { loadHierarchicalGeminiMemory } from '../config/config.js';
 import { LoadedSettings } from '../config/settings.js';
 import { Tips } from './components/Tips.js';
-import { useConsolePatcher } from './components/ConsolePatcher.js';
+import { ConsolePatcher } from './utils/ConsolePatcher.js';
+import { registerCleanup } from '../utils/cleanup.js';
 import { DetailedMessagesDisplay } from './components/DetailedMessagesDisplay.js';
 import { HistoryItemDisplay } from './components/HistoryItemDisplay.js';
 import { ContextSummaryDisplay } from './components/ContextSummaryDisplay.js';
@@ -111,6 +112,16 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     handleNewMessage,
     clearConsoleMessages: clearConsoleMessagesState,
   } = useConsoleMessages();
+
+  useEffect(() => {
+    const consolePatcher = new ConsolePatcher({
+      onNewMessage: handleNewMessage,
+      debugMode: config.getDebugMode(),
+    });
+    consolePatcher.patch();
+    registerCleanup(consolePatcher.cleanup);
+  }, [handleNewMessage, config]);
+
   const { stats: sessionStats } = useSessionStats();
   const [staticNeedsRefresh, setStaticNeedsRefresh] = useState(false);
   const [staticKey, setStaticKey] = useState(0);
@@ -470,11 +481,6 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
     }
   });
 
-  useConsolePatcher({
-    onNewMessage: handleNewMessage,
-    debugMode: config.getDebugMode(),
-  });
-
   useEffect(() => {
     if (config) {
       setGeminiMdFileCount(config.getGeminiMdFileCount());

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

@@ -1,60 +0,0 @@
-/**
- * @license
- * Copyright 2025 Google LLC
- * SPDX-License-Identifier: Apache-2.0
- */
-
-import { useEffect } from 'react';
-import util from 'util';
-import { ConsoleMessageItem } from '../types.js';
-
-interface UseConsolePatcherParams {
-  onNewMessage: (message: Omit<ConsoleMessageItem, 'id'>) => void;
-  debugMode: boolean;
-}
-
-export const useConsolePatcher = ({
-  onNewMessage,
-  debugMode,
-}: UseConsolePatcherParams): void => {
-  useEffect(() => {
-    const originalConsoleLog = console.log;
-    const originalConsoleWarn = console.warn;
-    const originalConsoleError = console.error;
-    const originalConsoleDebug = console.debug;
-
-    const formatArgs = (args: unknown[]): string => util.format(...args);
-
-    const patchConsoleMethod =
-      (
-        type: 'log' | 'warn' | 'error' | 'debug',
-        originalMethod: (...args: unknown[]) => void,
-      ) =>
-      (...args: unknown[]) => {
-        if (debugMode) {
-          originalMethod.apply(console, args);
-        }
-
-        // Then, if it's not a debug message or debugMode is on, pass to onNewMessage
-        if (type !== 'debug' || debugMode) {
-          onNewMessage({
-            type,
-            content: formatArgs(args),
-            count: 1,
-          });
-        }
-      };
-
-    console.log = patchConsoleMethod('log', originalConsoleLog);
-    console.warn = patchConsoleMethod('warn', originalConsoleWarn);
-    console.error = patchConsoleMethod('error', originalConsoleError);
-    console.debug = patchConsoleMethod('debug', originalConsoleDebug);
-
-    return () => {
-      console.log = originalConsoleLog;
-      console.warn = originalConsoleWarn;
-      console.error = originalConsoleError;
-      console.debug = originalConsoleDebug;
-    };
-  }, [onNewMessage, debugMode]);
-};

packages/cli/src/ui/hooks/useAuthCommand.ts

@@ -12,6 +12,7 @@ import {
   clearCachedCredentialFile,
   getErrorMessage,
 } from '@google/gemini-cli-core';
+import { runExitCleanup } from '../../utils/cleanup.js';
 
 export const useAuthCommand = (
   settings: LoadedSettings,
@@ -55,11 +56,22 @@ export const useAuthCommand = (
       if (authType) {
         await clearCachedCredentialFile();
         settings.setValue(scope, 'selectedAuthType', authType);
+        if (authType === AuthType.LOGIN_WITH_GOOGLE && config.getNoBrowser()) {
+          runExitCleanup();
+          console.log(
+            `
+----------------------------------------------------------------
+Logging in with Google... Please restart Gemini CLI to continue.
+----------------------------------------------------------------
+            `,
+          );
+          process.exit(0);
+        }
       }
       setIsAuthDialogOpen(false);
       setAuthError(null);
     },
-    [settings, setAuthError],
+    [settings, setAuthError, config],
   );
 
   const cancelAuthentication = useCallback(() => {

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

@@ -0,0 +1,61 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import util from 'util';
+import { ConsoleMessageItem } from '../types.js';
+
+interface ConsolePatcherParams {
+  onNewMessage: (message: Omit<ConsoleMessageItem, 'id'>) => void;
+  debugMode: boolean;
+}
+
+export class ConsolePatcher {
+  private originalConsoleLog = console.log;
+  private originalConsoleWarn = console.warn;
+  private originalConsoleError = console.error;
+  private originalConsoleDebug = console.debug;
+
+  private params: ConsolePatcherParams;
+
+  constructor(params: ConsolePatcherParams) {
+    this.params = params;
+  }
+
+  patch() {
+    console.log = this.patchConsoleMethod('log', this.originalConsoleLog);
+    console.warn = this.patchConsoleMethod('warn', this.originalConsoleWarn);
+    console.error = this.patchConsoleMethod('error', this.originalConsoleError);
+    console.debug = this.patchConsoleMethod('debug', this.originalConsoleDebug);
+  }
+
+  cleanup = () => {
+    console.log = this.originalConsoleLog;
+    console.warn = this.originalConsoleWarn;
+    console.error = this.originalConsoleError;
+    console.debug = this.originalConsoleDebug;
+  };
+
+  private formatArgs = (args: unknown[]): string => util.format(...args);
+
+  private patchConsoleMethod =
+    (
+      type: 'log' | 'warn' | 'error' | 'debug',
+      originalMethod: (...args: unknown[]) => void,
+    ) =>
+    (...args: unknown[]) => {
+      if (this.params.debugMode) {
+        originalMethod.apply(console, args);
+      }
+
+      if (type !== 'debug' || this.params.debugMode) {
+        this.params.onNewMessage({
+          type,
+          content: this.formatArgs(args),
+          count: 1,
+        });
+      }
+    };
+}

packages/cli/src/utils/cleanup.ts

@@ -8,6 +8,23 @@ import { promises as fs } from 'fs';
 import { join } from 'path';
 import { getProjectTempDir } from '@google/gemini-cli-core';
 
+const cleanupFunctions: Array<() => void> = [];
+
+export function registerCleanup(fn: () => void) {
+  cleanupFunctions.push(fn);
+}
+
+export function runExitCleanup() {
+  for (const fn of cleanupFunctions) {
+    try {
+      fn();
+    } catch (_) {
+      // Ignore errors during cleanup.
+    }
+  }
+  cleanupFunctions.length = 0; // Clear the array
+}
+
 export async function cleanupCheckpoints() {
   const tempDir = getProjectTempDir(process.cwd());
   const checkpointsDir = join(tempDir, 'checkpoints');

packages/core/src/code_assist/oauth2.test.ts

@@ -212,9 +212,7 @@ describe('oauth2', () => {
     };
     (readline.createInterface as Mock).mockReturnValue(mockReadline);
 
-    const consoleErrorSpy = vi
-      .spyOn(console, 'error')
-      .mockImplementation(() => {});
+    const consoleLogSpy = vi.spyOn(console, 'log').mockImplementation(() => {});
 
     const client = await getOauthClient(
       AuthType.LOGIN_WITH_GOOGLE,
@@ -226,7 +224,7 @@ describe('oauth2', () => {
     // Verify the auth flow
     expect(mockGenerateCodeVerifierAsync).toHaveBeenCalled();
     expect(mockGenerateAuthUrl).toHaveBeenCalled();
-    expect(consoleErrorSpy).toHaveBeenCalledWith(
+    expect(consoleLogSpy).toHaveBeenCalledWith(
       expect.stringContaining(mockAuthUrl),
     );
     expect(mockReadline.question).toHaveBeenCalledWith(
@@ -240,7 +238,7 @@ describe('oauth2', () => {
     });
     expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
 
-    consoleErrorSpy.mockRestore();
+    consoleLogSpy.mockRestore();
   });
 
   describe('in Cloud Shell', () => {

packages/core/src/code_assist/oauth2.ts

@@ -163,38 +163,35 @@ async function authWithUserCode(client: OAuth2Client): Promise<boolean> {
     code_challenge: codeVerifier.codeChallenge,
     state,
   });
-  console.error('Please visit the following URL to authorize the application:');
-  console.error('');
-  console.error(authUrl);
-  console.error('');
+  console.log('Please visit the following URL to authorize the application:');
+  console.log('');
+  console.log(authUrl);
+  console.log('');
 
   const code = await new Promise<string>((resolve) => {
     const rl = readline.createInterface({
       input: process.stdin,
       output: process.stdout,
     });
-    rl.question('Enter the authorization code: ', (answer) => {
+    rl.question('Enter the authorization code: ', (code) => {
       rl.close();
-      resolve(answer.trim());
+      resolve(code.trim());
     });
   });
 
   if (!code) {
     console.error('Authorization code is required.');
     return false;
-  } else {
-    console.error(`Received authorization code: "${code}"`);
   }
 
   try {
-    const response = await client.getToken({
+    const { tokens } = await client.getToken({
       code,
       codeVerifier: codeVerifier.codeVerifier,
       redirect_uri: redirectUri,
     });
-    client.setCredentials(response.tokens);
+    client.setCredentials(tokens);
   } catch (_error) {
-    // Consider logging the error.
     return false;
   }
   return true;