From c55b15f705d083e3dadcfb71494dcb0d6043e6c6 Mon Sep 17 00:00:00 2001
From: Tommaso Sciortino <sciortino@gmail.com>
Date: Thu, 26 Jun 2025 08:27:20 -0700
Subject: Improve LoadCodeAssist error handling (#1645)

---
 packages/cli/src/gemini.tsx                 | 12 +++--
 packages/cli/src/ui/hooks/useAuthCommand.ts | 12 +----
 packages/core/src/code_assist/setup.ts      | 70 ++++++++++++++++++++---------
 3 files changed, 58 insertions(+), 36 deletions(-)

(limited to 'packages')

diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx
index 95c8d4f2..4a0014e1 100644
--- a/packages/cli/src/gemini.tsx
+++ b/packages/cli/src/gemini.tsx
@@ -141,12 +141,16 @@ export async function main() {
     if (sandboxConfig) {
       if (settings.merged.selectedAuthType) {
         // Validate authentication here because the sandbox will interfere with the Oauth2 web redirect.
-        const err = validateAuthMethod(settings.merged.selectedAuthType);
-        if (err) {
-          console.error(err);
+        try {
+          const err = validateAuthMethod(settings.merged.selectedAuthType);
+          if (err) {
+            throw new Error(err);
+          }
+          await config.refreshAuth(settings.merged.selectedAuthType);
+        } catch (err) {
+          console.error('Error authenticating:', err);
           process.exit(1);
         }
-        await config.refreshAuth(settings.merged.selectedAuthType);
       }
       await start_sandbox(sandboxConfig, memoryArgs);
       process.exit(0);
diff --git a/packages/cli/src/ui/hooks/useAuthCommand.ts b/packages/cli/src/ui/hooks/useAuthCommand.ts
index 283572d2..766ebe2b 100644
--- a/packages/cli/src/ui/hooks/useAuthCommand.ts
+++ b/packages/cli/src/ui/hooks/useAuthCommand.ts
@@ -46,17 +46,7 @@ export const useAuthCommand = (
           config,
         );
       } catch (e) {
-        let errorMessage = `Failed to login.\nMessage: ${getErrorMessage(e)}`;
-        if (
-          settings.merged.selectedAuthType ===
-            AuthType.LOGIN_WITH_GOOGLE_PERSONAL &&
-          !process.env.GOOGLE_CLOUD_PROJECT
-        ) {
-          errorMessage =
-            'Failed to login. Workspace accounts and licensed Code Assist users must configure' +
-            ` GOOGLE_CLOUD_PROJECT (see https://goo.gle/gemini-cli-auth-docs#workspace-gca).\nMessage: ${getErrorMessage(e)}`;
-        }
-        setAuthError(errorMessage);
+        setAuthError(`Failed to login. Message: ${getErrorMessage(e)}`);
         openAuthDialog();
       } finally {
         setIsAuthenticating(false);
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. your-name@your-company.com)' +
-        ' must specify a GOOGLE_CLOUD_PROJECT environment variable.\n\n',
-    );
-    throw e;
   }
+  return {
+    name: '',
+    description: '',
+    id: UserTierId.LEGACY,
+    userDefinedCloudaicompanionProject: true,
+  };
 }
-- 
cgit v1.2.3