diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 2278b6b825..f103cf2db9 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -1738,6 +1738,11 @@ export const SETTINGS_SCHEMA_DEFINITIONS: Record< type: 'boolean', description: 'Whether to forward telemetry to an OTLP collector.', }, + useCliAuth: { + type: 'boolean', + description: + 'Whether to use CLI authentication for telemetry (only for in-process exporters).', + }, }, }, BugCommandSettings: { diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index 1b9b8fa0c9..f5033f8307 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/packages/core/src/code_assist/oauth2.test.ts @@ -278,6 +278,7 @@ describe('oauth2', () => { }; const mockAuthUrl = 'https://example.com/auth-user-code'; const mockCode = 'test-user-code'; + const mockTokens = { access_token: 'test-access-token-user-code', refresh_token: 'test-refresh-token-user-code', @@ -285,7 +286,6 @@ describe('oauth2', () => { const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); - const mockSetCredentials = vi.fn(); const mockGenerateCodeVerifierAsync = vi .fn() .mockResolvedValue(mockCodeVerifier); @@ -293,10 +293,13 @@ describe('oauth2', () => { const mockOAuth2Client = { generateAuthUrl: mockGenerateAuthUrl, getToken: mockGetToken, - setCredentials: mockSetCredentials, generateCodeVerifierAsync: mockGenerateCodeVerifierAsync, on: vi.fn(), + credentials: {}, } as unknown as OAuth2Client; + mockOAuth2Client.setCredentials = vi.fn().mockImplementation((creds) => { + mockOAuth2Client.credentials = creds; + }); vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client); const mockReadline = { @@ -328,7 +331,83 @@ describe('oauth2', () => { codeVerifier: mockCodeVerifier.codeVerifier, redirect_uri: 'https://codeassist.google.com/authcode', }); - expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); + expect(mockOAuth2Client.setCredentials).toHaveBeenCalledWith(mockTokens); + }); + + it('should cache Google Account when logging in with user code', async () => { + const mockConfigWithNoBrowser = { + getNoBrowser: () => true, + getProxy: () => 'http://test.proxy.com:8080', + isBrowserLaunchSuppressed: () => true, + } as unknown as Config; + + const mockCodeVerifier = { + codeChallenge: 'test-challenge', + codeVerifier: 'test-verifier', + }; + const mockAuthUrl = 'https://example.com/auth-user-code'; + const mockCode = 'test-user-code'; + const mockTokens = { + access_token: 'test-access-token-user-code', + refresh_token: 'test-refresh-token-user-code', + }; + + const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); + const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); + const mockGenerateCodeVerifierAsync = vi + .fn() + .mockResolvedValue(mockCodeVerifier); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'test-access-token-user-code' }); + + const mockOAuth2Client = { + generateAuthUrl: mockGenerateAuthUrl, + getToken: mockGetToken, + generateCodeVerifierAsync: mockGenerateCodeVerifierAsync, + getAccessToken: mockGetAccessToken, + on: vi.fn(), + credentials: {}, + } as unknown as OAuth2Client; + mockOAuth2Client.setCredentials = vi.fn().mockImplementation((creds) => { + mockOAuth2Client.credentials = creds; + }); + vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client); + + vi.spyOn(crypto, 'randomBytes').mockReturnValue('test-state' as never); + + const mockReadline = { + question: vi.fn((_query, callback) => callback(mockCode)), + close: vi.fn(), + on: vi.fn(), + }; + (readline.createInterface as Mock).mockReturnValue(mockReadline); + + // Mock User Info API + vi.mocked(global.fetch).mockResolvedValue({ + ok: true, + json: vi + .fn() + .mockResolvedValue({ email: 'test-user-code-account@gmail.com' }), + } as unknown as Response); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfigWithNoBrowser); + + // Verify Google Account was cached + const googleAccountPath = path.join( + tempHomeDir, + GEMINI_DIR, + 'google_accounts.json', + ); + + expect(fs.existsSync(googleAccountPath)).toBe(true); + if (fs.existsSync(googleAccountPath)) { + const cachedGoogleAccount = fs.readFileSync(googleAccountPath, 'utf-8'); + expect(JSON.parse(cachedGoogleAccount)).toEqual({ + active: 'test-user-code-account@gmail.com', + old: [], + }); + } }); describe('in Cloud Shell', () => { @@ -894,12 +973,17 @@ describe('oauth2', () => { const mockOAuth2Client = { generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), getToken: vi.fn().mockResolvedValue({ tokens: mockTokens }), - setCredentials: vi.fn(), getAccessToken: vi .fn() .mockResolvedValue({ token: 'test-access-token' }), on: vi.fn(), + credentials: {}, } as unknown as OAuth2Client; + mockOAuth2Client.setCredentials = vi + .fn() + .mockImplementation((creds) => { + mockOAuth2Client.credentials = creds; + }); vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client); vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index a99dde03f5..d2399083d4 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -56,6 +56,7 @@ async function triggerPostAuthCallbacks(tokens: Credentials) { client_secret: OAUTH_CLIENT_SECRET, refresh_token: tokens.refresh_token ?? undefined, // Ensure null is not passed type: 'authorized_user', + client_email: userAccountManager.getCachedGoogleAccount() ?? undefined, }; // Execute all registered post-authentication callbacks. @@ -255,6 +256,18 @@ async function initOauthClient( 'Failed to authenticate with user code.', ); } + + // Retrieve and cache Google Account ID after successful user code auth + try { + await fetchAndCacheUserInfo(client); + } catch (error) { + debugLogger.warn( + 'Failed to retrieve Google Account ID during authentication:', + getErrorMessage(error), + ); + } + + await triggerPostAuthCallbacks(client.credentials); } else { const webLogin = await authWithWeb(client); @@ -318,6 +331,8 @@ async function initOauthClient( severity: 'info', message: 'Authentication succeeded\n', }); + + await triggerPostAuthCallbacks(client.credentials); } return client; diff --git a/packages/core/src/telemetry/sdk.test.ts b/packages/core/src/telemetry/sdk.test.ts index caaf1e69a3..c17f8b159f 100644 --- a/packages/core/src/telemetry/sdk.test.ts +++ b/packages/core/src/telemetry/sdk.test.ts @@ -19,7 +19,7 @@ import { OTLPLogExporter as OTLPLogExporterHttp } from '@opentelemetry/exporter- import { OTLPMetricExporter as OTLPMetricExporterHttp } from '@opentelemetry/exporter-metrics-otlp-http'; import { ConsoleSpanExporter } from '@opentelemetry/sdk-trace-node'; import { NodeSDK } from '@opentelemetry/sdk-node'; -import { GoogleAuth } from 'google-auth-library'; +import { GoogleAuth, type JWTInput } from 'google-auth-library'; import { GcpTraceExporter, GcpLogExporter, @@ -327,8 +327,7 @@ describe('Telemetry SDK', () => { await vi.waitFor(() => { // Check if debugLogger was called, which indicates the listener ran expect(debugLogger.log).toHaveBeenCalledWith( - 'Telemetry reinit with credentials: ', - mockCredentials, + 'Telemetry reinit with credentials.', ); // Should use GCP exporters now with the project ID @@ -370,4 +369,36 @@ describe('Telemetry SDK', () => { ); expect(NodeSDK.prototype.start).not.toHaveBeenCalled(); }); + it('should log error when re-initializing with different credentials', async () => { + const creds1 = { client_email: 'user1@example.com' }; + const creds2 = { client_email: 'user2@example.com' }; + + // 1. Initialize with first account + await initializeTelemetry(mockConfig, creds1 as JWTInput); + + // 2. Attempt to initialize with second account + await initializeTelemetry(mockConfig, creds2 as JWTInput); + + // 3. Verify error log + expect(debugLogger.error).toHaveBeenCalledWith( + expect.stringContaining( + 'Telemetry credentials have changed (from user1@example.com to user2@example.com)', + ), + ); + }); + + it('should NOT log error when re-initializing with SAME credentials', async () => { + const creds1 = { client_email: 'user1@example.com' }; + + // 1. Initialize with first account + await initializeTelemetry(mockConfig, creds1 as JWTInput); + + // 2. Attempt to initialize with same account + await initializeTelemetry(mockConfig, creds1 as JWTInput); + + // 3. Verify NO error log + expect(debugLogger.error).not.toHaveBeenCalledWith( + expect.stringContaining('Telemetry credentials have changed'), + ); + }); }); diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts index 71e21a11a8..b3bc2d5969 100644 --- a/packages/core/src/telemetry/sdk.ts +++ b/packages/core/src/telemetry/sdk.ts @@ -87,6 +87,7 @@ let callbackRegistered = false; let authListener: ((newCredentials: JWTInput) => Promise) | undefined = undefined; const telemetryBuffer: Array<() => void | Promise> = []; +let activeTelemetryEmail: string | undefined; export function isTelemetrySdkInitialized(): boolean { return telemetryInitialized; @@ -144,7 +145,20 @@ export async function initializeTelemetry( config: Config, credentials?: JWTInput, ): Promise { - if (telemetryInitialized || !config.getTelemetryEnabled()) { + if (!config.getTelemetryEnabled()) { + return; + } + + if (telemetryInitialized) { + if ( + credentials?.client_email && + activeTelemetryEmail && + credentials.client_email !== activeTelemetryEmail + ) { + const message = `Telemetry credentials have changed (from ${activeTelemetryEmail} to ${credentials.client_email}), but telemetry cannot be re-initialized in this process. Please restart the CLI to use the new account for telemetry.`; + debugLogger.error(message); + console.error(message); + } return; } @@ -165,10 +179,7 @@ export async function initializeTelemetry( callbackRegistered = true; authListener = async (newCredentials: JWTInput) => { if (config.getTelemetryEnabled() && config.getTelemetryUseCliAuth()) { - debugLogger.log( - 'Telemetry reinit with credentials: ', - newCredentials, - ); + debugLogger.log('Telemetry reinit with credentials.'); await initializeTelemetry(config, newCredentials); } }; @@ -294,6 +305,7 @@ export async function initializeTelemetry( debugLogger.log('OpenTelemetry SDK started successfully.'); } telemetryInitialized = true; + activeTelemetryEmail = credentials?.client_email; initializeMetrics(config); void flushTelemetryBuffer(); } catch (error) { @@ -366,5 +378,6 @@ export async function shutdownTelemetry( authListener = undefined; } callbackRegistered = false; + activeTelemetryEmail = undefined; } } diff --git a/schemas/settings.schema.json b/schemas/settings.schema.json index 0dde0dbc63..93c95bb659 100644 --- a/schemas/settings.schema.json +++ b/schemas/settings.schema.json @@ -1609,6 +1609,10 @@ "useCollector": { "type": "boolean", "description": "Whether to forward telemetry to an OTLP collector." + }, + "useCliAuth": { + "type": "boolean", + "description": "Whether to use CLI authentication for telemetry (only for in-process exporters)." } } },