Support logging in with Application Default Credentials (#1157)

Co-authored-by: N. Taylor Mullen <[email protected]>

4 files changed, 36 insertions, 38 deletions

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<ContentGenerator> {
-  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<AuthClient> {
+  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<void> {
-  await fs.rm(getCachedCredentialPath());
-}
-
 export async function getOauthClient(): Promise<OAuth2Client> {
   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<number> {
 async function getCachedCredentialClient(): Promise<OAuth2Client> {
   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<OAuth2Client> {
   }
 }
 
+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<string> {
-  const caServer = new CodeAssistServer(oAuth2Client, projectId);
+export async function setupUser(authClient: AuthClient): Promise<string> {
+  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. [email protected])' +
+        ' must specify a GOOGLE_CLOUD_PROJECT environment variable.\n\n',
     );
     throw e;
   }