Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 134 additions & 2 deletions packages/core/src/code_assist/setup.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -332,9 +332,86 @@ describe('setupUser validation', () => {
vi.unstubAllEnvs();
});

it('should throw error if LoadCodeAssist returns ineligible tiers and no current tier', async () => {
it('should throw ineligible tier error when currentTier exists but no project ID available', async () => {
vi.stubEnv('GOOGLE_CLOUD_PROJECT', '');
mockLoad.mockResolvedValue({
currentTier: mockPaidTier,
cloudaicompanionProject: undefined,
ineligibleTiers: [
{
reasonMessage: 'User is not eligible',
reasonCode: 'INELIGIBLE_ACCOUNT',
tierId: 'free-tier',
tierName: 'free',
},
],
});

await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
'User is not eligible',
);
});

it('should continue if LoadCodeAssist returns ineligible tiers but has allowed tiers', async () => {
const mockOnboardUser = vi.fn().mockResolvedValue({
done: true,
response: {
cloudaicompanionProject: {
id: 'server-project',
},
},
});
vi.mocked(CodeAssistServer).mockImplementation(
() =>
({
loadCodeAssist: mockLoad,
onboardUser: mockOnboardUser,
}) as unknown as CodeAssistServer,
);

mockLoad.mockResolvedValue({
currentTier: null,
allowedTiers: [mockPaidTier],
ineligibleTiers: [
{
reasonMessage: 'Not eligible for free tier',
reasonCode: 'INELIGIBLE_ACCOUNT',
tierId: 'free-tier',
tierName: 'free',
},
],
});

// Should not throw - should proceed to onboarding with the allowed tier
const result = await setupUser({} as OAuth2Client);
expect(result).toEqual({
projectId: 'server-project',
userTier: 'standard-tier',
userTierName: 'paid',
});
expect(mockOnboardUser).toHaveBeenCalled();
});

it('should proceed to onboarding with LEGACY tier when no currentTier and no allowedTiers', async () => {
const mockOnboardUser = vi.fn().mockResolvedValue({
done: true,
response: {
cloudaicompanionProject: {
id: 'server-project',
},
},
});
vi.mocked(CodeAssistServer).mockImplementation(
() =>
({
loadCodeAssist: mockLoad,
onboardUser: mockOnboardUser,
}) as unknown as CodeAssistServer,
);

mockLoad.mockResolvedValue({
currentTier: null,
allowedTiers: undefined,
ineligibleTiers: [
{
reasonMessage: 'User is not eligible',
Expand All @@ -345,8 +422,63 @@ describe('setupUser validation', () => {
],
});

// Should proceed to onboarding with LEGACY tier, ignoring ineligible tier errors
const result = await setupUser({} as OAuth2Client);
expect(result).toEqual({
projectId: 'server-project',
userTier: 'legacy-tier',
userTierName: '',
});
expect(mockOnboardUser).toHaveBeenCalledWith(
expect.objectContaining({
tierId: 'legacy-tier',
}),
);
});

it('should throw ValidationRequiredError even if allowed tiers exist', async () => {
mockLoad.mockResolvedValue({
currentTier: null,
allowedTiers: [mockPaidTier],
ineligibleTiers: [
{
reasonMessage: 'Please verify your account',
reasonCode: 'VALIDATION_REQUIRED',
tierId: 'free-tier',
tierName: 'free',
validationUrl: 'https://example.com/verify',
},
],
});

await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
ValidationRequiredError,
);
});

it('should combine multiple ineligible tier messages when currentTier exists but no project ID', async () => {
vi.stubEnv('GOOGLE_CLOUD_PROJECT', '');
mockLoad.mockResolvedValue({
currentTier: mockPaidTier,
cloudaicompanionProject: undefined,
ineligibleTiers: [
{
reasonMessage: 'Not eligible for standard',
reasonCode: 'INELIGIBLE_ACCOUNT',
tierId: 'standard-tier',
tierName: 'standard',
},
{
reasonMessage: 'Not eligible for free',
reasonCode: 'INELIGIBLE_ACCOUNT',
tierId: 'free-tier',
tierName: 'free',
},
],
});

await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
'User is not eligible',
'Not eligible for standard, Not eligible for free',
);
});

Expand Down
34 changes: 27 additions & 7 deletions packages/core/src/code_assist/setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,26 @@ export interface UserData {
}

/**
* Sets up the user by loading their Code Assist configuration and onboarding if needed.
*
* @param projectId the user's project id, if any
* @returns the user's actual project id
* Tier eligibility:
* - FREE tier: Eligibility is determined by the Code Assist server response.
* - STANDARD tier: User is always eligible if they have a valid project ID.
*
* If no valid project ID is available (from env var or server response):
* - Surfaces ineligibility reasons for the FREE tier from the server.
* - Throws ProjectIdRequiredError if no ineligibility reasons are available.
*
* Handles VALIDATION_REQUIRED via the optional validation handler, allowing
* retry, auth change, or cancellation.
*
* @param client - The authenticated client to use for API calls
* @param validationHandler - Optional handler for account validation flow
* @returns The user's project ID, tier ID, and tier name
* @throws {ValidationRequiredError} If account validation is required
* @throws {ProjectIdRequiredError} If no project ID is available and required
* @throws {ValidationCancelledError} If user cancels validation
* @throws {ChangeAuthRequestedError} If user requests to change auth method
*/
export async function setupUser(
client: AuthClient,
Expand Down Expand Up @@ -108,6 +125,14 @@ export async function setupUser(
userTierName: loadRes.currentTier.name,
};
}

// If user is not setup for standard tier, inform them about all other tiers they are ineligible for.
if (loadRes.ineligibleTiers && loadRes.ineligibleTiers.length > 0) {
const reasons = loadRes.ineligibleTiers
.map((t) => t.reasonMessage)
.join(', ');
throw new Error(reasons);
}
throw new ProjectIdRequiredError();
}
return {
Expand Down Expand Up @@ -188,7 +213,6 @@ function validateLoadCodeAssistResponse(res: LoadCodeAssistResponse): void {
res.ineligibleTiers &&
res.ineligibleTiers.length > 0
) {
// Check for VALIDATION_REQUIRED first - this is a recoverable state
const validationTier = res.ineligibleTiers.find(
(t) =>
t.validationUrl &&
Expand All @@ -203,9 +227,5 @@ function validateLoadCodeAssistResponse(res: LoadCodeAssistResponse): void {
validationTier.reasonMessage,
);
}

// For other ineligibility reasons, throw a generic error
const reasons = res.ineligibleTiers.map((t) => t.reasonMessage).join(', ');
throw new Error(reasons);
}
}
Loading