From 63214428658da3608180fc07b811e7fd2d143e63 Mon Sep 17 00:00:00 2001 From: Gal Zahavi <38544478+galz10@users.noreply.github.com> Date: Fri, 25 Jul 2025 10:19:38 -0700 Subject: [PATCH] feat(auth): Enhance non-interactive gcp auth (#4811) --- packages/cli/src/utils/sandbox.ts | 8 ++ .../src/validateNonInterActiveAuth.test.ts | 30 ++++ .../cli/src/validateNonInterActiveAuth.ts | 23 +-- packages/core/src/code_assist/oauth2.test.ts | 136 ++++++++++++++++++ packages/core/src/code_assist/oauth2.ts | 11 ++ 5 files changed, 200 insertions(+), 8 deletions(-) diff --git a/packages/cli/src/utils/sandbox.ts b/packages/cli/src/utils/sandbox.ts index 1c2d7d08..84fdc8f7 100644 --- a/packages/cli/src/utils/sandbox.ts +++ b/packages/cli/src/utils/sandbox.ts @@ -525,6 +525,14 @@ export async function start_sandbox( ); } + // copy GOOGLE_GENAI_USE_GCA + if (process.env.GOOGLE_GENAI_USE_GCA) { + args.push( + '--env', + `GOOGLE_GENAI_USE_GCA=${process.env.GOOGLE_GENAI_USE_GCA}`, + ); + } + // copy GOOGLE_CLOUD_PROJECT if (process.env.GOOGLE_CLOUD_PROJECT) { args.push( diff --git a/packages/cli/src/validateNonInterActiveAuth.test.ts b/packages/cli/src/validateNonInterActiveAuth.test.ts index 9238bbe4..184a70e0 100644 --- a/packages/cli/src/validateNonInterActiveAuth.test.ts +++ b/packages/cli/src/validateNonInterActiveAuth.test.ts @@ -14,6 +14,7 @@ import { AuthType } from '@google/gemini-cli-core'; describe('validateNonInterActiveAuth', () => { let originalEnvGeminiApiKey: string | undefined; let originalEnvVertexAi: string | undefined; + let originalEnvGcp: string | undefined; let consoleErrorSpy: ReturnType; let processExitSpy: ReturnType; let refreshAuthMock: jest.MockedFunction< @@ -23,8 +24,10 @@ describe('validateNonInterActiveAuth', () => { beforeEach(() => { originalEnvGeminiApiKey = process.env.GEMINI_API_KEY; originalEnvVertexAi = process.env.GOOGLE_GENAI_USE_VERTEXAI; + originalEnvGcp = process.env.GOOGLE_GENAI_USE_GCA; delete process.env.GEMINI_API_KEY; delete process.env.GOOGLE_GENAI_USE_VERTEXAI; + delete process.env.GOOGLE_GENAI_USE_GCA; consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}); processExitSpy = vi.spyOn(process, 'exit').mockImplementation((code) => { throw new Error(`process.exit(${code}) called`); @@ -43,6 +46,11 @@ describe('validateNonInterActiveAuth', () => { } else { delete process.env.GOOGLE_GENAI_USE_VERTEXAI; } + if (originalEnvGcp !== undefined) { + process.env.GOOGLE_GENAI_USE_GCA = originalEnvGcp; + } else { + delete process.env.GOOGLE_GENAI_USE_GCA; + } vi.restoreAllMocks(); }); @@ -62,6 +70,15 @@ describe('validateNonInterActiveAuth', () => { expect(processExitSpy).toHaveBeenCalledWith(1); }); + it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set', async () => { + process.env.GOOGLE_GENAI_USE_GCA = 'true'; + const nonInteractiveConfig: NonInteractiveConfig = { + refreshAuth: refreshAuthMock, + }; + await validateNonInteractiveAuth(undefined, nonInteractiveConfig); + expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.LOGIN_WITH_GOOGLE); + }); + it('uses USE_GEMINI if GEMINI_API_KEY is set', async () => { process.env.GEMINI_API_KEY = 'fake-key'; const nonInteractiveConfig: NonInteractiveConfig = { @@ -92,6 +109,19 @@ describe('validateNonInterActiveAuth', () => { expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_VERTEX_AI); }); + it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set, even with other env vars', async () => { + process.env.GOOGLE_GENAI_USE_GCA = 'true'; + process.env.GEMINI_API_KEY = 'fake-key'; + process.env.GOOGLE_GENAI_USE_VERTEXAI = 'true'; + process.env.GOOGLE_CLOUD_PROJECT = 'test-project'; + process.env.GOOGLE_CLOUD_LOCATION = 'us-central1'; + const nonInteractiveConfig: NonInteractiveConfig = { + refreshAuth: refreshAuthMock, + }; + await validateNonInteractiveAuth(undefined, nonInteractiveConfig); + expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.LOGIN_WITH_GOOGLE); + }); + it('uses USE_VERTEX_AI if both GEMINI_API_KEY and GOOGLE_GENAI_USE_VERTEXAI are set', async () => { process.env.GEMINI_API_KEY = 'fake-key'; process.env.GOOGLE_GENAI_USE_VERTEXAI = 'true'; diff --git a/packages/cli/src/validateNonInterActiveAuth.ts b/packages/cli/src/validateNonInterActiveAuth.ts index 87a7f4ff..a85b7370 100644 --- a/packages/cli/src/validateNonInterActiveAuth.ts +++ b/packages/cli/src/validateNonInterActiveAuth.ts @@ -8,21 +8,28 @@ import { AuthType, Config } from '@google/gemini-cli-core'; import { USER_SETTINGS_PATH } from './config/settings.js'; import { validateAuthMethod } from './config/auth.js'; +function getAuthTypeFromEnv(): AuthType | undefined { + if (process.env.GOOGLE_GENAI_USE_GCA === 'true') { + return AuthType.LOGIN_WITH_GOOGLE; + } + if (process.env.GOOGLE_GENAI_USE_VERTEXAI === 'true') { + return AuthType.USE_VERTEX_AI; + } + if (process.env.GEMINI_API_KEY) { + return AuthType.USE_GEMINI; + } + return undefined; +} + export async function validateNonInteractiveAuth( configuredAuthType: AuthType | undefined, nonInteractiveConfig: Config, ) { - const effectiveAuthType = - configuredAuthType || - (process.env.GOOGLE_GENAI_USE_VERTEXAI === 'true' - ? AuthType.USE_VERTEX_AI - : process.env.GEMINI_API_KEY - ? AuthType.USE_GEMINI - : undefined); + const effectiveAuthType = configuredAuthType || getAuthTypeFromEnv(); if (!effectiveAuthType) { console.error( - `Please set an Auth method in your ${USER_SETTINGS_PATH} or specify either the GEMINI_API_KEY or GOOGLE_GENAI_USE_VERTEXAI environment variables before running`, + `Please set an Auth method in your ${USER_SETTINGS_PATH} or specify one of the following environment variables before running: GEMINI_API_KEY, GOOGLE_GENAI_USE_VERTEXAI, GOOGLE_GENAI_USE_GCA`, ); process.exit(1); } diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index ee63c02b..0dc77867 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/packages/core/src/code_assist/oauth2.test.ts @@ -57,6 +57,8 @@ describe('oauth2', () => { fs.rmSync(tempHomeDir, { recursive: true, force: true }); vi.clearAllMocks(); delete process.env.CLOUD_SHELL; + delete process.env.GOOGLE_GENAI_USE_GCA; + delete process.env.GOOGLE_CLOUD_ACCESS_TOKEN; }); it('should perform a web login', async () => { @@ -332,4 +334,138 @@ describe('oauth2', () => { ); }); }); + + describe('with GCP environment variables', () => { + it('should use GOOGLE_CLOUD_ACCESS_TOKEN when GOOGLE_GENAI_USE_GCA is true', async () => { + process.env.GOOGLE_GENAI_USE_GCA = 'true'; + process.env.GOOGLE_CLOUD_ACCESS_TOKEN = 'gcp-access-token'; + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'gcp-access-token' }); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Mock the UserInfo API response for fetchAndCacheUserInfo + (global.fetch as Mock).mockResolvedValue({ + ok: true, + json: vi + .fn() + .mockResolvedValue({ email: 'test-gcp-account@gmail.com' }), + } as unknown as Response); + + const writeFileSpy = vi + .spyOn(fs.promises, 'writeFile') + .mockResolvedValue(undefined); + + const client = await getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + + expect(client).toBe(mockOAuth2Client); + expect(mockSetCredentials).toHaveBeenCalledWith({ + access_token: 'gcp-access-token', + }); + + // Verify fetchAndCacheUserInfo was effectively called + expect(mockGetAccessToken).toHaveBeenCalled(); + expect(global.fetch).toHaveBeenCalledWith( + 'https://www.googleapis.com/oauth2/v2/userinfo', + { + headers: { + Authorization: 'Bearer gcp-access-token', + }, + }, + ); + + // Verify Google Account was cached + const googleAccountPath = path.join( + tempHomeDir, + '.gemini', + 'google_accounts.json', + ); + expect(writeFileSpy).toHaveBeenCalledWith( + googleAccountPath, + JSON.stringify( + { + active: 'test-gcp-account@gmail.com', + old: [], + }, + null, + 2, + ), + 'utf-8', + ); + }); + + it('should not use GCP token if GOOGLE_CLOUD_ACCESS_TOKEN is not set', async () => { + process.env.GOOGLE_GENAI_USE_GCA = 'true'; + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'cached-access-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Make it fall through to cached credentials path + const cachedCreds = { refresh_token: 'cached-token' }; + vi.spyOn(fs.promises, 'readFile').mockResolvedValue( + JSON.stringify(cachedCreds), + ); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // It should be called with the cached credentials, not the GCP access token. + expect(mockSetCredentials).toHaveBeenCalledTimes(1); + expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); + }); + + it('should not use GCP token if GOOGLE_GENAI_USE_GCA is not set', async () => { + process.env.GOOGLE_CLOUD_ACCESS_TOKEN = 'gcp-access-token'; + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'cached-access-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Make it fall through to cached credentials path + const cachedCreds = { refresh_token: 'cached-token' }; + vi.spyOn(fs.promises, 'readFile').mockResolvedValue( + JSON.stringify(cachedCreds), + ); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // It should be called with the cached credentials, not the GCP access token. + expect(mockSetCredentials).toHaveBeenCalledTimes(1); + expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); + }); + }); }); diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 5958625a..f1046416 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -78,6 +78,17 @@ export async function getOauthClient( }, }); + if ( + process.env.GOOGLE_GENAI_USE_GCA && + process.env.GOOGLE_CLOUD_ACCESS_TOKEN + ) { + client.setCredentials({ + access_token: process.env.GOOGLE_CLOUD_ACCESS_TOKEN, + }); + await fetchAndCacheUserInfo(client); + return client; + } + client.on('tokens', async (tokens: Credentials) => { await cacheCredentials(tokens); });