Allow telemetry exporters to GCP to utilize user's login credentials, if requested (#13778)

packages/core/src/code_assist/oauth2.test.ts

@@ -12,6 +12,7 @@ import {
   resetOauthClientForTesting,
   clearCachedCredentialFile,
   clearOauthClientCache,
+  authEvents,
 } from './oauth2.js';
 import { UserAccountManager } from '../utils/userAccountManager.js';
 import { OAuth2Client, Compute, GoogleAuth } from 'google-auth-library';
@@ -109,13 +110,18 @@ describe('oauth2', () => {
       const mockGetAccessToken = vi
         .fn()
         .mockResolvedValue({ token: 'mock-access-token' });
+      let tokensListener: ((tokens: Credentials) => void) | undefined;
       const mockOAuth2Client = {
         generateAuthUrl: mockGenerateAuthUrl,
         getToken: mockGetToken,
         setCredentials: mockSetCredentials,
         getAccessToken: mockGetAccessToken,
         credentials: mockTokens,
-        on: vi.fn(),
+        on: vi.fn((event, listener) => {
+          if (event === 'tokens') {
+            tokensListener = listener;
+          }
+        }),
       } as unknown as OAuth2Client;
       vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
 
@@ -195,6 +201,11 @@ describe('oauth2', () => {
       });
       expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
 
+      // Manually trigger the 'tokens' event listener
+      if (tokensListener) {
+        await tokensListener(mockTokens);
+      }
+
       // Verify Google Account was cached
       const googleAccountPath = path.join(
         tempHomeDir,
@@ -215,6 +226,45 @@ describe('oauth2', () => {
       );
     });
 
+    it('should clear credentials file', async () => {
+      // Setup initial state with files
+      const credsPath = path.join(tempHomeDir, GEMINI_DIR, 'oauth_creds.json');
+
+      await fs.promises.mkdir(path.dirname(credsPath), { recursive: true });
+      await fs.promises.writeFile(credsPath, '{}');
+
+      await clearCachedCredentialFile();
+
+      expect(fs.existsSync(credsPath)).toBe(false);
+    });
+
+    it('should emit post_auth event when loading cached credentials', async () => {
+      const cachedCreds = { refresh_token: 'cached-token' };
+      const credsPath = path.join(tempHomeDir, GEMINI_DIR, 'oauth_creds.json');
+      await fs.promises.mkdir(path.dirname(credsPath), { recursive: true });
+      await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds));
+
+      const mockClient = {
+        setCredentials: vi.fn(),
+        getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }),
+        getTokenInfo: vi.fn().mockResolvedValue({}),
+        on: vi.fn(),
+      };
+      vi.mocked(OAuth2Client).mockImplementation(
+        () => mockClient as unknown as OAuth2Client,
+      );
+
+      const eventPromise = new Promise<void>((resolve) => {
+        authEvents.once('post_auth', (creds) => {
+          expect(creds.refresh_token).toBe('cached-token');
+          resolve();
+        });
+      });
+
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+      await eventPromise;
+    });
+
     it('should perform login with user code', async () => {
       const mockConfigWithNoBrowser = {
         getNoBrowser: () => true,

packages/core/src/code_assist/oauth2.ts

@@ -15,6 +15,7 @@ import * as http from 'node:http';
 import url from 'node:url';
 import crypto from 'node:crypto';
 import * as net from 'node:net';
+import { EventEmitter } from 'node:events';
 import open from 'open';
 import path from 'node:path';
 import { promises as fs } from 'node:fs';
@@ -45,6 +46,22 @@ import {
 } from '../utils/terminal.js';
 import { coreEvents, CoreEvent } from '../utils/events.js';
 
+export const authEvents = new EventEmitter();
+
+async function triggerPostAuthCallbacks(tokens: Credentials) {
+  // Construct a JWTInput object to pass to callbacks, as this is the
+  // type expected by the downstream Google Cloud client libraries.
+  const jwtInput: JWTInput = {
+    client_id: OAUTH_CLIENT_ID,
+    client_secret: OAUTH_CLIENT_SECRET,
+    refresh_token: tokens.refresh_token ?? undefined, // Ensure null is not passed
+    type: 'authorized_user',
+  };
+
+  // Execute all registered post-authentication callbacks.
+  authEvents.emit('post_auth', jwtInput);
+}
+
 const userAccountManager = new UserAccountManager();
 
 //  OAuth Client ID used to initiate OAuth2Client class.
@@ -139,6 +156,8 @@ async function initOauthClient(
     } else {
       await cacheCredentials(tokens);
     }
+
+    await triggerPostAuthCallbacks(tokens);
   });
 
   if (credentials) {
@@ -162,6 +181,8 @@ async function initOauthClient(
           }
         }
         debugLogger.log('Loaded cached credentials.');
+        await triggerPostAuthCallbacks(credentials as Credentials);
+
         return client;
       }
     } catch (error) {
@@ -570,19 +591,6 @@ async function fetchCachedCredentials(): Promise<
   return null;
 }
 
-async function cacheCredentials(credentials: Credentials) {
-  const filePath = Storage.getOAuthCredsPath();
-  await fs.mkdir(path.dirname(filePath), { recursive: true });
-
-  const credString = JSON.stringify(credentials, null, 2);
-  await fs.writeFile(filePath, credString, { mode: 0o600 });
-  try {
-    await fs.chmod(filePath, 0o600);
-  } catch {
-    /* empty */
-  }
-}
-
 export function clearOauthClientCache() {
   oauthClientPromises.clear();
 }
@@ -640,3 +648,16 @@ async function fetchAndCacheUserInfo(client: OAuth2Client): Promise<void> {
 export function resetOauthClientForTesting() {
   oauthClientPromises.clear();
 }
+
+async function cacheCredentials(credentials: Credentials) {
+  const filePath = Storage.getOAuthCredsPath();
+  await fs.mkdir(path.dirname(filePath), { recursive: true });
+
+  const credString = JSON.stringify(credentials, null, 2);
+  await fs.writeFile(filePath, credString, { mode: 0o600 });
+  try {
+    await fs.chmod(filePath, 0o600);
+  } catch {
+    /* empty */
+  }
+}