2 files changed, 50 insertions, 99 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index d8cd525b..cec05121 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -5,7 +5,8 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
-import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js';
+import { getOauthClient } from './oauth2.js';
+import { getCachedGoogleAccount } from '../utils/user_account.js';
 import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
 import * as path from 'path';
@@ -66,30 +67,11 @@ describe('oauth2', () => {
     const mockGetAccessToken = vi
       .fn()
       .mockResolvedValue({ token: 'mock-access-token' });
-    const mockRefreshAccessToken = vi.fn().mockImplementation((callback) => {
-      // Mock the callback-style refreshAccessToken method
-      const mockTokensWithIdToken = {
-        access_token: 'test-access-token',
-        refresh_token: 'test-refresh-token',
-        id_token:
-          'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ0ZXN0LWdvb2dsZS1hY2NvdW50LWlkLTEyMyJ9.signature', // Mock JWT with sub: test-google-account-id-123
-      };
-      callback(null, mockTokensWithIdToken);
-    });
-    const mockVerifyIdToken = vi.fn().mockResolvedValue({
-      getPayload: () => ({
-        sub: 'test-google-account-id-123',
-        aud: 'test-audience',
-        iss: 'https://accounts.google.com',
-      }),
-    });
     const mockOAuth2Client = {
       generateAuthUrl: mockGenerateAuthUrl,
       getToken: mockGetToken,
       setCredentials: mockSetCredentials,
       getAccessToken: mockGetAccessToken,
-      refreshAccessToken: mockRefreshAccessToken,
-      verifyIdToken: mockVerifyIdToken,
       credentials: mockTokens,
       on: vi.fn(),
     } as unknown as OAuth2Client;
@@ -103,7 +85,9 @@ describe('oauth2', () => {
     // Mock the UserInfo API response
     (global.fetch as Mock).mockResolvedValue({
       ok: true,
-      json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }),
+      json: vi
+        .fn()
+        .mockResolvedValue({ email: '[email protected]' }),
     } as unknown as Response);
 
     let requestCallback!: http.RequestListener<
@@ -169,18 +153,21 @@ describe('oauth2', () => {
     });
     expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
 
-    // Verify Google Account ID was cached
-    const googleAccountIdPath = path.join(
+    // Verify Google Account was cached
+    const googleAccountPath = path.join(
       tempHomeDir,
       '.gemini',
-      'google_account_id',
+      'google_accounts.json',
     );
-    expect(fs.existsSync(googleAccountIdPath)).toBe(true);
-    const cachedGoogleAccountId = fs.readFileSync(googleAccountIdPath, 'utf-8');
-    expect(cachedGoogleAccountId).toBe('test-google-account-id-123');
+    expect(fs.existsSync(googleAccountPath)).toBe(true);
+    const cachedGoogleAccount = fs.readFileSync(googleAccountPath, 'utf-8');
+    expect(JSON.parse(cachedGoogleAccount)).toEqual({
+      active: '[email protected]',
+      old: [],
+    });
 
-    // Verify the getCachedGoogleAccountId function works
-    expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123');
+    // Verify the getCachedGoogleAccount function works
+    expect(getCachedGoogleAccount()).toBe('[email protected]');
   });
 
   describe('in Cloud Shell', () => {
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index 2d3c04d0..d5f28880 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -16,10 +16,15 @@ import crypto from 'crypto';
 import * as net from 'net';
 import open from 'open';
 import path from 'node:path';
-import { promises as fs, existsSync, readFileSync } from 'node:fs';
+import { promises as fs } from 'node:fs';
 import * as os from 'os';
 import { Config } from '../config/config.js';
 import { getErrorMessage } from '../utils/errors.js';
+import {
+  cacheGoogleAccount,
+  getCachedGoogleAccount,
+  clearCachedGoogleAccount,
+} from '../utils/user_account.js';
 import { AuthType } from '../core/contentGenerator.js';
 import readline from 'node:readline';
 
@@ -50,7 +55,6 @@ const SIGN_IN_FAILURE_URL =
 
 const GEMINI_DIR = '.gemini';
 const CREDENTIAL_FILENAME = 'oauth_creds.json';
-const GOOGLE_ACCOUNT_ID_FILENAME = 'google_account_id';
 
 /**
  * An Authentication URL for updating the credentials of a Oauth2Client
@@ -78,13 +82,10 @@ export async function getOauthClient(
   // If there are cached creds on disk, they always take precedence
   if (await loadCachedCredentials(client)) {
     // Found valid cached credentials.
-    // Check if we need to retrieve Google Account ID
-    if (!getCachedGoogleAccountId()) {
+    // Check if we need to retrieve Google Account ID or Email
+    if (!getCachedGoogleAccount()) {
       try {
-        const googleAccountId = await getRawGoogleAccountId(client);
-        if (googleAccountId) {
-          await cacheGoogleAccountId(googleAccountId);
-        }
+        await fetchAndCacheUserInfo(client);
       } catch {
         // Non-fatal, continue with existing auth.
       }
@@ -237,10 +238,7 @@ async function authWithWeb(client: OAuth2Client): Promise<OauthWebLogin> {
           client.setCredentials(tokens);
           // Retrieve and cache Google Account ID during authentication
           try {
-            const googleAccountId = await getRawGoogleAccountId(client);
-            if (googleAccountId) {
-              await cacheGoogleAccountId(googleAccountId);
-            }
+            await fetchAndCacheUserInfo(client);
           } catch (error) {
             console.error(
               'Failed to retrieve Google Account ID during authentication:',
@@ -326,80 +324,46 @@ function getCachedCredentialPath(): string {
   return path.join(os.homedir(), GEMINI_DIR, CREDENTIAL_FILENAME);
 }
 
-function getGoogleAccountIdCachePath(): string {
-  return path.join(os.homedir(), GEMINI_DIR, GOOGLE_ACCOUNT_ID_FILENAME);
-}
-
-async function cacheGoogleAccountId(googleAccountId: string): Promise<void> {
-  const filePath = getGoogleAccountIdCachePath();
-  await fs.mkdir(path.dirname(filePath), { recursive: true });
-  await fs.writeFile(filePath, googleAccountId, 'utf-8');
-}
-
-export function getCachedGoogleAccountId(): string | null {
-  try {
-    const filePath = getGoogleAccountIdCachePath();
-    if (existsSync(filePath)) {
-      return readFileSync(filePath, 'utf-8').trim() || null;
-    }
-    return null;
-  } catch (error) {
-    console.debug('Error reading cached Google Account ID:', error);
-    return null;
-  }
-}
-
 export async function clearCachedCredentialFile() {
   try {
     await fs.rm(getCachedCredentialPath(), { force: true });
     // Clear the Google Account ID cache when credentials are cleared
-    await fs.rm(getGoogleAccountIdCachePath(), { force: true });
+    await clearCachedGoogleAccount();
   } catch (_) {
     /* empty */
   }
 }
 
-/**
- * Retrieves the authenticated user's Google Account ID from Google's UserInfo API.
- * @param client - The authenticated OAuth2Client
- * @returns The user's Google Account ID or null if not available
- */
-export async function getRawGoogleAccountId(
-  client: OAuth2Client,
-): Promise<string | null> {
+async function fetchAndCacheUserInfo(client: OAuth2Client): Promise<void> {
   try {
-    // 1. Get a new Access Token including the id_token
-    const refreshedTokens = await new Promise<Credentials | null>(
-      (resolve, reject) => {
-        client.refreshAccessToken((err, tokens) => {
-          if (err) {
-            return reject(err);
-          }
-          resolve(tokens ?? null);
-        });
+    const { token } = await client.getAccessToken();
+    if (!token) {
+      return;
+    }
+
+    const response = await fetch(
+      'https://www.googleapis.com/oauth2/v2/userinfo',
+      {
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
       },
     );
 
-    if (!refreshedTokens?.id_token) {
-      console.warn('No id_token obtained after refreshing tokens.');
-      return null;
+    if (!response.ok) {
+      console.error(
+        'Failed to fetch user info:',
+        response.status,
+        response.statusText,
+      );
+      return;
     }
 
-    // 2. Verify the ID token to securely get the user's Google Account ID.
-    const ticket = await client.verifyIdToken({
-      idToken: refreshedTokens.id_token,
-      audience: OAUTH_CLIENT_ID,
-    });
-
-    const payload = ticket.getPayload();
-    if (!payload?.sub) {
-      console.warn('Could not extract sub claim from verified ID token.');
-      return null;
+    const userInfo = await response.json();
+    if (userInfo.email) {
+      await cacheGoogleAccount(userInfo.email);
     }
-
-    return payload.sub;
   } catch (error) {
-    console.error('Error retrieving or verifying Google Account ID:', error);
-    return null;
+    console.error('Error retrieving user info:', error);
   }
 }