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 { 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 { getCachedGoogleAccount } from '../utils/user_account.js';
import { OAuth2Client, Compute } from 'google-auth-library'; import { OAuth2Client, Compute } from 'google-auth-library';
import * as fs from 'fs'; import * as fs from 'fs';
@ -56,6 +56,7 @@ describe('oauth2', () => {
afterEach(() => { afterEach(() => {
fs.rmSync(tempHomeDir, { recursive: true, force: true }); fs.rmSync(tempHomeDir, { recursive: true, force: true });
vi.clearAllMocks(); vi.clearAllMocks();
resetOauthClientForTesting();
delete process.env.CLOUD_SHELL; delete process.env.CLOUD_SHELL;
delete process.env.GOOGLE_GENAI_USE_GCA; delete process.env.GOOGLE_GENAI_USE_GCA;
delete process.env.GOOGLE_CLOUD_ACCESS_TOKEN; delete process.env.GOOGLE_CLOUD_ACCESS_TOKEN;

View File

@ -66,7 +66,9 @@ export interface OauthWebLogin {
loginCompletePromise: Promise<void>; loginCompletePromise: Promise<void>;
} }
export async function getOauthClient( const oauthClientPromises = new Map<AuthType, Promise<OAuth2Client>>();
async function initOauthClient(
authType: AuthType, authType: AuthType,
config: Config, config: Config,
): Promise<OAuth2Client> { ): Promise<OAuth2Client> {
@ -187,6 +189,16 @@ export async function getOauthClient(
return client; 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> { async function authWithUserCode(client: OAuth2Client): Promise<boolean> {
const redirectUri = 'https://codeassist.google.com/authcode'; const redirectUri = 'https://codeassist.google.com/authcode';
const codeVerifier = await client.generateCodeVerifierAsync(); const codeVerifier = await client.generateCodeVerifierAsync();
@ -416,3 +428,8 @@ async function fetchAndCacheUserInfo(client: OAuth2Client): Promise<void> {
console.error('Error retrieving user info:', error); console.error('Error retrieving user info:', error);
} }
} }
// Helper to ensure test isolation
export function resetOauthClientForTesting() {
oauthClientPromises.clear();
}