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

docs/cli/authentication.md

@@ -74,6 +74,10 @@ The Gemini CLI requires you to authenticate with Google's AI services. On initia
           echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc
           source ~/.bashrc
           ```
+4.  **Cloud Shell:**
+    - This option is only available when running in a Google Cloud Shell environment.
+    - It automatically uses the credentials of the logged-in user in the Cloud Shell environment.
+    - This is the default authentication method when running in Cloud Shell and no other method is configured.
 
 ### Persisting Environment Variables with `.env` Files
 

docs/cli/configuration.md

@@ -248,6 +248,7 @@ The CLI automatically loads environment variables from an `.env` file. The loadi
   - Your Google Cloud Project ID.
   - Required for using Code Assist or Vertex AI.
   - If using Vertex AI, ensure you have the necessary permissions and set the `GOOGLE_GENAI_USE_VERTEXAI=true` environment variable.
+  - **Cloud Shell Note:** When running in a Cloud Shell environment, this variable defaults to a special project allocated for Cloud Shell users. If you have `GOOGLE_CLOUD_PROJECT` set in your global environment in Cloud Shell, it will be overridden by this default. To use a different project in Cloud Shell, you must define `GOOGLE_CLOUD_PROJECT` in a `.env` file.
   - Example: `export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"`.
 - **`GOOGLE_APPLICATION_CREDENTIALS`** (string):
   - **Description:** The path to your Google Application Credentials JSON file.

packages/cli/src/config/auth.ts

@@ -9,7 +9,7 @@ import { loadEnvironment } from './settings.js';
 
 export const validateAuthMethod = (authMethod: string): string | null => {
   loadEnvironment();
-  if (authMethod === AuthType.LOGIN_WITH_GOOGLE) {
+  if (authMethod === AuthType.LOGIN_WITH_GOOGLE || AuthType.CLOUD_SHELL) {
     return null;
   }
 

packages/cli/src/config/settings.ts

@@ -203,8 +203,35 @@ function findEnvFile(startDir: string): string | null {
   }
 }
 
+export function setUpCloudShellEnvironment(envFilePath: string | null): void {
+  // Special handling for GOOGLE_CLOUD_PROJECT in Cloud Shell:
+  // Because GOOGLE_CLOUD_PROJECT in Cloud Shell tracks the project
+  // set by the user using "gcloud config set project" we do not want to
+  // use its value. So, unless the user overrides GOOGLE_CLOUD_PROJECT in
+  // one of the .env files, we set the Cloud Shell-specific default here.
+  if (envFilePath && fs.existsSync(envFilePath)) {
+    const envFileContent = fs.readFileSync(envFilePath);
+    const parsedEnv = dotenv.parse(envFileContent);
+    if (parsedEnv.GOOGLE_CLOUD_PROJECT) {
+      // .env file takes precedence in Cloud Shell
+      process.env.GOOGLE_CLOUD_PROJECT = parsedEnv.GOOGLE_CLOUD_PROJECT;
+    } else {
+      // If not in .env, set to default and override global
+      process.env.GOOGLE_CLOUD_PROJECT = 'cloudshell-gca';
+    }
+  } else {
+    // If no .env file, set to default and override global
+    process.env.GOOGLE_CLOUD_PROJECT = 'cloudshell-gca';
+  }
+}
+
 export function loadEnvironment(): void {
   const envFilePath = findEnvFile(process.cwd());
+
+  if (process.env.CLOUD_SHELL === 'true') {
+    setUpCloudShellEnvironment(envFilePath);
+  }
+
   if (envFilePath) {
     dotenv.config({ path: envFilePath, quiet: true });
   }

packages/cli/src/gemini.tsx

@@ -103,14 +103,21 @@ export async function main() {
   const extensions = loadExtensions(workspaceRoot);
   const config = await loadCliConfig(settings.merged, extensions, sessionId);
 
-  // set default fallback to gemini api key
-  // this has to go after load cli because that's where the env is set
-  if (!settings.merged.selectedAuthType && process.env.GEMINI_API_KEY) {
-    settings.setValue(
-      SettingScope.User,
-      'selectedAuthType',
-      AuthType.USE_GEMINI,
-    );
+  // Set a default auth type if one isn't set for a couple of known cases.
+  if (!settings.merged.selectedAuthType) {
+    if (process.env.GEMINI_API_KEY) {
+      settings.setValue(
+        SettingScope.User,
+        'selectedAuthType',
+        AuthType.USE_GEMINI,
+      );
+    } else if (process.env.CLOUD_SHELL === 'true') {
+      settings.setValue(
+        SettingScope.User,
+        'selectedAuthType',
+        AuthType.CLOUD_SHELL,
+      );
+    }
   }
 
   setMaxSizedBoxDebugging(config.getDebugMode());

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

@@ -27,8 +27,22 @@ export function AuthDialog({
     initialErrorMessage || null,
   );
   const items = [
-    { label: 'Login with Google', value: AuthType.LOGIN_WITH_GOOGLE },
-    { label: 'Gemini API Key (AI Studio)', value: AuthType.USE_GEMINI },
+    {
+      label: 'Login with Google',
+      value: AuthType.LOGIN_WITH_GOOGLE,
+    },
+    ...(process.env.CLOUD_SHELL === 'true'
+      ? [
+          {
+            label: 'Use Cloud Shell user credentials',
+            value: AuthType.CLOUD_SHELL,
+          },
+        ]
+      : []),
+    {
+      label: 'Use Gemini API Key',
+      value: AuthType.USE_GEMINI,
+    },
     { label: 'Vertex AI', value: AuthType.USE_VERTEX_AI },
   ];
 

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(

packages/core/src/core/contentGenerator.ts

@@ -38,6 +38,7 @@ export enum AuthType {
   LOGIN_WITH_GOOGLE = 'oauth-personal',
   USE_GEMINI = 'gemini-api-key',
   USE_VERTEX_AI = 'vertex-ai',
+  CLOUD_SHELL = 'cloud-shell',
 }
 
 export type ContentGeneratorConfig = {
@@ -64,8 +65,11 @@ export async function createContentGeneratorConfig(
     authType,
   };
 
-  // if we are using google auth nothing else to validate for now
-  if (authType === AuthType.LOGIN_WITH_GOOGLE) {
+  // If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
+  if (
+    authType === AuthType.LOGIN_WITH_GOOGLE ||
+    authType === AuthType.CLOUD_SHELL
+  ) {
     return contentGeneratorConfig;
   }
 
@@ -108,7 +112,10 @@ export async function createContentGenerator(
       'User-Agent': `GeminiCLI/${version} (${process.platform}; ${process.arch})`,
     },
   };
-  if (config.authType === AuthType.LOGIN_WITH_GOOGLE) {
+  if (
+    config.authType === AuthType.LOGIN_WITH_GOOGLE ||
+    config.authType === AuthType.CLOUD_SHELL
+  ) {
     return createCodeAssistContentGenerator(
       httpOptions,
       config.authType,