From a255529c6b5c333c82b12336e85c3f8f5f4e73bc Mon Sep 17 00:00:00 2001 From: Spencer Date: Mon, 30 Mar 2026 08:26:15 -0400 Subject: [PATCH] fix(a2a-server): prioritize ADC before evaluating headless constraints for auth initialization (#23614) --- packages/a2a-server/src/config/config.test.ts | 188 +++++------------- packages/a2a-server/src/config/config.ts | 82 +++----- packages/core/src/code_assist/oauth2.ts | 5 +- 3 files changed, 80 insertions(+), 195 deletions(-) diff --git a/packages/a2a-server/src/config/config.test.ts b/packages/a2a-server/src/config/config.test.ts index 1c553d7539..f4d5fbd330 100644 --- a/packages/a2a-server/src/config/config.test.ts +++ b/packages/a2a-server/src/config/config.test.ts @@ -424,7 +424,22 @@ describe('loadConfig', () => { }); }); - describe('authentication fallback', () => { + describe('authentication logic', () => { + const setupConfigMock = (refreshAuthMock: ReturnType) => { + vi.mocked(Config).mockImplementation( + (params: unknown) => + ({ + ...(params as object), + initialize: vi.fn(), + waitForMcpInit: vi.fn(), + refreshAuth: refreshAuthMock, + getExperiments: vi.fn().mockReturnValue({ flags: {} }), + getRemoteAdminSettings: vi.fn(), + setRemoteAdminSettings: vi.fn(), + }) as unknown as Config, + ); + }; + beforeEach(() => { vi.stubEnv('USE_CCPA', 'true'); vi.stubEnv('GEMINI_API_KEY', ''); @@ -434,182 +449,77 @@ describe('loadConfig', () => { vi.unstubAllEnvs(); }); - it('should fall back to COMPUTE_ADC in Cloud Shell if LOGIN_WITH_GOOGLE fails', async () => { - vi.stubEnv('CLOUD_SHELL', 'true'); - vi.mocked(isHeadlessMode).mockReturnValue(false); - const refreshAuthMock = vi.fn().mockImplementation((authType) => { - if (authType === AuthType.LOGIN_WITH_GOOGLE) { - throw new FatalAuthenticationError('Non-interactive session'); - } - return Promise.resolve(); - }); - - // Update the mock implementation for this test - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); - - await loadConfig(mockSettings, mockExtensionLoader, taskId); - - expect(refreshAuthMock).toHaveBeenCalledWith( - AuthType.LOGIN_WITH_GOOGLE, - ); - expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); - }); - - it('should not fall back to COMPUTE_ADC if not in cloud environment', async () => { - vi.mocked(isHeadlessMode).mockReturnValue(false); - const refreshAuthMock = vi.fn().mockImplementation((authType) => { - if (authType === AuthType.LOGIN_WITH_GOOGLE) { - throw new FatalAuthenticationError('Non-interactive session'); - } - return Promise.resolve(); - }); - - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); - - await expect( - loadConfig(mockSettings, mockExtensionLoader, taskId), - ).rejects.toThrow('Non-interactive session'); - - expect(refreshAuthMock).toHaveBeenCalledWith( - AuthType.LOGIN_WITH_GOOGLE, - ); - expect(refreshAuthMock).not.toHaveBeenCalledWith(AuthType.COMPUTE_ADC); - }); - - it('should skip LOGIN_WITH_GOOGLE and use COMPUTE_ADC directly in headless Cloud Shell', async () => { - vi.stubEnv('CLOUD_SHELL', 'true'); - vi.mocked(isHeadlessMode).mockReturnValue(true); - + it('should attempt COMPUTE_ADC by default and bypass LOGIN_WITH_GOOGLE if successful', async () => { const refreshAuthMock = vi.fn().mockResolvedValue(undefined); - - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); + setupConfigMock(refreshAuthMock); await loadConfig(mockSettings, mockExtensionLoader, taskId); + expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); expect(refreshAuthMock).not.toHaveBeenCalledWith( AuthType.LOGIN_WITH_GOOGLE, ); - expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); }); - it('should skip LOGIN_WITH_GOOGLE and use COMPUTE_ADC directly if GEMINI_CLI_USE_COMPUTE_ADC is true', async () => { - vi.stubEnv('GEMINI_CLI_USE_COMPUTE_ADC', 'true'); - vi.mocked(isHeadlessMode).mockReturnValue(false); // Even if not headless - - const refreshAuthMock = vi.fn().mockResolvedValue(undefined); - - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); + it('should fallback to LOGIN_WITH_GOOGLE if COMPUTE_ADC fails and interactive mode is available', async () => { + vi.mocked(isHeadlessMode).mockReturnValue(false); + const refreshAuthMock = vi.fn().mockImplementation((authType) => { + if (authType === AuthType.COMPUTE_ADC) { + return Promise.reject(new Error('ADC failed')); + } + return Promise.resolve(); + }); + setupConfigMock(refreshAuthMock); await loadConfig(mockSettings, mockExtensionLoader, taskId); - expect(refreshAuthMock).not.toHaveBeenCalledWith( + expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); + expect(refreshAuthMock).toHaveBeenCalledWith( AuthType.LOGIN_WITH_GOOGLE, ); - expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); }); - it('should throw FatalAuthenticationError in headless mode if no ADC fallback available', async () => { + it('should throw FatalAuthenticationError in headless mode if COMPUTE_ADC fails', async () => { vi.mocked(isHeadlessMode).mockReturnValue(true); - const refreshAuthMock = vi.fn().mockResolvedValue(undefined); - - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); + const refreshAuthMock = vi.fn().mockImplementation((authType) => { + if (authType === AuthType.COMPUTE_ADC) { + return Promise.reject(new Error('ADC not found')); + } + return Promise.resolve(); + }); + setupConfigMock(refreshAuthMock); await expect( loadConfig(mockSettings, mockExtensionLoader, taskId), ).rejects.toThrow( - 'Interactive terminal required for LOGIN_WITH_GOOGLE. Run in an interactive terminal or set GEMINI_CLI_USE_COMPUTE_ADC=true to use Application Default Credentials.', + 'COMPUTE_ADC failed: ADC not found. (LOGIN_WITH_GOOGLE fallback skipped due to headless mode. Run in an interactive terminal to use OAuth.)', ); - expect(refreshAuthMock).not.toHaveBeenCalled(); + expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC); + expect(refreshAuthMock).not.toHaveBeenCalledWith( + AuthType.LOGIN_WITH_GOOGLE, + ); }); - it('should include both original and fallback error when COMPUTE_ADC fallback fails', async () => { - vi.stubEnv('CLOUD_SHELL', 'true'); + it('should include both original and fallback error when LOGIN_WITH_GOOGLE fallback fails', async () => { vi.mocked(isHeadlessMode).mockReturnValue(false); const refreshAuthMock = vi.fn().mockImplementation((authType) => { - if (authType === AuthType.LOGIN_WITH_GOOGLE) { - throw new FatalAuthenticationError('OAuth failed'); - } if (authType === AuthType.COMPUTE_ADC) { throw new Error('ADC failed'); } + if (authType === AuthType.LOGIN_WITH_GOOGLE) { + throw new FatalAuthenticationError('OAuth failed'); + } return Promise.resolve(); }); - - vi.mocked(Config).mockImplementation( - (params: unknown) => - ({ - ...(params as object), - initialize: vi.fn(), - waitForMcpInit: vi.fn(), - refreshAuth: refreshAuthMock, - getExperiments: vi.fn().mockReturnValue({ flags: {} }), - getRemoteAdminSettings: vi.fn(), - setRemoteAdminSettings: vi.fn(), - }) as unknown as Config, - ); + setupConfigMock(refreshAuthMock); await expect( loadConfig(mockSettings, mockExtensionLoader, taskId), ).rejects.toThrow( - 'OAuth failed. Fallback to COMPUTE_ADC also failed: ADC failed', + 'OAuth failed. The initial COMPUTE_ADC attempt also failed: ADC failed', ); }); }); diff --git a/packages/a2a-server/src/config/config.ts b/packages/a2a-server/src/config/config.ts index cd4f5df25f..3badd3ff79 100644 --- a/packages/a2a-server/src/config/config.ts +++ b/packages/a2a-server/src/config/config.ts @@ -25,7 +25,6 @@ import { ExperimentFlags, isHeadlessMode, FatalAuthenticationError, - isCloudShell, PolicyDecision, PRIORITY_YOLO_ALLOW_ALL, type TelemetryTarget, @@ -43,7 +42,6 @@ export async function loadConfig( taskId: string, ): Promise { const workspaceDir = process.cwd(); - const adcFilePath = process.env['GOOGLE_APPLICATION_CREDENTIALS']; const folderTrust = settings.folderTrust === true || @@ -192,7 +190,7 @@ export async function loadConfig( await config.waitForMcpInit(); startupProfiler.flush(config); - await refreshAuthentication(config, adcFilePath, 'Config'); + await refreshAuthentication(config, 'Config'); return config; } @@ -263,75 +261,51 @@ function findEnvFile(startDir: string): string | null { async function refreshAuthentication( config: Config, - adcFilePath: string | undefined, logPrefix: string, ): Promise { if (process.env['USE_CCPA']) { logger.info(`[${logPrefix}] Using CCPA Auth:`); + + logger.info(`[${logPrefix}] Attempting COMPUTE_ADC first.`); try { - if (adcFilePath) { - path.resolve(adcFilePath); - } - } catch (e) { - logger.error( - `[${logPrefix}] USE_CCPA env var is true but unable to resolve GOOGLE_APPLICATION_CREDENTIALS file path ${adcFilePath}. Error ${e}`, + await config.refreshAuth(AuthType.COMPUTE_ADC); + logger.info(`[${logPrefix}] COMPUTE_ADC successful.`); + } catch (adcError) { + const adcMessage = + adcError instanceof Error ? adcError.message : String(adcError); + logger.info( + `[${logPrefix}] COMPUTE_ADC failed or not available: ${adcMessage}`, ); - } - const useComputeAdc = process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true'; - const isHeadless = isHeadlessMode(); - const shouldSkipOauth = isHeadless || useComputeAdc; + const useComputeAdc = + process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true'; + const isHeadless = isHeadlessMode(); - if (shouldSkipOauth) { - if (isCloudShell() || useComputeAdc) { - logger.info( - `[${logPrefix}] Skipping LOGIN_WITH_GOOGLE due to ${isHeadless ? 'headless mode' : 'GEMINI_CLI_USE_COMPUTE_ADC'}. Attempting COMPUTE_ADC.`, - ); - try { - await config.refreshAuth(AuthType.COMPUTE_ADC); - logger.info(`[${logPrefix}] COMPUTE_ADC successful.`); - } catch (adcError) { - const adcMessage = - adcError instanceof Error ? adcError.message : String(adcError); - throw new FatalAuthenticationError( - `COMPUTE_ADC failed: ${adcMessage}. (Skipped LOGIN_WITH_GOOGLE due to ${isHeadless ? 'headless mode' : 'GEMINI_CLI_USE_COMPUTE_ADC'})`, - ); - } - } else { + if (isHeadless || useComputeAdc) { + const reason = isHeadless + ? 'headless mode' + : 'GEMINI_CLI_USE_COMPUTE_ADC=true'; throw new FatalAuthenticationError( - `Interactive terminal required for LOGIN_WITH_GOOGLE. Run in an interactive terminal or set GEMINI_CLI_USE_COMPUTE_ADC=true to use Application Default Credentials.`, + `COMPUTE_ADC failed: ${adcMessage}. (LOGIN_WITH_GOOGLE fallback skipped due to ${reason}. Run in an interactive terminal to use OAuth.)`, ); } - } else { + + logger.info( + `[${logPrefix}] COMPUTE_ADC failed, falling back to LOGIN_WITH_GOOGLE.`, + ); try { await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE); } catch (e) { - if ( - e instanceof FatalAuthenticationError && - (isCloudShell() || useComputeAdc) - ) { - logger.warn( - `[${logPrefix}] LOGIN_WITH_GOOGLE failed. Attempting COMPUTE_ADC fallback.`, + if (e instanceof FatalAuthenticationError) { + const originalMessage = e instanceof Error ? e.message : String(e); + throw new FatalAuthenticationError( + `${originalMessage}. The initial COMPUTE_ADC attempt also failed: ${adcMessage}`, ); - try { - await config.refreshAuth(AuthType.COMPUTE_ADC); - logger.info(`[${logPrefix}] COMPUTE_ADC fallback successful.`); - } catch (adcError) { - logger.error( - `[${logPrefix}] COMPUTE_ADC fallback failed: ${adcError}`, - ); - const originalMessage = e instanceof Error ? e.message : String(e); - const adcMessage = - adcError instanceof Error ? adcError.message : String(adcError); - throw new FatalAuthenticationError( - `${originalMessage}. Fallback to COMPUTE_ADC also failed: ${adcMessage}`, - ); - } - } else { - throw e; } + throw e; } } + logger.info( `[${logPrefix}] GOOGLE_CLOUD_PROJECT: ${process.env['GOOGLE_CLOUD_PROJECT']}`, ); diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 0ae523dc94..cb4b645ab3 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -119,7 +119,8 @@ async function initOauthClient( credentials && typeof credentials === 'object' && 'type' in credentials && - credentials.type === 'external_account_authorized_user' + (credentials.type === 'external_account_authorized_user' || + credentials.type === 'service_account') ) { const auth = new GoogleAuth({ scopes: OAUTH_SCOPE, @@ -130,7 +131,7 @@ async function initOauthClient( }); const token = await byoidClient.getAccessToken(); if (token) { - debugLogger.debug('Created BYOID auth client.'); + debugLogger.debug(`Created ${credentials.type} auth client.`); return byoidClient; } }