Enable Gemini CLI to reuse user's auth in Cloud Shell (#3070)

packages/core/src/code_assist/codeAssist.ts

@@ -14,8 +14,11 @@ export async function createCodeAssistContentGenerator(
   authType: AuthType,
   sessionId?: string,
 ): Promise<ContentGenerator> {
-  if (authType === AuthType.LOGIN_WITH_GOOGLE) {
-    const authClient = await getOauthClient();
+  if (
+    authType === AuthType.LOGIN_WITH_GOOGLE ||
+    authType === AuthType.CLOUD_SHELL
+  ) {
+    const authClient = await getOauthClient(authType);
     const projectId = await setupUser(authClient);
     return new CodeAssistServer(authClient, projectId, httpOptions, sessionId);
   }

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

@@ -4,15 +4,16 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
 import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js';
-import { OAuth2Client } from 'google-auth-library';
+import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
 import * as path from 'path';
 import http from 'http';
 import open from 'open';
 import crypto from 'crypto';
 import * as os from 'os';
+import { AuthType } from '../core/contentGenerator.js';
 
 vi.mock('os', async (importOriginal) => {
   const os = await importOriginal<typeof import('os')>();
@@ -37,10 +38,12 @@ describe('oauth2', () => {
     tempHomeDir = fs.mkdtempSync(
       path.join(os.tmpdir(), 'gemini-cli-test-home-'),
     );
-    vi.mocked(os.homedir).mockReturnValue(tempHomeDir);
+    (os.homedir as Mock).mockReturnValue(tempHomeDir);
   });
   afterEach(() => {
     fs.rmSync(tempHomeDir, { recursive: true, force: true });
+    vi.clearAllMocks();
+    delete process.env.CLOUD_SHELL;
   });
 
   it('should perform a web login', async () => {
@@ -85,13 +88,15 @@ describe('oauth2', () => {
       credentials: mockTokens,
       on: vi.fn(),
     } as unknown as OAuth2Client;
-    vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
+    (OAuth2Client as unknown as Mock).mockImplementation(
+      () => mockOAuth2Client,
+    );
 
     vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
-    vi.mocked(open).mockImplementation(async () => ({}) as never);
+    (open as Mock).mockImplementation(async () => ({}) as never);
 
     // Mock the UserInfo API response
-    vi.mocked(global.fetch).mockResolvedValue({
+    (global.fetch as Mock).mockResolvedValue({
       ok: true,
       json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }),
     } as unknown as Response);
@@ -123,7 +128,7 @@ describe('oauth2', () => {
       on: vi.fn(),
       address: () => ({ port: capturedPort }),
     };
-    vi.mocked(http.createServer).mockImplementation((cb) => {
+    (http.createServer as Mock).mockImplementation((cb) => {
       requestCallback = cb as http.RequestListener<
         typeof http.IncomingMessage,
         typeof http.ServerResponse
@@ -131,7 +136,7 @@ describe('oauth2', () => {
       return mockHttpServer as unknown as http.Server;
     });
 
-    const clientPromise = getOauthClient();
+    const clientPromise = getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
 
     // wait for server to start listening.
     await serverListeningPromise;
@@ -169,4 +174,87 @@ describe('oauth2', () => {
     // Verify the getCachedGoogleAccountId function works
     expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123');
   });
+
+  describe('in Cloud Shell', () => {
+    const mockGetAccessToken = vi.fn();
+    let mockComputeClient: Compute;
+
+    beforeEach(() => {
+      vi.spyOn(os, 'homedir').mockReturnValue('/user/home');
+      vi.spyOn(fs.promises, 'mkdir').mockResolvedValue(undefined);
+      vi.spyOn(fs.promises, 'writeFile').mockResolvedValue(undefined);
+      vi.spyOn(fs.promises, 'readFile').mockRejectedValue(
+        new Error('File not found'),
+      ); // Default to no cached creds
+
+      mockGetAccessToken.mockResolvedValue({ token: 'test-access-token' });
+      mockComputeClient = {
+        credentials: { refresh_token: 'test-refresh-token' },
+        getAccessToken: mockGetAccessToken,
+      } as unknown as Compute;
+
+      (Compute as unknown as Mock).mockImplementation(() => mockComputeClient);
+    });
+
+    it('should attempt to load cached credentials first', async () => {
+      const cachedCreds = { refresh_token: 'cached-token' };
+      vi.spyOn(fs.promises, 'readFile').mockResolvedValue(
+        JSON.stringify(cachedCreds),
+      );
+
+      const mockClient = {
+        setCredentials: vi.fn(),
+        getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }),
+        getTokenInfo: vi.fn().mockResolvedValue({}),
+        on: vi.fn(),
+      };
+
+      // To mock the new OAuth2Client() inside the function
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockClient as unknown as OAuth2Client,
+      );
+
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
+
+      expect(fs.promises.readFile).toHaveBeenCalledWith(
+        '/user/home/.gemini/oauth_creds.json',
+        'utf-8',
+      );
+      expect(mockClient.setCredentials).toHaveBeenCalledWith(cachedCreds);
+      expect(mockClient.getAccessToken).toHaveBeenCalled();
+      expect(mockClient.getTokenInfo).toHaveBeenCalled();
+      expect(Compute).not.toHaveBeenCalled(); // Should not fetch new client if cache is valid
+    });
+
+    it('should use Compute to get a client if no cached credentials exist', async () => {
+      await getOauthClient(AuthType.CLOUD_SHELL);
+
+      expect(Compute).toHaveBeenCalledWith({});
+      expect(mockGetAccessToken).toHaveBeenCalled();
+    });
+
+    it('should not cache the credentials after fetching them via ADC', async () => {
+      const newCredentials = { refresh_token: 'new-adc-token' };
+      mockComputeClient.credentials = newCredentials;
+      mockGetAccessToken.mockResolvedValue({ token: 'new-adc-token' });
+
+      await getOauthClient(AuthType.CLOUD_SHELL);
+
+      expect(fs.promises.writeFile).not.toHaveBeenCalled();
+    });
+
+    it('should return the Compute client on successful ADC authentication', async () => {
+      const client = await getOauthClient(AuthType.CLOUD_SHELL);
+      expect(client).toBe(mockComputeClient);
+    });
+
+    it('should throw an error if ADC fails', async () => {
+      const testError = new Error('ADC Failed');
+      mockGetAccessToken.mockRejectedValue(testError);
+
+      await expect(getOauthClient(AuthType.CLOUD_SHELL)).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',
+      );
+    });
+  });
 });

packages/core/src/code_assist/oauth2.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { OAuth2Client, Credentials } from 'google-auth-library';
+import { OAuth2Client, Credentials, Compute } from 'google-auth-library';
 import * as http from 'http';
 import url from 'url';
 import crypto from 'crypto';
@@ -13,6 +13,8 @@ import open from 'open';
 import path from 'node:path';
 import { promises as fs, existsSync, readFileSync } from 'node:fs';
 import * as os from 'os';
+import { getErrorMessage } from '../utils/errors.js';
+import { AuthType } from '../core/contentGenerator.js';
 
 //  OAuth Client ID used to initiate OAuth2Client class.
 const OAUTH_CLIENT_ID =
@@ -53,15 +55,19 @@ export interface OauthWebLogin {
   loginCompletePromise: Promise<void>;
 }
 
-export async function getOauthClient(): Promise<OAuth2Client> {
+export async function getOauthClient(
+  authType: AuthType,
+): Promise<OAuth2Client> {
   const client = new OAuth2Client({
     clientId: OAUTH_CLIENT_ID,
     clientSecret: OAUTH_CLIENT_SECRET,
   });
+
   client.on('tokens', async (tokens: Credentials) => {
     await cacheCredentials(tokens);
   });
 
+  // If there are cached creds on disk, they always take precedence
   if (await loadCachedCredentials(client)) {
     // Found valid cached credentials.
     // Check if we need to retrieve Google Account ID
@@ -71,17 +77,39 @@ export async function getOauthClient(): Promise<OAuth2Client> {
         if (googleAccountId) {
           await cacheGoogleAccountId(googleAccountId);
         }
-      } catch (error) {
-        console.error(
-          'Failed to retrieve Google Account ID for existing credentials:',
-          error,
-        );
-        // Continue with existing auth flow
+      } catch {
+        // Non-fatal, continue with existing auth.
       }
     }
+    console.log('Loaded cached credentials.');
     return client;
   }
 
+  // In Google Cloud Shell, we can use Application Default Credentials (ADC)
+  // provided via its metadata server to authenticate non-interactively using
+  // the identity of the user logged into Cloud Shell.
+  if (authType === AuthType.CLOUD_SHELL) {
+    try {
+      console.log("Attempting to authenticate via Cloud Shell VM's ADC.");
+      const computeClient = new Compute({
+        // We can leave this empty, since the metadata server will provide
+        // the service account email.
+      });
+      await computeClient.getAccessToken();
+      console.log('Authentication successful.');
+
+      // Do not cache creds in this case; note that Compute client will handle its own refresh
+      return computeClient;
+    } catch (e) {
+      throw new Error(
+        `Could not authenticate using Cloud Shell credentials. Please select a different authentication method or ensure you are in a properly configured environment. Error: ${getErrorMessage(
+          e,
+        )}`,
+      );
+    }
+  }
+
+  // Otherwise, obtain creds using standard web flow
   const webLogin = await authWithWeb(client);
 
   console.log(