feat(oauth): Make oauth client a singleton to survive cache failures (#6348)

This commit is contained in:
Marat Boshernitsan 2025-08-15 22:05:59 -07:00 committed by GitHub
parent 6c1373c332
commit bc60257e22
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 2 deletions

View File

@ -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;

View File

@ -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();
}