From 24c61147b839b3173fa1ad79781f3c4c0f4144fa Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Wed, 11 Jun 2025 13:26:41 -0700 Subject: [PATCH] Cache oauth credentials (#927) --- packages/cli/src/gemini.tsx | 4 ++ packages/core/src/code_assist/ccpaServer.ts | 2 +- packages/core/src/code_assist/codeAssist.ts | 4 +- packages/core/src/code_assist/oauth2.ts | 61 +++++++++++++++++++-- packages/core/src/code_assist/setup.ts | 30 +++++++--- packages/core/src/index.ts | 2 + 6 files changed, 87 insertions(+), 16 deletions(-) diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 555a7c11..0478a40e 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -80,6 +80,10 @@ export async function main() { } } + // When using Code Assist this triggers the Oauth login. + // Do this now, before sandboxing, so web redirect works. + await config.getGeminiClient().getChat(); + // hop into sandbox if we are outside and sandboxing is enabled if (!process.env.SANDBOX) { const sandbox = sandbox_command(config.getSandbox()); 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 { - 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 { +const GEMINI_DIR = '.gemini'; +const CREDENTIAL_FILENAME = 'oauth_creds.json'; + +export async function getCachedCredentialClient(): Promise { + 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 { + await fs.rm(path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME)); +} + +export async function getOauthClient(): Promise { + 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 { const port = await getAvailablePort(); const oAuth2Client = new OAuth2Client({ clientId: OAUTH_CLIENT_ID, @@ -51,33 +96,37 @@ export async function loginWithOauth(): Promise { 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';