From 3453b977b8d200ebb5195cc2db38d89dbc2a7323 Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Wed, 18 Jun 2025 09:49:13 -0700 Subject: Support logging in with Application Default Credentials (#1157) Co-authored-by: N. Taylor Mullen --- packages/core/src/code_assist/codeAssist.ts | 20 ++++++++++++++------ packages/core/src/code_assist/oauth2.ts | 24 ++++++++++-------------- packages/core/src/code_assist/server.ts | 4 ++-- packages/core/src/code_assist/setup.ts | 26 ++++++++++---------------- 4 files changed, 36 insertions(+), 38 deletions(-) (limited to 'packages/core/src') diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 5df1502b..92b53104 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/packages/core/src/code_assist/codeAssist.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import { GoogleAuth, AuthClient } from 'google-auth-library'; import { ContentGenerator } from '../core/contentGenerator.js'; import { getOauthClient } from './oauth2.js'; import { setupUser } from './setup.js'; @@ -12,10 +13,17 @@ import { CodeAssistServer, HttpOptions } from './server.js'; export async function createCodeAssistContentGenerator( httpOptions: HttpOptions, ): Promise { - const oauth2Client = await getOauthClient(); - const projectId = await setupUser( - oauth2Client, - process.env.GOOGLE_CLOUD_PROJECT, - ); - return new CodeAssistServer(oauth2Client, projectId, httpOptions); + const authClient = await getAuthClient(); + const projectId = await setupUser(authClient); + return new CodeAssistServer(authClient, projectId, httpOptions); +} + +async function getAuthClient(): Promise { + try { + // Try for Application Default Credentials. + return await new GoogleAuth().getClient(); + } catch (_) { + // No Application Default Credentials so try Oauth. + return await getOauthClient(); + } } diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 84c72fca..9e15f65b 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { OAuth2Client } from 'google-auth-library'; +import { OAuth2Client, Credentials } from 'google-auth-library'; import * as http from 'http'; import url from 'url'; import crypto from 'crypto'; @@ -42,23 +42,12 @@ const SIGN_IN_FAILURE_URL = const GEMINI_DIR = '.gemini'; const CREDENTIAL_FILENAME = 'oauth_creds.json'; -export async function clearCachedCredentials(): Promise { - await fs.rm(getCachedCredentialPath()); -} - export async function getOauthClient(): Promise { try { return await getCachedCredentialClient(); } catch (_) { const loggedInClient = await webLoginClient(); - - await fs.mkdir(path.dirname(getCachedCredentialPath()), { - recursive: true, - }); - await fs.writeFile( - getCachedCredentialPath(), - JSON.stringify(loggedInClient.credentials, null, 2), - ); + await setCachedCredentials(loggedInClient.credentials); return loggedInClient; } } @@ -149,7 +138,6 @@ function getAvailablePort(): Promise { async function getCachedCredentialClient(): Promise { try { const creds = await fs.readFile(getCachedCredentialPath(), 'utf-8'); - const oAuth2Client = new OAuth2Client({ clientId: OAUTH_CLIENT_ID, clientSecret: OAUTH_CLIENT_SECRET, @@ -165,6 +153,14 @@ async function getCachedCredentialClient(): Promise { } } +async function setCachedCredentials(credentials: Credentials) { + const filePath = getCachedCredentialPath(); + await fs.mkdir(path.dirname(filePath), { recursive: true }); + + const credString = JSON.stringify(credentials, null, 2); + await fs.writeFile(filePath, credString); +} + function getCachedCredentialPath(): string { return path.join(os.homedir(), GEMINI_DIR, CREDENTIAL_FILENAME); } diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 2c5d6db9..d700353c 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { OAuth2Client } from 'google-auth-library'; +import { AuthClient } from 'google-auth-library'; import { LoadCodeAssistResponse, LoadCodeAssistRequest, @@ -42,7 +42,7 @@ export const CODE_ASSIST_API_VERSION = 'v1internal'; export class CodeAssistServer implements ContentGenerator { constructor( - readonly auth: OAuth2Client, + readonly auth: AuthClient, readonly projectId?: string, readonly httpOptions: HttpOptions = {}, ) {} diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts index c2f0ef87..cd1b7ffe 100644 --- a/packages/core/src/code_assist/setup.ts +++ b/packages/core/src/code_assist/setup.ts @@ -6,31 +6,27 @@ import { ClientMetadata, OnboardUserRequest } from './types.js'; import { CodeAssistServer } from './server.js'; -import { OAuth2Client } from 'google-auth-library'; -import { clearCachedCredentials } from './oauth2.js'; +import { AuthClient } from 'google-auth-library'; /** * * @param projectId the user's project id, if any * @returns the user's actual project id */ -export async function setupUser( - oAuth2Client: OAuth2Client, - projectId?: string, -): Promise { - const caServer = new CodeAssistServer(oAuth2Client, projectId); +export async function setupUser(authClient: AuthClient): Promise { + const projectId = process.env.GOOGLE_CLOUD_PROJECT; + const caServer = new CodeAssistServer(authClient, projectId); + const clientMetadata: ClientMetadata = { ideType: 'IDE_UNSPECIFIED', platform: 'PLATFORM_UNSPECIFIED', pluginType: 'GEMINI', + duetProject: projectId, }; - if (process.env.GOOGLE_CLOUD_PROJECT) { - clientMetadata.duetProject = process.env.GOOGLE_CLOUD_PROJECT; - } // TODO: Support Free Tier user without projectId. const loadRes = await caServer.loadCodeAssist({ - cloudaicompanionProject: process.env.GOOGLE_CLOUD_PROJECT, + cloudaicompanionProject: projectId, metadata: clientMetadata, }); @@ -39,7 +35,7 @@ export async function setupUser( const onboardReq: OnboardUserRequest = { tierId: onboardTier, - cloudaicompanionProject: loadRes.cloudaicompanionProject || '', + cloudaicompanionProject: loadRes.cloudaicompanionProject || projectId || '', metadata: clientMetadata, }; try { @@ -49,14 +45,12 @@ export async function setupUser( await new Promise((f) => setTimeout(f, 5000)); lroRes = await caServer.onboardUser(onboardReq); } - return lroRes.response?.cloudaicompanionProject?.id || ''; } catch (e) { - await clearCachedCredentials(); console.log( '\n\nError onboarding with Code Assist.\n' + - 'Enterprise users must specify GOOGLE_CLOUD_PROJECT ' + - 'in their environment variables or .env file.\n\n', + 'Google Workspace Account (e.g. your-name@your-company.com)' + + ' must specify a GOOGLE_CLOUD_PROJECT environment variable.\n\n', ); throw e; } -- cgit v1.2.3