fix(a2a-server): prioritize ADC before evaluating headless constraints for auth initialization (#23614)

This commit is contained in:
Spencer
2026-03-30 08:26:15 -04:00
committed by GitHub
parent d9d2ce36f2
commit a255529c6b
3 changed files with 80 additions and 195 deletions
+49 -139
View File
@@ -424,7 +424,22 @@ describe('loadConfig', () => {
}); });
}); });
describe('authentication fallback', () => { describe('authentication logic', () => {
const setupConfigMock = (refreshAuthMock: ReturnType<typeof vi.fn>) => {
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(() => { beforeEach(() => {
vi.stubEnv('USE_CCPA', 'true'); vi.stubEnv('USE_CCPA', 'true');
vi.stubEnv('GEMINI_API_KEY', ''); vi.stubEnv('GEMINI_API_KEY', '');
@@ -434,182 +449,77 @@ describe('loadConfig', () => {
vi.unstubAllEnvs(); vi.unstubAllEnvs();
}); });
it('should fall back to COMPUTE_ADC in Cloud Shell if LOGIN_WITH_GOOGLE fails', async () => { it('should attempt COMPUTE_ADC by default and bypass LOGIN_WITH_GOOGLE if successful', 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);
const refreshAuthMock = vi.fn().mockResolvedValue(undefined); const refreshAuthMock = vi.fn().mockResolvedValue(undefined);
setupConfigMock(refreshAuthMock);
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); await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC);
expect(refreshAuthMock).not.toHaveBeenCalledWith( expect(refreshAuthMock).not.toHaveBeenCalledWith(
AuthType.LOGIN_WITH_GOOGLE, 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 () => { it('should fallback to LOGIN_WITH_GOOGLE if COMPUTE_ADC fails and interactive mode is available', async () => {
vi.stubEnv('GEMINI_CLI_USE_COMPUTE_ADC', 'true'); vi.mocked(isHeadlessMode).mockReturnValue(false);
vi.mocked(isHeadlessMode).mockReturnValue(false); // Even if not headless const refreshAuthMock = vi.fn().mockImplementation((authType) => {
if (authType === AuthType.COMPUTE_ADC) {
const refreshAuthMock = vi.fn().mockResolvedValue(undefined); return Promise.reject(new Error('ADC failed'));
}
vi.mocked(Config).mockImplementation( return Promise.resolve();
(params: unknown) => });
({ setupConfigMock(refreshAuthMock);
...(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); await loadConfig(mockSettings, mockExtensionLoader, taskId);
expect(refreshAuthMock).not.toHaveBeenCalledWith( expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.COMPUTE_ADC);
expect(refreshAuthMock).toHaveBeenCalledWith(
AuthType.LOGIN_WITH_GOOGLE, 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); vi.mocked(isHeadlessMode).mockReturnValue(true);
const refreshAuthMock = vi.fn().mockResolvedValue(undefined); const refreshAuthMock = vi.fn().mockImplementation((authType) => {
if (authType === AuthType.COMPUTE_ADC) {
vi.mocked(Config).mockImplementation( return Promise.reject(new Error('ADC not found'));
(params: unknown) => }
({ return Promise.resolve();
...(params as object), });
initialize: vi.fn(), setupConfigMock(refreshAuthMock);
waitForMcpInit: vi.fn(),
refreshAuth: refreshAuthMock,
getExperiments: vi.fn().mockReturnValue({ flags: {} }),
getRemoteAdminSettings: vi.fn(),
setRemoteAdminSettings: vi.fn(),
}) as unknown as Config,
);
await expect( await expect(
loadConfig(mockSettings, mockExtensionLoader, taskId), loadConfig(mockSettings, mockExtensionLoader, taskId),
).rejects.toThrow( ).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 () => { it('should include both original and fallback error when LOGIN_WITH_GOOGLE fallback fails', async () => {
vi.stubEnv('CLOUD_SHELL', 'true');
vi.mocked(isHeadlessMode).mockReturnValue(false); vi.mocked(isHeadlessMode).mockReturnValue(false);
const refreshAuthMock = vi.fn().mockImplementation((authType) => { const refreshAuthMock = vi.fn().mockImplementation((authType) => {
if (authType === AuthType.LOGIN_WITH_GOOGLE) {
throw new FatalAuthenticationError('OAuth failed');
}
if (authType === AuthType.COMPUTE_ADC) { if (authType === AuthType.COMPUTE_ADC) {
throw new Error('ADC failed'); throw new Error('ADC failed');
} }
if (authType === AuthType.LOGIN_WITH_GOOGLE) {
throw new FatalAuthenticationError('OAuth failed');
}
return Promise.resolve(); return Promise.resolve();
}); });
setupConfigMock(refreshAuthMock);
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( await expect(
loadConfig(mockSettings, mockExtensionLoader, taskId), loadConfig(mockSettings, mockExtensionLoader, taskId),
).rejects.toThrow( ).rejects.toThrow(
'OAuth failed. Fallback to COMPUTE_ADC also failed: ADC failed', 'OAuth failed. The initial COMPUTE_ADC attempt also failed: ADC failed',
); );
}); });
}); });
+28 -54
View File
@@ -25,7 +25,6 @@ import {
ExperimentFlags, ExperimentFlags,
isHeadlessMode, isHeadlessMode,
FatalAuthenticationError, FatalAuthenticationError,
isCloudShell,
PolicyDecision, PolicyDecision,
PRIORITY_YOLO_ALLOW_ALL, PRIORITY_YOLO_ALLOW_ALL,
type TelemetryTarget, type TelemetryTarget,
@@ -43,7 +42,6 @@ export async function loadConfig(
taskId: string, taskId: string,
): Promise<Config> { ): Promise<Config> {
const workspaceDir = process.cwd(); const workspaceDir = process.cwd();
const adcFilePath = process.env['GOOGLE_APPLICATION_CREDENTIALS'];
const folderTrust = const folderTrust =
settings.folderTrust === true || settings.folderTrust === true ||
@@ -192,7 +190,7 @@ export async function loadConfig(
await config.waitForMcpInit(); await config.waitForMcpInit();
startupProfiler.flush(config); startupProfiler.flush(config);
await refreshAuthentication(config, adcFilePath, 'Config'); await refreshAuthentication(config, 'Config');
return config; return config;
} }
@@ -263,75 +261,51 @@ function findEnvFile(startDir: string): string | null {
async function refreshAuthentication( async function refreshAuthentication(
config: Config, config: Config,
adcFilePath: string | undefined,
logPrefix: string, logPrefix: string,
): Promise<void> { ): Promise<void> {
if (process.env['USE_CCPA']) { if (process.env['USE_CCPA']) {
logger.info(`[${logPrefix}] Using CCPA Auth:`); logger.info(`[${logPrefix}] Using CCPA Auth:`);
logger.info(`[${logPrefix}] Attempting COMPUTE_ADC first.`);
try { try {
if (adcFilePath) { await config.refreshAuth(AuthType.COMPUTE_ADC);
path.resolve(adcFilePath); logger.info(`[${logPrefix}] COMPUTE_ADC successful.`);
} } catch (adcError) {
} catch (e) { const adcMessage =
logger.error( adcError instanceof Error ? adcError.message : String(adcError);
`[${logPrefix}] USE_CCPA env var is true but unable to resolve GOOGLE_APPLICATION_CREDENTIALS file path ${adcFilePath}. Error ${e}`, logger.info(
`[${logPrefix}] COMPUTE_ADC failed or not available: ${adcMessage}`,
); );
}
const useComputeAdc = process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true'; const useComputeAdc =
const isHeadless = isHeadlessMode(); process.env['GEMINI_CLI_USE_COMPUTE_ADC'] === 'true';
const shouldSkipOauth = isHeadless || useComputeAdc; const isHeadless = isHeadlessMode();
if (shouldSkipOauth) { if (isHeadless || useComputeAdc) {
if (isCloudShell() || useComputeAdc) { const reason = isHeadless
logger.info( ? 'headless mode'
`[${logPrefix}] Skipping LOGIN_WITH_GOOGLE due to ${isHeadless ? 'headless mode' : 'GEMINI_CLI_USE_COMPUTE_ADC'}. Attempting COMPUTE_ADC.`, : 'GEMINI_CLI_USE_COMPUTE_ADC=true';
);
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 {
throw new FatalAuthenticationError( 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 { try {
await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE); await config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
} catch (e) { } catch (e) {
if ( if (e instanceof FatalAuthenticationError) {
e instanceof FatalAuthenticationError && const originalMessage = e instanceof Error ? e.message : String(e);
(isCloudShell() || useComputeAdc) throw new FatalAuthenticationError(
) { `${originalMessage}. The initial COMPUTE_ADC attempt also failed: ${adcMessage}`,
logger.warn(
`[${logPrefix}] LOGIN_WITH_GOOGLE failed. Attempting COMPUTE_ADC fallback.`,
); );
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( logger.info(
`[${logPrefix}] GOOGLE_CLOUD_PROJECT: ${process.env['GOOGLE_CLOUD_PROJECT']}`, `[${logPrefix}] GOOGLE_CLOUD_PROJECT: ${process.env['GOOGLE_CLOUD_PROJECT']}`,
); );
+3 -2
View File
@@ -119,7 +119,8 @@ async function initOauthClient(
credentials && credentials &&
typeof credentials === 'object' && typeof credentials === 'object' &&
'type' in credentials && 'type' in credentials &&
credentials.type === 'external_account_authorized_user' (credentials.type === 'external_account_authorized_user' ||
credentials.type === 'service_account')
) { ) {
const auth = new GoogleAuth({ const auth = new GoogleAuth({
scopes: OAUTH_SCOPE, scopes: OAUTH_SCOPE,
@@ -130,7 +131,7 @@ async function initOauthClient(
}); });
const token = await byoidClient.getAccessToken(); const token = await byoidClient.getAccessToken();
if (token) { if (token) {
debugLogger.debug('Created BYOID auth client.'); debugLogger.debug(`Created ${credentials.type} auth client.`);
return byoidClient; return byoidClient;
} }
} }