refactor(auth): Refactor non-interactive mode auth validation & refresh (#15679)

This commit is contained in:
Shreya Keshive
2025-12-30 11:09:00 -05:00
committed by GitHub
parent fb22f5b8ee
commit 1e08b150f7
5 changed files with 56 additions and 65 deletions
+3
View File
@@ -281,6 +281,7 @@ describe('gemini.tsx main function', () => {
getOutputFormat: () => 'text', getOutputFormat: () => 'text',
getExtensions: () => [], getExtensions: () => [],
getUsageStatisticsEnabled: () => false, getUsageStatisticsEnabled: () => false,
refreshAuth: vi.fn(),
setTerminalBackground: vi.fn(), setTerminalBackground: vi.fn(),
} as unknown as Config; } as unknown as Config;
}); });
@@ -783,6 +784,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getFileFilteringRespectGitIgnore: () => true, getFileFilteringRespectGitIgnore: () => true,
getOutputFormat: () => 'text', getOutputFormat: () => 'text',
getUsageStatisticsEnabled: () => false, getUsageStatisticsEnabled: () => false,
refreshAuth: vi.fn(),
setTerminalBackground: vi.fn(), setTerminalBackground: vi.fn(),
} as any); // eslint-disable-line @typescript-eslint/no-explicit-any } as any); // eslint-disable-line @typescript-eslint/no-explicit-any
@@ -1019,6 +1021,7 @@ describe('gemini.tsx main function kitty protocol', () => {
getFileFilteringRespectGitIgnore: () => true, getFileFilteringRespectGitIgnore: () => true,
getOutputFormat: () => 'text', getOutputFormat: () => 'text',
getUsageStatisticsEnabled: () => false, getUsageStatisticsEnabled: () => false,
refreshAuth: vi.fn(),
setTerminalBackground: vi.fn(), setTerminalBackground: vi.fn(),
} as any); // eslint-disable-line @typescript-eslint/no-explicit-any } as any); // eslint-disable-line @typescript-eslint/no-explicit-any
+23 -12
View File
@@ -401,18 +401,28 @@ export async function main() {
settings.merged.security?.auth?.selectedType && settings.merged.security?.auth?.selectedType &&
!settings.merged.security?.auth?.useExternal !settings.merged.security?.auth?.useExternal
) { ) {
// Validate authentication here because the sandbox will interfere with the Oauth2 web redirect.
try { try {
const err = validateAuthMethod( if (partialConfig.isInteractive()) {
settings.merged.security.auth.selectedType, // Validate authentication here because the sandbox will interfere with the Oauth2 web redirect.
); const err = validateAuthMethod(
if (err) { settings.merged.security.auth.selectedType,
throw new Error(err); );
} if (err) {
throw new Error(err);
}
await partialConfig.refreshAuth( await partialConfig.refreshAuth(
settings.merged.security.auth.selectedType, settings.merged.security.auth.selectedType,
); );
} else {
const authType = await validateNonInteractiveAuth(
settings.merged.security?.auth?.selectedType,
settings.merged.security?.auth?.useExternal,
partialConfig,
settings,
);
await partialConfig.refreshAuth(authType);
}
} catch (err) { } catch (err) {
debugLogger.error('Error authenticating:', err); debugLogger.error('Error authenticating:', err);
await runExitCleanup(); await runExitCleanup();
@@ -667,12 +677,13 @@ export async function main() {
), ),
); );
const nonInteractiveConfig = await validateNonInteractiveAuth( const authType = await validateNonInteractiveAuth(
settings.merged.security?.auth?.selectedType, settings.merged.security?.auth?.selectedType,
settings.merged.security?.auth?.useExternal, settings.merged.security?.auth?.useExternal,
config, config,
settings, settings,
); );
await config.refreshAuth(authType);
if (config.getDebugMode()) { if (config.getDebugMode()) {
debugLogger.log('Session ID: %s', sessionId); debugLogger.log('Session ID: %s', sessionId);
@@ -684,7 +695,7 @@ export async function main() {
initializeOutputListenersAndFlush(); initializeOutputListenersAndFlush();
await runNonInteractive({ await runNonInteractive({
config: nonInteractiveConfig, config,
settings, settings,
input, input,
prompt_id, prompt_id,
+1
View File
@@ -213,6 +213,7 @@ describe('gemini.tsx main function cleanup', () => {
getOutputFormat: vi.fn(() => 'text'), getOutputFormat: vi.fn(() => 'text'),
getUsageStatisticsEnabled: vi.fn(() => false), getUsageStatisticsEnabled: vi.fn(() => false),
setTerminalBackground: vi.fn(), setTerminalBackground: vi.fn(),
refreshAuth: vi.fn(),
} as any); // eslint-disable-line @typescript-eslint/no-explicit-any } as any); // eslint-disable-line @typescript-eslint/no-explicit-any
try { try {
@@ -12,7 +12,6 @@ import {
beforeEach, beforeEach,
afterEach, afterEach,
type MockInstance, type MockInstance,
type Mock,
} from 'vitest'; } from 'vitest';
import { validateNonInteractiveAuth } from './validateNonInterActiveAuth.js'; import { validateNonInteractiveAuth } from './validateNonInterActiveAuth.js';
import { import {
@@ -40,7 +39,6 @@ describe('validateNonInterActiveAuth', () => {
let debugLoggerErrorSpy: ReturnType<typeof vi.spyOn>; let debugLoggerErrorSpy: ReturnType<typeof vi.spyOn>;
let coreEventsEmitFeedbackSpy: MockInstance; let coreEventsEmitFeedbackSpy: MockInstance;
let processExitSpy: MockInstance; let processExitSpy: MockInstance;
let refreshAuthMock: Mock;
let mockSettings: LoadedSettings; let mockSettings: LoadedSettings;
beforeEach(() => { beforeEach(() => {
@@ -62,7 +60,6 @@ describe('validateNonInterActiveAuth', () => {
throw new Error(`process.exit(${code}) called`); throw new Error(`process.exit(${code}) called`);
}); });
vi.spyOn(auth, 'validateAuthMethod').mockReturnValue(null); vi.spyOn(auth, 'validateAuthMethod').mockReturnValue(null);
refreshAuthMock = vi.fn().mockImplementation(async () => 'refreshed');
mockSettings = { mockSettings = {
system: { path: '', settings: {} }, system: { path: '', settings: {} },
systemDefaults: { path: '', settings: {} }, systemDefaults: { path: '', settings: {} },
@@ -105,7 +102,6 @@ describe('validateNonInterActiveAuth', () => {
it('exits if no auth type is configured or env vars set', async () => { it('exits if no auth type is configured or env vars set', async () => {
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT),
getContentGeneratorConfig: vi getContentGeneratorConfig: vi
.fn() .fn()
@@ -134,61 +130,57 @@ describe('validateNonInterActiveAuth', () => {
it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set', async () => { it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set', async () => {
process.env['GOOGLE_GENAI_USE_GCA'] = 'true'; process.env['GOOGLE_GENAI_USE_GCA'] = 'true';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.LOGIN_WITH_GOOGLE); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses USE_GEMINI if GEMINI_API_KEY is set', async () => { it('uses USE_GEMINI if GEMINI_API_KEY is set', async () => {
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_GEMINI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses USE_VERTEX_AI if GOOGLE_GENAI_USE_VERTEXAI is true (with GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION)', async () => { it('uses USE_VERTEX_AI if GOOGLE_GENAI_USE_VERTEXAI is true (with GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION)', async () => {
process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project'; process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project';
process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1'; process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_VERTEX_AI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses USE_VERTEX_AI if GOOGLE_GENAI_USE_VERTEXAI is true and GOOGLE_API_KEY is set', async () => { it('uses USE_VERTEX_AI if GOOGLE_GENAI_USE_VERTEXAI is true and GOOGLE_API_KEY is set', async () => {
process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
process.env['GOOGLE_API_KEY'] = 'vertex-api-key'; process.env['GOOGLE_API_KEY'] = 'vertex-api-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_VERTEX_AI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set, even with other env vars', async () => { it('uses LOGIN_WITH_GOOGLE if GOOGLE_GENAI_USE_GCA is set, even with other env vars', async () => {
@@ -197,16 +189,15 @@ describe('validateNonInterActiveAuth', () => {
process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project'; process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project';
process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1'; process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.LOGIN_WITH_GOOGLE); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses USE_VERTEX_AI if both GEMINI_API_KEY and GOOGLE_GENAI_USE_VERTEXAI are set', async () => { it('uses USE_VERTEX_AI if both GEMINI_API_KEY and GOOGLE_GENAI_USE_VERTEXAI are set', async () => {
@@ -214,16 +205,15 @@ describe('validateNonInterActiveAuth', () => {
process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project'; process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project';
process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1'; process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_VERTEX_AI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses USE_GEMINI if GOOGLE_GENAI_USE_VERTEXAI is false, GEMINI_API_KEY is set, and project/location are available', async () => { it('uses USE_GEMINI if GOOGLE_GENAI_USE_VERTEXAI is false, GEMINI_API_KEY is set, and project/location are available', async () => {
@@ -231,37 +221,34 @@ describe('validateNonInterActiveAuth', () => {
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project'; process.env['GOOGLE_CLOUD_PROJECT'] = 'test-project';
process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1'; process.env['GOOGLE_CLOUD_LOCATION'] = 'us-central1';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_GEMINI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('uses configuredAuthType over environment variables', async () => { it('uses configuredAuthType over environment variables', async () => {
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
AuthType.LOGIN_WITH_GOOGLE, AuthType.LOGIN_WITH_GOOGLE,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.LOGIN_WITH_GOOGLE); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('exits if validateAuthMethod returns error', async () => { it('exits if validateAuthMethod returns error', async () => {
// Mock validateAuthMethod to return error // Mock validateAuthMethod to return error
vi.spyOn(auth, 'validateAuthMethod').mockReturnValue('Auth error!'); vi.spyOn(auth, 'validateAuthMethod').mockReturnValue('Auth error!');
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT),
getContentGeneratorConfig: vi getContentGeneratorConfig: vi
.fn() .fn()
@@ -291,9 +278,7 @@ describe('validateNonInterActiveAuth', () => {
const validateAuthMethodSpy = vi const validateAuthMethodSpy = vi
.spyOn(auth, 'validateAuthMethod') .spyOn(auth, 'validateAuthMethod')
.mockReturnValue('Auth error!'); .mockReturnValue('Auth error!');
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
// Even with an invalid auth type, it should not exit // Even with an invalid auth type, it should not exit
// because validation is skipped. // because validation is skipped.
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
@@ -307,30 +292,26 @@ describe('validateNonInterActiveAuth', () => {
expect(debugLoggerErrorSpy).not.toHaveBeenCalled(); expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
expect(coreEventsEmitFeedbackSpy).not.toHaveBeenCalled(); expect(coreEventsEmitFeedbackSpy).not.toHaveBeenCalled();
expect(processExitSpy).not.toHaveBeenCalled(); expect(processExitSpy).not.toHaveBeenCalled();
// We still expect refreshAuth to be called with the (invalid) type
expect(refreshAuthMock).toHaveBeenCalledWith('invalid-auth-type');
}); });
it('succeeds if effectiveAuthType matches enforcedAuthType', async () => { it('succeeds if effectiveAuthType matches enforcedAuthType', async () => {
mockSettings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI; mockSettings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI;
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({});
refreshAuth: refreshAuthMock,
});
await validateNonInteractiveAuth( await validateNonInteractiveAuth(
undefined, undefined,
undefined, undefined,
nonInteractiveConfig, nonInteractiveConfig,
mockSettings, mockSettings,
); );
expect(refreshAuthMock).toHaveBeenCalledWith(AuthType.USE_GEMINI); expect(processExitSpy).not.toHaveBeenCalled();
expect(debugLoggerErrorSpy).not.toHaveBeenCalled();
}); });
it('exits if configuredAuthType does not match enforcedAuthType', async () => { it('exits if configuredAuthType does not match enforcedAuthType', async () => {
mockSettings.merged.security!.auth!.enforcedType = mockSettings.merged.security!.auth!.enforcedType =
AuthType.LOGIN_WITH_GOOGLE; AuthType.LOGIN_WITH_GOOGLE;
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT),
}); });
try { try {
@@ -359,7 +340,6 @@ describe('validateNonInterActiveAuth', () => {
AuthType.LOGIN_WITH_GOOGLE; AuthType.LOGIN_WITH_GOOGLE;
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.TEXT),
}); });
try { try {
@@ -386,7 +366,6 @@ describe('validateNonInterActiveAuth', () => {
describe('JSON output mode', () => { describe('JSON output mode', () => {
it(`prints JSON error when no auth is configured and exits with code ${ExitCodes.FATAL_AUTHENTICATION_ERROR}`, async () => { it(`prints JSON error when no auth is configured and exits with code ${ExitCodes.FATAL_AUTHENTICATION_ERROR}`, async () => {
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON),
getContentGeneratorConfig: vi getContentGeneratorConfig: vi
.fn() .fn()
@@ -421,7 +400,6 @@ describe('validateNonInterActiveAuth', () => {
it(`prints JSON error when enforced auth mismatches current auth and exits with code ${ExitCodes.FATAL_AUTHENTICATION_ERROR}`, async () => { it(`prints JSON error when enforced auth mismatches current auth and exits with code ${ExitCodes.FATAL_AUTHENTICATION_ERROR}`, async () => {
mockSettings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI; mockSettings.merged.security!.auth!.enforcedType = AuthType.USE_GEMINI;
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON),
getContentGeneratorConfig: vi getContentGeneratorConfig: vi
.fn() .fn()
@@ -460,7 +438,6 @@ describe('validateNonInterActiveAuth', () => {
process.env['GEMINI_API_KEY'] = 'fake-key'; process.env['GEMINI_API_KEY'] = 'fake-key';
const nonInteractiveConfig = createLocalMockConfig({ const nonInteractiveConfig = createLocalMockConfig({
refreshAuth: refreshAuthMock,
getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON), getOutputFormat: vi.fn().mockReturnValue(OutputFormat.JSON),
getContentGeneratorConfig: vi getContentGeneratorConfig: vi
.fn() .fn()
@@ -61,8 +61,7 @@ export async function validateNonInteractiveAuth(
} }
} }
await nonInteractiveConfig.refreshAuth(authType); return authType;
return nonInteractiveConfig;
} catch (error) { } catch (error) {
if (nonInteractiveConfig.getOutputFormat() === OutputFormat.JSON) { if (nonInteractiveConfig.getOutputFormat() === OutputFormat.JSON) {
handleError( handleError(