Add NO_BROWSER environment variable to trigger offline oauth flow (#3713)

3 files changed, 105 insertions, 19 deletions

diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index 5f3f843e..23dfe403 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/packages/core/src/code_assist/codeAssist.ts
@@ -8,17 +8,19 @@ import { AuthType, ContentGenerator } from '../core/contentGenerator.js';
 import { getOauthClient } from './oauth2.js';
 import { setupUser } from './setup.js';
 import { CodeAssistServer, HttpOptions } from './server.js';
+import { Config } from '../config/config.js';
 
 export async function createCodeAssistContentGenerator(
   httpOptions: HttpOptions,
   authType: AuthType,
+  config: Config,
   sessionId?: string,
 ): Promise<ContentGenerator> {
   if (
     authType === AuthType.LOGIN_WITH_GOOGLE ||
     authType === AuthType.CLOUD_SHELL
   ) {
-    const authClient = await getOauthClient(authType);
+    const authClient = await getOauthClient(authType, config);
     const projectId = await setupUser(authClient);
     return new CodeAssistServer(authClient, projectId, httpOptions, sessionId);
   }
diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index 76d43726..d8cd525b 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -14,6 +14,7 @@ import open from 'open';
 import crypto from 'crypto';
 import * as os from 'os';
 import { AuthType } from '../core/contentGenerator.js';
+import { Config } from '../config/config.js';
 
 vi.mock('os', async (importOriginal) => {
   const os = await importOriginal<typeof import('os')>();
@@ -28,6 +29,10 @@ vi.mock('http');
 vi.mock('open');
 vi.mock('crypto');
 
+const mockConfig = {
+  getNoBrowser: () => false,
+} as unknown as Config;
+
 // Mock fetch globally
 global.fetch = vi.fn();
 
@@ -136,7 +141,10 @@ describe('oauth2', () => {
       return mockHttpServer as unknown as http.Server;
     });
 
-    const clientPromise = getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
+    const clientPromise = getOauthClient(
+      AuthType.LOGIN_WITH_GOOGLE,
+      mockConfig,
+    );
 
     // wait for server to start listening.
     await serverListeningPromise;
@@ -214,7 +222,7 @@ describe('oauth2', () => {
         () => mockClient as unknown as OAuth2Client,
       );
 
-      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
 
       expect(fs.promises.readFile).toHaveBeenCalledWith(
         '/user/home/.gemini/oauth_creds.json',
@@ -227,7 +235,7 @@ describe('oauth2', () => {
     });
 
     it('should use Compute to get a client if no cached credentials exist', async () => {
-      await getOauthClient(AuthType.CLOUD_SHELL);
+      await getOauthClient(AuthType.CLOUD_SHELL, mockConfig);
 
       expect(Compute).toHaveBeenCalledWith({});
       expect(mockGetAccessToken).toHaveBeenCalled();
@@ -238,13 +246,13 @@ describe('oauth2', () => {
       mockComputeClient.credentials = newCredentials;
       mockGetAccessToken.mockResolvedValue({ token: 'new-adc-token' });
 
-      await getOauthClient(AuthType.CLOUD_SHELL);
+      await getOauthClient(AuthType.CLOUD_SHELL, mockConfig);
 
       expect(fs.promises.writeFile).not.toHaveBeenCalled();
     });
 
     it('should return the Compute client on successful ADC authentication', async () => {
-      const client = await getOauthClient(AuthType.CLOUD_SHELL);
+      const client = await getOauthClient(AuthType.CLOUD_SHELL, mockConfig);
       expect(client).toBe(mockComputeClient);
     });
 
@@ -252,7 +260,9 @@ describe('oauth2', () => {
       const testError = new Error('ADC Failed');
       mockGetAccessToken.mockRejectedValue(testError);
 
-      await expect(getOauthClient(AuthType.CLOUD_SHELL)).rejects.toThrow(
+      await expect(
+        getOauthClient(AuthType.CLOUD_SHELL, mockConfig),
+      ).rejects.toThrow(
         'Could not authenticate using Cloud Shell credentials. Please select a different authentication method or ensure you are in a properly configured environment. Error: ADC Failed',
       );
     });
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index 93d0e28b..2d3c04d0 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -4,7 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { OAuth2Client, Credentials, Compute } from 'google-auth-library';
+import {
+  OAuth2Client,
+  Credentials,
+  Compute,
+  CodeChallengeMethod,
+} from 'google-auth-library';
 import * as http from 'http';
 import url from 'url';
 import crypto from 'crypto';
@@ -13,8 +18,10 @@ import open from 'open';
 import path from 'node:path';
 import { promises as fs, existsSync, readFileSync } from 'node:fs';
 import * as os from 'os';
+import { Config } from '../config/config.js';
 import { getErrorMessage } from '../utils/errors.js';
 import { AuthType } from '../core/contentGenerator.js';
+import readline from 'node:readline';
 
 //  OAuth Client ID used to initiate OAuth2Client class.
 const OAUTH_CLIENT_ID =
@@ -57,6 +64,7 @@ export interface OauthWebLogin {
 
 export async function getOauthClient(
   authType: AuthType,
+  config: Config,
 ): Promise<OAuth2Client> {
   const client = new OAuth2Client({
     clientId: OAUTH_CLIENT_ID,
@@ -109,27 +117,93 @@ export async function getOauthClient(
     }
   }
 
-  // Otherwise, obtain creds using standard web flow
-  const webLogin = await authWithWeb(client);
+  if (config.getNoBrowser()) {
+    let success = false;
+    const maxRetries = 2;
+    for (let i = 0; !success && i < maxRetries; i++) {
+      success = await authWithUserCode(client);
+      if (!success) {
+        console.error(
+          '\nFailed to authenticate with user code.',
+          i === maxRetries - 1 ? '' : 'Retrying...\n',
+        );
+      }
+    }
+    if (!success) {
+      process.exit(1);
+    }
+  } else {
+    const webLogin = await authWithWeb(client);
 
-  console.log(
-    `\n\nCode Assist login required.\n` +
-      `Attempting to open authentication page in your browser.\n` +
-      `Otherwise navigate to:\n\n${webLogin.authUrl}\n\n`,
-  );
-  await open(webLogin.authUrl);
-  console.log('Waiting for authentication...');
+    // This does basically nothing, as it isn't show to the user.
+    console.log(
+      `\n\nCode Assist login required.\n` +
+        `Attempting to open authentication page in your browser.\n` +
+        `Otherwise navigate to:\n\n${webLogin.authUrl}\n\n`,
+    );
+    await open(webLogin.authUrl);
+    console.log('Waiting for authentication...');
 
-  await webLogin.loginCompletePromise;
+    await webLogin.loginCompletePromise;
+  }
 
   return client;
 }
 
+async function authWithUserCode(client: OAuth2Client): Promise<boolean> {
+  const redirectUri = 'https://sdk.cloud.google.com/authcode_cloudcode.html';
+  const codeVerifier = await client.generateCodeVerifierAsync();
+  const state = crypto.randomBytes(32).toString('hex');
+  const authUrl: string = client.generateAuthUrl({
+    redirect_uri: redirectUri,
+    access_type: 'offline',
+    scope: OAUTH_SCOPE,
+    code_challenge_method: CodeChallengeMethod.S256,
+    code_challenge: codeVerifier.codeChallenge,
+    state,
+  });
+  console.error('Please visit the following URL to authorize the application:');
+  console.error('');
+  console.error(authUrl);
+  console.error('');
+
+  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.close();
+      resolve(answer.trim());
+    });
+  });
+
+  if (!code) {
+    console.error('Authorization code is required.');
+    return false;
+  } else {
+    console.error(`Received authorization code: "${code}"`);
+  }
+
+  try {
+    const response = await client.getToken({
+      code,
+      codeVerifier: codeVerifier.codeVerifier,
+      redirect_uri: redirectUri,
+    });
+    client.setCredentials(response.tokens);
+  } catch (_error) {
+    // Consider logging the error.
+    return false;
+  }
+  return true;
+}
+
 async function authWithWeb(client: OAuth2Client): Promise<OauthWebLogin> {
   const port = await getAvailablePort();
   const redirectUri = `http://localhost:${port}/oauth2callback`;
   const state = crypto.randomBytes(32).toString('hex');
-  const authUrl: string = client.generateAuthUrl({
+  const authUrl = client.generateAuthUrl({
     redirect_uri: redirectUri,
     access_type: 'offline',
     scope: OAUTH_SCOPE,