1 files changed, 10 insertions, 16 deletions

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;
   }