Refac: Centralize storage file management (#4078)

Co-authored-by: Taylor Mullen <ntaylormullen@google.com>

packages/core/src/code_assist/oauth2.test.ts

@@ -11,7 +11,7 @@ import {
   clearCachedCredentialFile,
   clearOauthClientCache,
 } from './oauth2.js';
-import { getCachedGoogleAccount } from '../utils/user_account.js';
+import { UserAccountManager } from '../utils/userAccountManager.js';
 import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
 import * as path from 'path';
@@ -180,7 +180,10 @@ describe('oauth2', () => {
     });
 
     // Verify the getCachedGoogleAccount function works
-    expect(getCachedGoogleAccount()).toBe('test-google-account@gmail.com');
+    const userAccountManager = new UserAccountManager();
+    expect(userAccountManager.getCachedGoogleAccount()).toBe(
+      'test-google-account@gmail.com',
+    );
   });
 
   it('should perform login with user code', async () => {
@@ -533,14 +536,17 @@ describe('oauth2', () => {
         googleAccountPath,
         JSON.stringify(accountData),
       );
+      const userAccountManager = new UserAccountManager();
 
       expect(fs.existsSync(credsPath)).toBe(true);
       expect(fs.existsSync(googleAccountPath)).toBe(true);
-      expect(getCachedGoogleAccount()).toBe('test@example.com');
+      expect(userAccountManager.getCachedGoogleAccount()).toBe(
+        'test@example.com',
+      );
 
       await clearCachedCredentialFile();
       expect(fs.existsSync(credsPath)).toBe(false);
-      expect(getCachedGoogleAccount()).toBeNull();
+      expect(userAccountManager.getCachedGoogleAccount()).toBeNull();
       const updatedAccountData = JSON.parse(
         fs.readFileSync(googleAccountPath, 'utf-8'),
       );

packages/core/src/code_assist/oauth2.ts

@@ -17,16 +17,14 @@ import * as net from 'net';
 import open from 'open';
 import path from 'node:path';
 import { promises as fs } from 'node:fs';
-import * as os from 'os';
 import { Config } from '../config/config.js';
 import { getErrorMessage } from '../utils/errors.js';
-import {
-  cacheGoogleAccount,
-  getCachedGoogleAccount,
-  clearCachedGoogleAccount,
-} from '../utils/user_account.js';
+import { UserAccountManager } from '../utils/userAccountManager.js';
 import { AuthType } from '../core/contentGenerator.js';
 import readline from 'node:readline';
+import { Storage } from '../config/storage.js';
+
+const userAccountManager = new UserAccountManager();
 
 //  OAuth Client ID used to initiate OAuth2Client class.
 const OAUTH_CLIENT_ID =
@@ -53,9 +51,6 @@ const SIGN_IN_SUCCESS_URL =
 const SIGN_IN_FAILURE_URL =
   'https://developers.google.com/gemini-code-assist/auth_failure_gemini';
 
-const GEMINI_DIR = '.gemini';
-const CREDENTIAL_FILENAME = 'oauth_creds.json';
-
 /**
  * An Authentication URL for updating the credentials of a Oauth2Client
  * as well as a promise that will resolve when the credentials have
@@ -99,7 +94,7 @@ async function initOauthClient(
   if (await loadCachedCredentials(client)) {
     // Found valid cached credentials.
     // Check if we need to retrieve Google Account ID or Email
-    if (!getCachedGoogleAccount()) {
+    if (!userAccountManager.getCachedGoogleAccount()) {
       try {
         await fetchAndCacheUserInfo(client);
       } catch {
@@ -352,7 +347,7 @@ export function getAvailablePort(): Promise<number> {
 
 async function loadCachedCredentials(client: OAuth2Client): Promise<boolean> {
   const pathsToTry = [
-    getCachedCredentialPath(),
+    Storage.getOAuthCredsPath(),
     process.env['GOOGLE_APPLICATION_CREDENTIALS'],
   ].filter((p): p is string => !!p);
 
@@ -380,26 +375,22 @@ async function loadCachedCredentials(client: OAuth2Client): Promise<boolean> {
 }
 
 async function cacheCredentials(credentials: Credentials) {
-  const filePath = getCachedCredentialPath();
+  const filePath = Storage.getOAuthCredsPath();
   await fs.mkdir(path.dirname(filePath), { recursive: true });
 
   const credString = JSON.stringify(credentials, null, 2);
   await fs.writeFile(filePath, credString, { mode: 0o600 });
 }
 
-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 });
+    await fs.rm(Storage.getOAuthCredsPath(), { force: true });
     // Clear the Google Account ID cache when credentials are cleared
-    await clearCachedGoogleAccount();
+    await userAccountManager.clearCachedGoogleAccount();
     // Clear the in-memory OAuth client cache to force re-authentication
     clearOauthClientCache();
   } catch (e) {
@@ -433,9 +424,7 @@ async function fetchAndCacheUserInfo(client: OAuth2Client): Promise<void> {
     }
 
     const userInfo = await response.json();
-    if (userInfo.email) {
-      await cacheGoogleAccount(userInfo.email);
-    }
+    await userAccountManager.cacheGoogleAccount(userInfo.email);
   } catch (error) {
     console.error('Error retrieving user info:', error);
   }