From bc60257e220a77c2d6e57ea4bfd1f0a483a1344c Mon Sep 17 00:00:00 2001
From: Marat Boshernitsan <maratb@google.com>
Date: Fri, 15 Aug 2025 22:05:59 -0700
Subject: [PATCH] feat(oauth): Make oauth client a singleton to survive cache
 failures (#6348)

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

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index 32c8ad3c..77a8fb90 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, Mock } from 'vitest';
-import { getOauthClient } from './oauth2.js';
+import { getOauthClient, resetOauthClientForTesting } from './oauth2.js';
 import { getCachedGoogleAccount } from '../utils/user_account.js';
 import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
@@ -56,6 +56,7 @@ describe('oauth2', () => {
   afterEach(() => {
     fs.rmSync(tempHomeDir, { recursive: true, force: true });
     vi.clearAllMocks();
+    resetOauthClientForTesting();
     delete process.env.CLOUD_SHELL;
     delete process.env.GOOGLE_GENAI_USE_GCA;
     delete process.env.GOOGLE_CLOUD_ACCESS_TOKEN;
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index dc1ec490..f9518cbe 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -66,7 +66,9 @@ export interface OauthWebLogin {
   loginCompletePromise: Promise<void>;
 }
 
-export async function getOauthClient(
+const oauthClientPromises = new Map<AuthType, Promise<OAuth2Client>>();
+
+async function initOauthClient(
   authType: AuthType,
   config: Config,
 ): Promise<OAuth2Client> {
@@ -187,6 +189,16 @@ export async function getOauthClient(
   return client;
 }
 
+export async function getOauthClient(
+  authType: AuthType,
+  config: Config,
+): Promise<OAuth2Client> {
+  if (!oauthClientPromises.has(authType)) {
+    oauthClientPromises.set(authType, initOauthClient(authType, config));
+  }
+  return oauthClientPromises.get(authType)!;
+}
+
 async function authWithUserCode(client: OAuth2Client): Promise<boolean> {
   const redirectUri = 'https://codeassist.google.com/authcode';
   const codeVerifier = await client.generateCodeVerifierAsync();
@@ -416,3 +428,8 @@ async function fetchAndCacheUserInfo(client: OAuth2Client): Promise<void> {
     console.error('Error retrieving user info:', error);
   }
 }
+
+// Helper to ensure test isolation
+export function resetOauthClientForTesting() {
+  oauthClientPromises.clear();
+}