2 files changed, 124 insertions, 2 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index 0f5b791b..7e3c38f0 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -5,7 +5,7 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { getOauthClient } from './oauth2.js';
+import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js';
 import { OAuth2Client } from 'google-auth-library';
 import * as fs from 'fs';
 import * as path from 'path';
@@ -27,6 +27,9 @@ vi.mock('http');
 vi.mock('open');
 vi.mock('crypto');
 
+// Mock fetch globally
+global.fetch = vi.fn();
+
 describe('oauth2', () => {
   let tempHomeDir: string;
 
@@ -52,10 +55,14 @@ describe('oauth2', () => {
     const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl);
     const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens });
     const mockSetCredentials = vi.fn();
+    const mockGetAccessToken = vi
+      .fn()
+      .mockResolvedValue({ token: 'mock-access-token' });
     const mockOAuth2Client = {
       generateAuthUrl: mockGenerateAuthUrl,
       getToken: mockGetToken,
       setCredentials: mockSetCredentials,
+      getAccessToken: mockGetAccessToken,
       credentials: mockTokens,
     } as unknown as OAuth2Client;
     vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
@@ -63,6 +70,12 @@ describe('oauth2', () => {
     vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
     vi.mocked(open).mockImplementation(async () => ({}) as never);
 
+    // Mock the UserInfo API response
+    vi.mocked(global.fetch).mockResolvedValue({
+      ok: true,
+      json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }),
+    } as unknown as Response);
+
     let requestCallback!: http.RequestListener<
       typeof http.IncomingMessage,
       typeof http.ServerResponse
@@ -126,5 +139,18 @@ describe('oauth2', () => {
     const tokenPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json');
     const tokenData = JSON.parse(fs.readFileSync(tokenPath, 'utf-8'));
     expect(tokenData).toEqual(mockTokens);
+
+    // Verify Google Account ID was cached
+    const googleAccountIdPath = path.join(
+      tempHomeDir,
+      '.gemini',
+      'google_account_id',
+    );
+    expect(fs.existsSync(googleAccountIdPath)).toBe(true);
+    const cachedGoogleAccountId = fs.readFileSync(googleAccountIdPath, 'utf-8');
+    expect(cachedGoogleAccountId).toBe('test-google-account-id-123');
+
+    // Verify the getCachedGoogleAccountId function works
+    expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123');
   });
 });
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index 68f6e137..d07c8560 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -41,6 +41,7 @@ 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
@@ -60,6 +61,21 @@ export async function getOauthClient(): Promise<OAuth2Client> {
 
   if (await loadCachedCredentials(client)) {
     // Found valid cached credentials.
+    // Check if we need to retrieve Google Account ID
+    if (!getCachedGoogleAccountId()) {
+      try {
+        const googleAccountId = await getGoogleAccountId(client);
+        if (googleAccountId) {
+          await cacheGoogleAccountId(googleAccountId);
+        }
+      } catch (error) {
+        console.error(
+          'Failed to retrieve Google Account ID for existing credentials:',
+          error,
+        );
+        // Continue with existing auth flow
+      }
+    }
     return client;
   }
 
@@ -116,6 +132,20 @@ async function authWithWeb(client: OAuth2Client): Promise<OauthWebLogin> {
           client.setCredentials(tokens);
           await cacheCredentials(client.credentials);
 
+          // Retrieve and cache Google Account ID during authentication
+          try {
+            const googleAccountId = await getGoogleAccountId(client);
+            if (googleAccountId) {
+              await cacheGoogleAccountId(googleAccountId);
+            }
+          } catch (error) {
+            console.error(
+              'Failed to retrieve Google Account ID during authentication:',
+              error,
+            );
+            // Don't fail the auth flow if Google Account ID retrieval fails
+          }
+
           res.writeHead(HTTP_REDIRECT, { Location: SIGN_IN_SUCCESS_URL });
           res.end();
           resolve();
@@ -193,10 +223,76 @@ 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();
+    // eslint-disable-next-line @typescript-eslint/no-require-imports, no-restricted-syntax
+    const fs_sync = require('fs');
+    if (fs_sync.existsSync(filePath)) {
+      return fs_sync.readFileSync(filePath, 'utf-8').trim() || null;
+    }
+    return null;
+  } catch (_error) {
+    return null;
+  }
+}
+
 export async function clearCachedCredentialFile() {
   try {
-    await fs.rm(getCachedCredentialPath());
+    await fs.rm(getCachedCredentialPath(), { force: true });
+    // Clear the Google Account ID cache when credentials are cleared
+    await fs.rm(getGoogleAccountIdCachePath(), { force: true });
   } 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 getGoogleAccountId(
+  client: OAuth2Client,
+): Promise<string | null> {
+  try {
+    const { token } = await client.getAccessToken();
+    if (!token) {
+      return null;
+    }
+
+    const response = await fetch(
+      'https://www.googleapis.com/oauth2/v2/userinfo',
+      {
+        headers: {
+          Authorization: `Bearer ${token}`,
+        },
+      },
+    );
+
+    if (!response.ok) {
+      console.error(
+        'Failed to fetch user info:',
+        response.status,
+        response.statusText,
+      );
+      return null;
+    }
+
+    const userInfo = await response.json();
+    return userInfo.id || null;
+  } catch (error) {
+    console.error('Error retrieving Google Account ID:', error);
+    return null;
+  }
+}