diff --git a/docs/cli/authentication.md b/docs/cli/authentication.md index a3d3c0f3..3dddb897 100644 --- a/docs/cli/authentication.md +++ b/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 diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 6e724ab1..31931391 100644 --- a/docs/cli/configuration.md +++ b/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. diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts index 439bd36b..4561f5d6 100644 --- a/packages/cli/src/config/auth.ts +++ b/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; } diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts index 93f042aa..c2d03167 100644 --- a/packages/cli/src/config/settings.ts +++ b/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 }); } diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 39d3bbe3..7e86a8ca 100644 --- a/packages/cli/src/gemini.tsx +++ b/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()); diff --git a/packages/cli/src/ui/components/AuthDialog.tsx b/packages/cli/src/ui/components/AuthDialog.tsx index d0601dc1..cea6727f 100644 --- a/packages/cli/src/ui/components/AuthDialog.tsx +++ b/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 }, ]; diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 80d95ca9..5f3f843e 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/packages/core/src/code_assist/codeAssist.ts @@ -14,8 +14,11 @@ export async function createCodeAssistContentGenerator( authType: AuthType, sessionId?: string, ): Promise { - 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); } diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index ae5d80d6..76d43726 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/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(); @@ -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', + ); + }); + }); }); diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 043572d7..93d0e28b 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/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; } -export async function getOauthClient(): Promise { +export async function getOauthClient( + authType: AuthType, +): Promise { 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 { 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( diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 1b22333a..ce3c11a9 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/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,