Improve error messages on failed onboarding (#17357)

packages/core/src/code_assist/codeAssist.test.ts

@@ -35,7 +35,10 @@ describe('codeAssist', () => {
 
   describe('createCodeAssistContentGenerator', () => {
     const httpOptions = {};
-    const mockConfig = {} as Config;
+    const mockValidationHandler = vi.fn();
+    const mockConfig = {
+      getValidationHandler: () => mockValidationHandler,
+    } as unknown as Config;
     const mockAuthClient = { a: 'client' };
     const mockUserData = {
       projectId: 'test-project',
@@ -57,7 +60,10 @@ describe('codeAssist', () => {
         AuthType.LOGIN_WITH_GOOGLE,
         mockConfig,
       );
-      expect(setupUser).toHaveBeenCalledWith(mockAuthClient);
+      expect(setupUser).toHaveBeenCalledWith(
+        mockAuthClient,
+        mockValidationHandler,
+      );
       expect(MockedCodeAssistServer).toHaveBeenCalledWith(
         mockAuthClient,
         'test-project',
@@ -83,7 +89,10 @@ describe('codeAssist', () => {
         AuthType.COMPUTE_ADC,
         mockConfig,
       );
-      expect(setupUser).toHaveBeenCalledWith(mockAuthClient);
+      expect(setupUser).toHaveBeenCalledWith(
+        mockAuthClient,
+        mockValidationHandler,
+      );
       expect(MockedCodeAssistServer).toHaveBeenCalledWith(
         mockAuthClient,
         'test-project',

packages/core/src/code_assist/codeAssist.ts

@@ -24,7 +24,7 @@ export async function createCodeAssistContentGenerator(
     authType === AuthType.COMPUTE_ADC
   ) {
     const authClient = await getOauthClient(authType, config);
-    const userData = await setupUser(authClient);
+    const userData = await setupUser(authClient, config.getValidationHandler());
     return new CodeAssistServer(
       authClient,
       userData.projectId,

packages/core/src/code_assist/setup.test.ts

@@ -5,7 +5,13 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { setupUser, ProjectIdRequiredError } from './setup.js';
+import {
+  ProjectIdRequiredError,
+  setupUser,
+  ValidationCancelledError,
+} from './setup.js';
+import { ValidationRequiredError } from '../utils/googleQuotaErrors.js';
+import { ChangeAuthRequestedError } from '../utils/errors.js';
 import { CodeAssistServer } from '../code_assist/server.js';
 import type { OAuth2Client } from 'google-auth-library';
 import type { GeminiUserTier } from './types.js';
@@ -307,3 +313,215 @@ describe('setupUser for new user', () => {
     });
   });
 });
+
+describe('setupUser validation', () => {
+  let mockLoad: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+    mockLoad = vi.fn();
+    vi.mocked(CodeAssistServer).mockImplementation(
+      () =>
+        ({
+          loadCodeAssist: mockLoad,
+        }) as unknown as CodeAssistServer,
+    );
+  });
+
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  it('should throw error if LoadCodeAssist returns ineligible tiers and no current tier', async () => {
+    mockLoad.mockResolvedValue({
+      currentTier: null,
+      ineligibleTiers: [
+        {
+          reasonMessage: 'User is not eligible',
+          reasonCode: 'INELIGIBLE_ACCOUNT',
+          tierId: 'standard-tier',
+          tierName: 'standard',
+        },
+      ],
+    });
+
+    await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
+      'User is not eligible',
+    );
+  });
+
+  it('should retry if validation handler returns verify', async () => {
+    // First call fails
+    mockLoad.mockResolvedValueOnce({
+      currentTier: null,
+      ineligibleTiers: [
+        {
+          reasonMessage: 'User is not eligible',
+          reasonCode: 'VALIDATION_REQUIRED',
+          tierId: 'standard-tier',
+          tierName: 'standard',
+          validationUrl: 'https://example.com/verify',
+          validationLearnMoreUrl: 'https://example.com/learn',
+        },
+      ],
+    });
+    // Second call succeeds
+    mockLoad.mockResolvedValueOnce({
+      currentTier: mockPaidTier,
+      cloudaicompanionProject: 'test-project',
+    });
+
+    const mockValidationHandler = vi.fn().mockResolvedValue('verify');
+
+    const result = await setupUser({} as OAuth2Client, mockValidationHandler);
+
+    expect(mockValidationHandler).toHaveBeenCalledWith(
+      'https://example.com/verify',
+      'User is not eligible',
+    );
+    expect(mockLoad).toHaveBeenCalledTimes(2);
+    expect(result).toEqual({
+      projectId: 'test-project',
+      userTier: 'standard-tier',
+      userTierName: 'paid',
+    });
+  });
+
+  it('should throw if validation handler returns cancel', async () => {
+    mockLoad.mockResolvedValue({
+      currentTier: null,
+      ineligibleTiers: [
+        {
+          reasonMessage: 'User is not eligible',
+          reasonCode: 'VALIDATION_REQUIRED',
+          tierId: 'standard-tier',
+          tierName: 'standard',
+          validationUrl: 'https://example.com/verify',
+        },
+      ],
+    });
+
+    const mockValidationHandler = vi.fn().mockResolvedValue('cancel');
+
+    await expect(
+      setupUser({} as OAuth2Client, mockValidationHandler),
+    ).rejects.toThrow(ValidationCancelledError);
+    expect(mockValidationHandler).toHaveBeenCalled();
+    expect(mockLoad).toHaveBeenCalledTimes(1);
+  });
+
+  it('should throw ChangeAuthRequestedError if validation handler returns change_auth', async () => {
+    mockLoad.mockResolvedValue({
+      currentTier: null,
+      ineligibleTiers: [
+        {
+          reasonMessage: 'User is not eligible',
+          reasonCode: 'VALIDATION_REQUIRED',
+          tierId: 'standard-tier',
+          tierName: 'standard',
+          validationUrl: 'https://example.com/verify',
+        },
+      ],
+    });
+
+    const mockValidationHandler = vi.fn().mockResolvedValue('change_auth');
+
+    await expect(
+      setupUser({} as OAuth2Client, mockValidationHandler),
+    ).rejects.toThrow(ChangeAuthRequestedError);
+    expect(mockValidationHandler).toHaveBeenCalled();
+    expect(mockLoad).toHaveBeenCalledTimes(1);
+  });
+
+  it('should throw ValidationRequiredError without handler', async () => {
+    mockLoad.mockResolvedValue({
+      currentTier: null,
+      ineligibleTiers: [
+        {
+          reasonMessage: 'Please verify your account',
+          reasonCode: 'VALIDATION_REQUIRED',
+          tierId: 'standard-tier',
+          tierName: 'standard',
+          validationUrl: 'https://example.com/verify',
+        },
+      ],
+    });
+
+    await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
+      ValidationRequiredError,
+    );
+    expect(mockLoad).toHaveBeenCalledTimes(1);
+  });
+
+  it('should throw error if LoadCodeAssist returns empty response', async () => {
+    mockLoad.mockResolvedValue(null);
+
+    await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
+      'LoadCodeAssist returned empty response',
+    );
+  });
+
+  it('should retry multiple times when validation handler keeps returning verify', async () => {
+    // First two calls fail with validation required
+    mockLoad
+      .mockResolvedValueOnce({
+        currentTier: null,
+        ineligibleTiers: [
+          {
+            reasonMessage: 'Verify 1',
+            reasonCode: 'VALIDATION_REQUIRED',
+            tierId: 'standard-tier',
+            tierName: 'standard',
+            validationUrl: 'https://example.com/verify',
+          },
+        ],
+      })
+      .mockResolvedValueOnce({
+        currentTier: null,
+        ineligibleTiers: [
+          {
+            reasonMessage: 'Verify 2',
+            reasonCode: 'VALIDATION_REQUIRED',
+            tierId: 'standard-tier',
+            tierName: 'standard',
+            validationUrl: 'https://example.com/verify',
+          },
+        ],
+      })
+      .mockResolvedValueOnce({
+        currentTier: mockPaidTier,
+        cloudaicompanionProject: 'test-project',
+      });
+
+    const mockValidationHandler = vi.fn().mockResolvedValue('verify');
+
+    const result = await setupUser({} as OAuth2Client, mockValidationHandler);
+
+    expect(mockValidationHandler).toHaveBeenCalledTimes(2);
+    expect(mockLoad).toHaveBeenCalledTimes(3);
+    expect(result).toEqual({
+      projectId: 'test-project',
+      userTier: 'standard-tier',
+      userTierName: 'paid',
+    });
+  });
+});
+
+describe('ValidationRequiredError', () => {
+  const error = new ValidationRequiredError(
+    'Account validation required: Please verify',
+    undefined,
+    'https://example.com/verify',
+    'Please verify',
+  );
+
+  it('should be an instance of Error', () => {
+    expect(error).toBeInstanceOf(Error);
+    expect(error).toBeInstanceOf(ValidationRequiredError);
+  });
+
+  it('should have the correct properties', () => {
+    expect(error.validationLink).toBe('https://example.com/verify');
+    expect(error.validationDescription).toBe('Please verify');
+  });
+});

packages/core/src/code_assist/setup.ts

@@ -10,9 +10,12 @@ import type {
   LoadCodeAssistResponse,
   OnboardUserRequest,
 } from './types.js';
-import { UserTierId } from './types.js';
+import { UserTierId, IneligibleTierReasonCode } from './types.js';
 import { CodeAssistServer } from './server.js';
 import type { AuthClient } from 'google-auth-library';
+import type { ValidationHandler } from '../fallback/types.js';
+import { ChangeAuthRequestedError } from '../utils/errors.js';
+import { ValidationRequiredError } from '../utils/googleQuotaErrors.js';
 
 export class ProjectIdRequiredError extends Error {
   constructor() {
@@ -22,6 +25,16 @@ export class ProjectIdRequiredError extends Error {
   }
 }
 
+/**
+ * Error thrown when user cancels the validation process.
+ * This is a non-recoverable error that should result in auth failure.
+ */
+export class ValidationCancelledError extends Error {
+  constructor() {
+    super('User cancelled account validation');
+  }
+}
+
 export interface UserData {
   projectId: string;
   userTier: UserTierId;
@@ -33,7 +46,10 @@ export interface UserData {
  * @param projectId the user's project id, if any
  * @returns the user's actual project id
  */
-export async function setupUser(client: AuthClient): Promise<UserData> {
+export async function setupUser(
+  client: AuthClient,
+  validationHandler?: ValidationHandler,
+): Promise<UserData> {
   const projectId =
     process.env['GOOGLE_CLOUD_PROJECT'] ||
     process.env['GOOGLE_CLOUD_PROJECT_ID'] ||
@@ -52,13 +68,36 @@ export async function setupUser(client: AuthClient): Promise<UserData> {
     pluginType: 'GEMINI',
   };
 
-  const loadRes = await caServer.loadCodeAssist({
-    cloudaicompanionProject: projectId,
-    metadata: {
-      ...coreClientMetadata,
-      duetProject: projectId,
-    },
-  });
+  let loadRes: LoadCodeAssistResponse;
+  while (true) {
+    loadRes = await caServer.loadCodeAssist({
+      cloudaicompanionProject: projectId,
+      metadata: {
+        ...coreClientMetadata,
+        duetProject: projectId,
+      },
+    });
+
+    try {
+      validateLoadCodeAssistResponse(loadRes);
+      break;
+    } catch (e) {
+      if (e instanceof ValidationRequiredError && validationHandler) {
+        const intent = await validationHandler(
+          e.validationLink,
+          e.validationDescription,
+        );
+        if (intent === 'verify') {
+          continue;
+        }
+        if (intent === 'change_auth') {
+          throw new ChangeAuthRequestedError();
+        }
+        throw new ValidationCancelledError();
+      }
+      throw e;
+    }
+  }
 
   if (loadRes.currentTier) {
     if (!loadRes.cloudaicompanionProject) {
@@ -139,3 +178,34 @@ function getOnboardTier(res: LoadCodeAssistResponse): GeminiUserTier {
     userDefinedCloudaicompanionProject: true,
   };
 }
+
+function validateLoadCodeAssistResponse(res: LoadCodeAssistResponse): void {
+  if (!res) {
+    throw new Error('LoadCodeAssist returned empty response');
+  }
+  if (
+    !res.currentTier &&
+    res.ineligibleTiers &&
+    res.ineligibleTiers.length > 0
+  ) {
+    // Check for VALIDATION_REQUIRED first - this is a recoverable state
+    const validationTier = res.ineligibleTiers.find(
+      (t) =>
+        t.validationUrl &&
+        t.reasonCode === IneligibleTierReasonCode.VALIDATION_REQUIRED,
+    );
+    const validationUrl = validationTier?.validationUrl;
+    if (validationTier && validationUrl) {
+      throw new ValidationRequiredError(
+        `Account validation required: ${validationTier.reasonMessage}`,
+        undefined,
+        validationUrl,
+        validationTier.reasonMessage,
+      );
+    }
+
+    // For other ineligibility reasons, throw a generic error
+    const reasons = res.ineligibleTiers.map((t) => t.reasonMessage).join(', ');
+    throw new Error(reasons);
+  }
+}

packages/core/src/code_assist/types.ts

@@ -82,6 +82,11 @@ export interface IneligibleTier {
   reasonMessage: string;
   tierId: UserTierId;
   tierName: string;
+  validationErrorMessage?: string;
+  validationUrl?: string;
+  validationUrlLinkText?: string;
+  validationLearnMoreUrl?: string;
+  validationLearnMoreLinkText?: string;
 }
 
 /**
@@ -98,6 +103,7 @@ export enum IneligibleTierReasonCode {
   UNKNOWN = 'UNKNOWN',
   UNKNOWN_LOCATION = 'UNKNOWN_LOCATION',
   UNSUPPORTED_LOCATION = 'UNSUPPORTED_LOCATION',
+  VALIDATION_REQUIRED = 'VALIDATION_REQUIRED',
   // go/keep-sorted end
 }
 /**