Improve LoadCodeAssist error handling (#1645)

1 files changed, 49 insertions, 21 deletions

diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts
index c49a5efa..f0ea60b3 100644
--- a/packages/core/src/code_assist/setup.ts
+++ b/packages/core/src/code_assist/setup.ts
@@ -4,17 +4,31 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ClientMetadata, OnboardUserRequest } from './types.js';
+import {
+  ClientMetadata,
+  GeminiUserTier,
+  LoadCodeAssistResponse,
+  OnboardUserRequest,
+  UserTierId,
+} from './types.js';
 import { CodeAssistServer } from './server.js';
 import { OAuth2Client } from 'google-auth-library';
 
+export class ProjectIdRequiredError extends Error {
+  constructor() {
+    super(
+      'This account requires setting the GOOGLE_CLOUD_PROJECT env var. See https://goo.gle/gemini-cli-auth-docs#workspace-gca',
+    );
+  }
+}
+
 /**
  *
  * @param projectId the user's project id, if any
  * @returns the user's actual project id
  */
 export async function setupUser(authClient: OAuth2Client): Promise<string> {
-  const projectId = process.env.GOOGLE_CLOUD_PROJECT;
+  let projectId = process.env.GOOGLE_CLOUD_PROJECT;
   const caServer = new CodeAssistServer(authClient, projectId);
 
   const clientMetadata: ClientMetadata = {
@@ -24,34 +38,48 @@ export async function setupUser(authClient: OAuth2Client): Promise<string> {
     duetProject: projectId,
   };
 
-  // TODO: Support Free Tier user without projectId.
   const loadRes = await caServer.loadCodeAssist({
     cloudaicompanionProject: projectId,
     metadata: clientMetadata,
   });
 
-  const onboardTier: string =
-    loadRes.allowedTiers?.find((tier) => tier.isDefault)?.id ?? 'legacy-tier';
+  if (!projectId && loadRes.cloudaicompanionProject) {
+    projectId = loadRes.cloudaicompanionProject;
+  }
+
+  const tier = getOnboardTier(loadRes);
+  if (tier.userDefinedCloudaicompanionProject && !projectId) {
+    throw new ProjectIdRequiredError();
+  }
 
   const onboardReq: OnboardUserRequest = {
-    tierId: onboardTier,
-    cloudaicompanionProject: loadRes.cloudaicompanionProject || projectId || '',
+    tierId: tier.id,
+    cloudaicompanionProject: projectId,
     metadata: clientMetadata,
   };
-  try {
-    // Poll onboardUser until long running operation is complete.
-    let lroRes = await caServer.onboardUser(onboardReq);
-    while (!lroRes.done) {
-      await new Promise((f) => setTimeout(f, 5000));
-      lroRes = await caServer.onboardUser(onboardReq);
+
+  // Poll onboardUser until long running operation is complete.
+  let lroRes = await caServer.onboardUser(onboardReq);
+  while (!lroRes.done) {
+    await new Promise((f) => setTimeout(f, 5000));
+    lroRes = await caServer.onboardUser(onboardReq);
+  }
+  return lroRes.response?.cloudaicompanionProject?.id || '';
+}
+
+function getOnboardTier(res: LoadCodeAssistResponse): GeminiUserTier {
+  if (res.currentTier) {
+    return res.currentTier;
+  }
+  for (const tier of res.allowedTiers || []) {
+    if (tier.isDefault) {
+      return tier;
     }
-    return lroRes.response?.cloudaicompanionProject?.id || '';
-  } catch (e) {
-    console.log(
-      '\n\nError onboarding with Code Assist.\n' +
-        'Google Workspace Account (e.g. [email protected])' +
-        ' must specify a GOOGLE_CLOUD_PROJECT environment variable.\n\n',
-    );
-    throw e;
   }
+  return {
+    name: '',
+    description: '',
+    id: UserTierId.LEGACY,
+    userDefinedCloudaicompanionProject: true,
+  };
 }