refactor(telemetry): Improve previous PR that allows telemetry to use the CLI auth and add testing (#14589)

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: {

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);

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;

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'),
+    );
+  });
 });

packages/core/src/telemetry/sdk.ts

@@ -87,6 +87,7 @@ let callbackRegistered = false;
 let authListener: ((newCredentials: JWTInput) => Promise<void>) | undefined =
   undefined;
 const telemetryBuffer: Array<() => void | Promise<void>> = [];
+let activeTelemetryEmail: string | undefined;
 
 export function isTelemetrySdkInitialized(): boolean {
   return telemetryInitialized;
@@ -144,7 +145,20 @@ export async function initializeTelemetry(
   config: Config,
   credentials?: JWTInput,
 ): Promise<void> {
-  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;
   }
 }

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)."
         }
       }
     },