Vertex base url update (#28145)

DavidAPierce · web-flow · commit b14416447e93 · 2026-06-25T20:40:55.000Z
diff --git 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<FallbackIntent | null>;
+        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<FallbackIntent | null>;
+        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
@@ -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
@@ -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
@@ -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);
