Cache oauth credentials (#927)

5 files changed, 83 insertions, 16 deletions

diff --git a/packages/core/src/code_assist/ccpaServer.ts b/packages/core/src/code_assist/ccpaServer.ts
index acfec90f..7a542db4 100644
--- a/packages/core/src/code_assist/ccpaServer.ts
+++ b/packages/core/src/code_assist/ccpaServer.ts
@@ -27,7 +27,7 @@ import { ContentGenerator } from '../core/contentGenerator.js';
 // TODO: Use production endpoint once it supports our methods.
 export const CCPA_ENDPOINT =
   'https://staging-cloudcode-pa.sandbox.googleapis.com';
-export const CCPA_API_VERSION = '/v1internal';
+export const CCPA_API_VERSION = 'v1internal';
 
 export class CcpaServer implements ContentGenerator {
   constructor(
diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index dd5c2ddd..6467b416 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/packages/core/src/code_assist/codeAssist.ts
@@ -5,12 +5,12 @@
  */
 
 import { ContentGenerator } from '../core/contentGenerator.js';
-import { loginWithOauth } from './oauth2.js';
+import { getOauthClient } from './oauth2.js';
 import { setupUser } from './setup.js';
 import { CcpaServer } from './ccpaServer.js';
 
 export async function createCodeAssistContentGenerator(): Promise<ContentGenerator> {
-  const oauth2Client = await loginWithOauth();
+  const oauth2Client = await getOauthClient();
   const projectId = await setupUser(
     oauth2Client,
     process.env.GOOGLE_CLOUD_PROJECT,
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index af87caea..7d65d260 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -10,6 +10,8 @@ import url from 'url';
 import crypto from 'crypto';
 import * as net from 'net';
 import open from 'open';
+import path from 'node:path';
+import { promises as fs } from 'node:fs';
 
 //  OAuth Client ID used to initiate OAuth2Client class.
 const OAUTH_CLIENT_ID =
@@ -36,7 +38,50 @@ const SIGN_IN_SUCCESS_URL =
 const SIGN_IN_FAILURE_URL =
   'https://developers.google.com/gemini-code-assist/auth_failure_gemini';
 
-export async function loginWithOauth(): Promise<OAuth2Client> {
+const GEMINI_DIR = '.gemini';
+const CREDENTIAL_FILENAME = 'oauth_creds.json';
+
+export async function getCachedCredentialClient(): Promise<OAuth2Client> {
+  try {
+    const creds = await fs.readFile(
+      path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME),
+      'utf-8',
+    );
+
+    const oAuth2Client = new OAuth2Client({
+      clientId: OAUTH_CLIENT_ID,
+      clientSecret: OAUTH_CLIENT_SECRET,
+    });
+    oAuth2Client.setCredentials(JSON.parse(creds));
+    // This will either return the existing token or refresh it.
+    await oAuth2Client.getAccessToken();
+    // If we are here, the token is valid.
+    return oAuth2Client;
+  } catch (_) {
+    // Could not load credentials.
+    throw new Error('Could not load credentials');
+  }
+}
+
+export async function clearCachedCredentials(): Promise<void> {
+  await fs.rm(path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME));
+}
+
+export async function getOauthClient(): Promise<OAuth2Client> {
+  try {
+    return await getCachedCredentialClient();
+  } catch (_) {
+    const loggedInClient = await webLoginClient();
+    await fs.mkdir(path.join(process.cwd(), GEMINI_DIR), { recursive: true });
+    await fs.writeFile(
+      path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME),
+      JSON.stringify(loggedInClient.credentials, null, 2),
+    );
+    return loggedInClient;
+  }
+}
+
+export async function webLoginClient(): Promise<OAuth2Client> {
   const port = await getAvailablePort();
   const oAuth2Client = new OAuth2Client({
     clientId: OAUTH_CLIENT_ID,
@@ -51,33 +96,37 @@ export async function loginWithOauth(): Promise<OAuth2Client> {
       scope: OAUTH_SCOPE,
       state,
     });
+    console.log(
+      `\n\nCode Assist login required.\n` +
+        `Attempting to open authentication page in your browser.\n` +
+        `Otherwise navigate to:\n\n${authURL}\n\n`,
+    );
     open(authURL);
+    console.log('Waiting for authentication...');
 
     const server = http.createServer(async (req, res) => {
       try {
         if (req.url!.indexOf('/oauth2callback') === -1) {
-          console.log('Unexpected request:', req.url);
           res.writeHead(HTTP_REDIRECT, { Location: SIGN_IN_FAILURE_URL });
           res.end();
           reject(new Error('Unexpected request: ' + req.url));
         }
         // acquire the code from the querystring, and close the web server.
         const qs = new url.URL(req.url!, 'http://localhost:3000').searchParams;
-        console.log('Processing request:', qs);
         if (qs.get('error')) {
           res.writeHead(HTTP_REDIRECT, { Location: SIGN_IN_FAILURE_URL });
           res.end();
+
           reject(new Error(`Error during authentication: ${qs.get('error')}`));
         } else if (qs.get('state') !== state) {
           res.end('State mismatch. Possible CSRF attack');
+
           reject(new Error('State mismatch. Possible CSRF attack'));
         } else if (qs.get('code')) {
           const code: string = qs.get('code')!;
-          console.log();
           const { tokens } = await oAuth2Client.getToken(code);
-          console.log('Logged in! Tokens:\n\n', tokens);
-
           oAuth2Client.setCredentials(tokens);
+
           res.writeHead(HTTP_REDIRECT, { Location: SIGN_IN_SUCCESS_URL });
           res.end();
           resolve(oAuth2Client);
diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts
index a3162c81..0efe5f0c 100644
--- a/packages/core/src/code_assist/setup.ts
+++ b/packages/core/src/code_assist/setup.ts
@@ -7,6 +7,8 @@
 import { ClientMetadata, OnboardUserRequest } from './types.js';
 import { CcpaServer } from './ccpaServer.js';
 import { OAuth2Client } from 'google-auth-library';
+import { GaxiosError } from 'gaxios';
+import { clearCachedCredentials } from './oauth2.js';
 
 /**
  *
@@ -38,13 +40,27 @@ export async function setupUser(
     cloudaicompanionProject: loadRes.cloudaicompanionProject || '',
     metadata: clientMetadata,
   };
+  try {
+    // Poll onboardUser until long running operation is complete.
+    let lroRes = await ccpaServer.onboardUser(onboardReq);
+    while (!lroRes.done) {
+      await new Promise((f) => setTimeout(f, 5000));
+      lroRes = await ccpaServer.onboardUser(onboardReq);
+    }
 
-  // Poll onboardUser until long running operation is complete.
-  let lroRes = await ccpaServer.onboardUser(onboardReq);
-  while (!lroRes.done) {
-    await new Promise((f) => setTimeout(f, 5000));
-    lroRes = await ccpaServer.onboardUser(onboardReq);
+    return lroRes.response?.cloudaicompanionProject?.id || '';
+  } catch (e) {
+    if (e instanceof GaxiosError) {
+      const detail = e.response?.data?.error?.details[0].detail;
+      if (detail && detail.includes('projectID is empty')) {
+        await clearCachedCredentials();
+        console.log(
+          '\n\nEnterprise users must specify GOOGLE_CLOUD_PROJECT ' +
+            'in your environment variables or .env file.\n\n',
+        );
+        process.exit(1);
+      }
+    }
+    throw e;
   }
-
-  return lroRes.response?.cloudaicompanionProject?.id || '';
 }
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 09ad1e92..99d9efc0 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -18,6 +18,8 @@ export * from './core/geminiRequest.js';
 export * from './core/coreToolScheduler.js';
 export * from './core/nonInteractiveToolExecutor.js';
 
+export * from './code_assist/codeAssist.js';
+
 // Export utilities
 export * from './utils/paths.js';
 export * from './utils/schemaValidator.js';