From 6505b0c8e16f876ba7ca80b19aa3893dc858ce59 Mon Sep 17 00:00:00 2001
From: Arya Gummadi <aryagummadi@google.com>
Date: Tue, 19 Aug 2025 17:06:25 -0700
Subject: [PATCH] fix: allow re-auth with another google account (#6544)

---
 packages/core/src/code_assist/oauth2.test.ts | 79 +++++++++++++++++++-
 packages/core/src/code_assist/oauth2.ts      | 10 ++-
 2 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index 25718cb1..4334bd36 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -5,7 +5,12 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
-import { getOauthClient, resetOauthClientForTesting } from './oauth2.js';
+import {
+  getOauthClient,
+  resetOauthClientForTesting,
+  clearCachedCredentialFile,
+  clearOauthClientCache,
+} from './oauth2.js';
 import { getCachedGoogleAccount } from '../utils/user_account.js';
 import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
@@ -510,4 +515,76 @@ describe('oauth2', () => {
       expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds);
     });
   });
+
+  describe('clearCachedCredentialFile', () => {
+    it('should clear cached credentials and Google account', async () => {
+      const cachedCreds = { refresh_token: 'test-token' };
+      const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json');
+      await fs.promises.mkdir(path.dirname(credsPath), { recursive: true });
+      await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds));
+
+      const googleAccountPath = path.join(
+        tempHomeDir,
+        '.gemini',
+        'google_accounts.json',
+      );
+      const accountData = { active: 'test@example.com', old: [] };
+      await fs.promises.writeFile(
+        googleAccountPath,
+        JSON.stringify(accountData),
+      );
+
+      expect(fs.existsSync(credsPath)).toBe(true);
+      expect(fs.existsSync(googleAccountPath)).toBe(true);
+      expect(getCachedGoogleAccount()).toBe('test@example.com');
+
+      await clearCachedCredentialFile();
+      expect(fs.existsSync(credsPath)).toBe(false);
+      expect(getCachedGoogleAccount()).toBeNull();
+      const updatedAccountData = JSON.parse(
+        fs.readFileSync(googleAccountPath, 'utf-8'),
+      );
+      expect(updatedAccountData.active).toBeNull();
+      expect(updatedAccountData.old).toContain('test@example.com');
+    });
+
+    it('should clear the in-memory OAuth client cache', async () => {
+      const mockSetCredentials = vi.fn();
+      const mockGetAccessToken = vi
+        .fn()
+        .mockResolvedValue({ token: 'test-token' });
+      const mockGetTokenInfo = vi.fn().mockResolvedValue({});
+      const mockOAuth2Client = {
+        setCredentials: mockSetCredentials,
+        getAccessToken: mockGetAccessToken,
+        getTokenInfo: mockGetTokenInfo,
+        on: vi.fn(),
+      } as unknown as OAuth2Client;
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockOAuth2Client,
+      );
+
+      // Pre-populate credentials to make getOauthClient resolve quickly
+      const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json');
+      await fs.promises.mkdir(path.dirname(credsPath), { recursive: true });
+      await fs.promises.writeFile(
+        credsPath,
+        JSON.stringify({ refresh_token: 'token' }),
+      );
+
+      // First call, should create a client
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+      expect(OAuth2Client).toHaveBeenCalledTimes(1);
+
+      // Second call, should use cached client
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+      expect(OAuth2Client).toHaveBeenCalledTimes(1);
+
+      clearOauthClientCache();
+
+      // Third call, after clearing cache, should create a new client
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+      expect(OAuth2Client).toHaveBeenCalledTimes(2);
+    });
+  });
 });
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index f479dd1b..38238b0b 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -391,13 +391,19 @@ function getCachedCredentialPath(): string {
   return path.join(os.homedir(), GEMINI_DIR, CREDENTIAL_FILENAME);
 }
 
+export function clearOauthClientCache() {
+  oauthClientPromises.clear();
+}
+
 export async function clearCachedCredentialFile() {
   try {
     await fs.rm(getCachedCredentialPath(), { force: true });
     // Clear the Google Account ID cache when credentials are cleared
     await clearCachedGoogleAccount();
-  } catch (_) {
-    /* empty */
+    // Clear the in-memory OAuth client cache to force re-authentication
+    clearOauthClientCache();
+  } catch (e) {
+    console.error('Failed to clear cached credentials:', e);
   }
 }