diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts index 4883789659..b032dc28ec 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.test.ts @@ -525,6 +525,102 @@ Your admin might have disabled the access. Contact them to enable the Preview Re expect(result.current.proQuotaRequest).toBeNull(); }); + it('should handle ModelNotFoundError with Vertex AI by displaying region-specific availability message and documentation link', async () => { + vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({ + authType: AuthType.USE_VERTEX_AI, + }); + vi.stubEnv('GOOGLE_CLOUD_LOCATION', 'us-central1'); + + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + + let promise: Promise; + const error = new ModelNotFoundError('model not found', 404); + + act(() => { + promise = handler('gemini-3.5-flash', 'gemini-1.5-flash', error); + }); + + const request = result.current.proQuotaRequest; + expect(request).not.toBeNull(); + expect(request?.failedModel).toBe('gemini-3.5-flash'); + expect(request?.isModelNotFoundError).toBe(true); + + const message = request!.message; + expect(message).toBe( + `Model "gemini-3.5-flash" is not available in region "us-central1".\n` + + `To see which models are available in this region, please visit:\n` + + `https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations\n` + + `/model to switch models.`, + ); + + act(() => { + result.current.handleProQuotaChoice('retry_always'); + }); + + const intent = await promise!; + expect(intent).toBe('retry_always'); + }); + + it('should handle ModelNotFoundError with Vertex AI and invalid model by displaying generic not found error message', async () => { + vi.spyOn(mockConfig, 'getContentGeneratorConfig').mockReturnValue({ + authType: AuthType.USE_VERTEX_AI, + }); + vi.stubEnv('GOOGLE_CLOUD_LOCATION', 'us-central1'); + + const { result } = await renderHook(() => + useQuotaAndFallback({ + config: mockConfig, + historyManager: mockHistoryManager, + userTier: UserTierId.FREE, + setModelSwitchedFromQuotaError: mockSetModelSwitchedFromQuotaError, + onShowAuthSelection: mockOnShowAuthSelection, + paidTier: null, + settings: mockSettings, + }), + ); + + const handler = setFallbackHandlerSpy.mock + .calls[0][0] as FallbackModelHandler; + + let promise: Promise; + const error = new ModelNotFoundError('model not found', 404); + + act(() => { + promise = handler('invalid-model-name', 'gemini-1.5-flash', error); + }); + + const request = result.current.proQuotaRequest; + expect(request).not.toBeNull(); + expect(request?.failedModel).toBe('invalid-model-name'); + expect(request?.isModelNotFoundError).toBe(true); + + const message = request!.message; + expect(message).toBe( + `Model "invalid-model-name" was not found or is invalid.\n` + + `/model to switch models.`, + ); + + act(() => { + result.current.handleProQuotaChoice('retry_always'); + }); + + const intent = await promise!; + expect(intent).toBe('retry_always'); + }); + it('should handle ModelNotFoundError with invalid model correctly', async () => { const { result } = await renderHook(() => useQuotaAndFallback({ diff --git a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts index 533eefa676..a8e757cca4 100644 --- a/packages/cli/src/ui/hooks/useQuotaAndFallback.ts +++ b/packages/cli/src/ui/hooks/useQuotaAndFallback.ts @@ -135,7 +135,20 @@ export function useQuotaAndFallback({ message = messageLines.join('\n'); } else if (error instanceof ModelNotFoundError) { isModelNotFoundError = true; - if (VALID_GEMINI_MODELS.has(failedModel)) { + if ( + contentGeneratorConfig?.authType === AuthType.USE_VERTEX_AI && + VALID_GEMINI_MODELS.has(failedModel) + ) { + const location = + process.env['GOOGLE_CLOUD_LOCATION'] || 'your configured region'; + const messageLines = [ + `Model "${failedModel}" is not available in region "${location}".`, + `To see which models are available in this region, please visit:`, + `https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations`, + `/model to switch models.`, + ]; + message = messageLines.join('\n'); + } else if (VALID_GEMINI_MODELS.has(failedModel)) { const messageLines = [ `It seems like you don't have access to ${getDisplayString(failedModel)}.`, `Your admin might have disabled the access. Contact them to enable the Preview Release Channel.`, diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 3fef648d87..50de05b89a 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -93,6 +93,7 @@ describe('createContentGenerator', () => { resetVersionCache(); vi.clearAllMocks(); vi.stubEnv('ANTIGRAVITY_CLI_ALIAS', ''); + vi.stubEnv('GOOGLE_CLOUD_LOCATION', ''); }); afterEach(() => { @@ -483,6 +484,82 @@ describe('createContentGenerator', () => { ); }); + it('should use US REP endpoint for Vertex AI when location is us and no baseUrl is provided', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getClientName: vi.fn().mockReturnValue(undefined), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + + vi.stubEnv('GOOGLE_CLOUD_LOCATION', 'us'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + vertexai: true, + authType: AuthType.USE_VERTEX_AI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + googleAuthOptions: expect.objectContaining({ + clientOptions: expect.objectContaining({ + apiEndpoint: 'https://aiplatform.us.rep.googleapis.com', + }), + }), + httpOptions: expect.objectContaining({ + baseUrl: 'https://aiplatform.us.rep.googleapis.com', + }), + }), + ); + }); + + it('should use EU REP endpoint for Vertex AI when location is eu and no baseUrl is provided', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getClientName: vi.fn().mockReturnValue(undefined), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + + vi.stubEnv('GOOGLE_CLOUD_LOCATION', 'eu'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + vertexai: true, + authType: AuthType.USE_VERTEX_AI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.objectContaining({ + googleAuthOptions: expect.objectContaining({ + clientOptions: expect.objectContaining({ + apiEndpoint: 'https://aiplatform.eu.rep.googleapis.com', + }), + }), + httpOptions: expect.objectContaining({ + baseUrl: 'https://aiplatform.eu.rep.googleapis.com', + }), + }), + ); + }); + it('should inject HttpsProxyAgent into googleAuthOptions when proxy URL uses https://', async () => { const mockConfigWithProxy = { getModel: vi.fn().mockReturnValue('gemini-pro'), diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index cf7cdcd37f..fb4ba249a1 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -121,6 +121,15 @@ const VERTEX_AI_REQUEST_TYPE_HEADER = 'X-Vertex-AI-LLM-Request-Type'; const VERTEX_AI_SHARED_REQUEST_TYPE_HEADER = 'X-Vertex-AI-LLM-Shared-Request-Type'; +/** + * Vertex AI Representative Endpoints (REP) for US and EU multi-regions. + * These are used as a workaround for the client dynamically + * constructing default legacy hostnames (e.g., 'us-aiplatform.googleapis.com') + * instead of routing to the official REP endpoints. + */ +const VERTEX_AI_US_REP_ENDPOINT = 'https://aiplatform.us.rep.googleapis.com'; +const VERTEX_AI_EU_REP_ENDPOINT = 'https://aiplatform.eu.rep.googleapis.com'; + function validateBaseUrl(baseUrl: string): void { try { new URL(baseUrl); @@ -341,6 +350,13 @@ export async function createContentGenerator( if (envBaseUrl) { validateBaseUrl(envBaseUrl); baseUrl = envBaseUrl; + } else if (config.authType === AuthType.USE_VERTEX_AI) { + const location = process.env['GOOGLE_CLOUD_LOCATION']; + if (location === 'us') { + baseUrl = VERTEX_AI_US_REP_ENDPOINT; + } else if (location === 'eu') { + baseUrl = VERTEX_AI_EU_REP_ENDPOINT; + } } } else { validateBaseUrl(baseUrl);