Vertex ai model mapping fix

packages/core/src/config/config.ts

@@ -3564,7 +3564,7 @@ export class Config implements McpContext, AgentLoopContext {
       // Gemini API key users should have the ability to manually select the
       // old preview flash model.
       if (authType === AuthType.USE_GEMINI) {
-        setFlashModels('gemini-3-flash-preview', 'gemini-3.5-flash');
+        setFlashModels('gemini-3-flash-preview', 'gemini-3-flash');
       } else {
         setFlashModels('gemini-3-flash', 'gemini-3-flash');
       }

packages/core/src/config/models.ts

@@ -59,16 +59,16 @@
 // cleaned up.
 export let PREVIEW_GEMINI_FLASH_MODEL = 'gemini-3-flash-preview';
 export const DEFAULT_GEMINI_MODEL = 'gemini-2.5-pro';
 // TODO: Set to const and update to 'gemini-3.5-flash' once the experiment for
 // 3_5 flash rollut can be cleaned up.
 // This is set to either the same as the DEFAULT_GEMINI_3_5_FLASH_MODEL const
 // OR the SECONDARY_GEMINI_3_5_FLASH_MODEL depending on which is needed for
 // the user's backend as determined by hasGemini35FlashGAAccess in
 // packages/core/src/config/config.ts
 export let DEFAULT_GEMINI_FLASH_MODEL = 'gemini-2.5-flash';
 export const DEFAULT_GEMINI_3_5_FLASH_MODEL = 'gemini-3.5-flash';
 // This is resolved to 3.5 flash in backends where it is used,
 // however those backends do not expect to see the string gemini-3.5-flash
 // so we need to provide this model as an alternative name in certain instances.
 export const SECONDARY_GEMINI_3_5_FLASH_MODEL = 'gemini-3-flash';
 
@@ -458,7 +458,7 @@
 export function isGemini2Model(model: string): boolean {
   // This is legacy behavior, will remove this when gemini 2 models are no
   // longer needed.
   return /^gemini-2(\.|$)/.test(model);
 }
 
 /**
@@ -574,3 +574,11 @@
     );
   }
 }
+
+export const VERTEX_AI_MODEL_MAPPINGS: Record<string, string> = {
+  SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL,
+};
+
+export const GEMINI_API_MODEL_MAPPINGS: Record<string, string> = {
+  SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL,
+};

packages/core/src/core/contentGenerator.test.ts

@@ -22,6 +22,9 @@
 import { FakeContentGenerator } from './fakeContentGenerator.js';
 import { RecordingContentGenerator } from './recordingContentGenerator.js';
 import { resetVersionCache } from '../utils/version.js';
+import type { LlmRole } from '../telemetry/llmRole.js';
+import { ModelMappingContentGenerator } from './modelMappingContentGenerator.js';
+import { GEMINI_API_MODEL_MAPPINGS } from '../config/models.js';
 
 vi.mock('../code_assist/codeAssist.js');
 vi.mock('@google/genai');
@@ -36,6 +39,14 @@
   getProxy: vi.fn().mockReturnValue(undefined),
   getUsageStatisticsEnabled: vi.fn().mockReturnValue(true),
   getClientName: vi.fn().mockReturnValue(undefined),
+  getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(true),
+  getTelemetryTracesEnabled: vi.fn().mockReturnValue(true),
+  getSessionId: vi.fn().mockReturnValue('test-session-id'),
+  refreshUserQuotaIfStale: vi.fn().mockResolvedValue(undefined),
+  setLatestApiRequest: vi.fn(),
+  getContentGeneratorConfig: vi.fn().mockReturnValue({}),
+  isInteractive: vi.fn().mockReturnValue(false),
+  getExperiments: vi.fn().mockReturnValue(undefined),
 } as unknown as Config;
 
 describe('getAuthTypeFromEnv', () => {
@@ -201,7 +212,13 @@
       }),
     });
     expect(generator).toEqual(
-      new LoggingContentGenerator(mockGenerator.models, mockConfig),
+      new LoggingContentGenerator(
+        new ModelMappingContentGenerator(
+          mockGenerator.models,
+          GEMINI_API_MODEL_MAPPINGS,
+        ),
+        mockConfig,
+      ),
     );
   });
 
@@ -738,7 +755,13 @@
       }),
     });
     expect(generator).toEqual(
-      new LoggingContentGenerator(mockGenerator.models, mockConfig),
+      new LoggingContentGenerator(
+        new ModelMappingContentGenerator(
+          mockGenerator.models,
+          GEMINI_API_MODEL_MAPPINGS,
+        ),
+        mockConfig,
+      ),
     );
   });
 
@@ -1095,6 +1118,77 @@
       }),
     );
   });
+
+  it('should apply model mapping for Vertex AI', async () => {
+    const mockModels = {
+      generateContent: vi.fn().mockResolvedValue({}),
+    };
+    const mockGenerator = {
+      models: mockModels,
+    } as unknown as GoogleGenAI;
+    vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
+
+    const generator = await createContentGenerator(
+      {
+        apiKey: 'test-api-key',
+        authType: AuthType.USE_VERTEX_AI,
+        vertexai: true,
+      },
+      mockConfig,
+    );
+
+    await generator.generateContent(
+      {
+        model: 'gemini-3-flash',
+        contents: [],
+      },
+      'prompt-id',
+      'user' as LlmRole,
+    );
+
+    expect(mockModels.generateContent).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'gemini-3.5-flash',
+      }),
+      'prompt-id',
+      'user',
+    );
+  });
+
+  it('should apply model mapping for Gemini API', async () => {
+    const mockModels = {
+      generateContent: vi.fn().mockResolvedValue({}),
+    };
+    const mockGenerator = {
+      models: mockModels,
+    } as unknown as GoogleGenAI;
+    vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
+
+    const generator = await createContentGenerator(
+      {
+        apiKey: 'test-api-key',
+        authType: AuthType.USE_GEMINI,
+      },
+      mockConfig,
+    );
+
+    await generator.generateContent(
+      {
+        model: 'gemini-3-flash',
+        contents: [],
+      },
+      'prompt-id',
+      'user' as LlmRole,
+    );
+
+    expect(mockModels.generateContent).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'gemini-3.5-flash',
+      }),
+      'prompt-id',
+      'user',
+    );
+  });
 });
 
 describe('createContentGeneratorConfig', () => {

packages/core/src/core/contentGenerator.ts

@@ -30,6 +30,11 @@ import { determineSurface } from '../utils/surface.js';
 import { RecordingContentGenerator } from './recordingContentGenerator.js';
 import { getVersion, resolveModel } from '../../index.js';
 import type { LlmRole } from '../telemetry/llmRole.js';
+import { ModelMappingContentGenerator } from './modelMappingContentGenerator.js';
+import {
+  VERTEX_AI_MODEL_MAPPINGS,
+  GEMINI_API_MODEL_MAPPINGS,
+} from '../config/models.js';
 
 /**
  * Interface abstracting the core functionalities for generating content and counting tokens.
@@ -375,7 +380,19 @@ export async function createContentGenerator(
           },
         }),
       });
-      return new LoggingContentGenerator(googleGenAI.models, gcConfig);
+      let generator: ContentGenerator = googleGenAI.models;
+      if (config.authType === AuthType.USE_VERTEX_AI) {
+        generator = new ModelMappingContentGenerator(
+          generator,
+          VERTEX_AI_MODEL_MAPPINGS,
+        );
+      } else if (config.authType === AuthType.USE_GEMINI) {
+        generator = new ModelMappingContentGenerator(
+          generator,
+          GEMINI_API_MODEL_MAPPINGS,
+        );
+      }
+      return new LoggingContentGenerator(generator, gcConfig);
     }
     throw new Error(
       `Error creating contentGenerator: Unsupported authType: ${config.authType}`,

packages/core/src/core/modelMappingContentGenerator.ts

@@ -0,0 +1,78 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  type CountTokensResponse,
+  type GenerateContentResponse,
+  type GenerateContentParameters,
+  type CountTokensParameters,
+  type EmbedContentResponse,
+  type EmbedContentParameters,
+} from '@google/genai';
+import { type ContentGenerator } from './contentGenerator.js';
+import type { LlmRole } from '../telemetry/llmRole.js';
+import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
+
+export class ModelMappingContentGenerator implements ContentGenerator {
+  constructor(
+    private readonly wrapped: ContentGenerator,
+    private readonly mappings: Record<string, string>,
+  ) {}
+
+  get userTier(): UserTierId | undefined {
+    return this.wrapped.userTier;
+  }
+
+  get userTierName(): string | undefined {
+    return this.wrapped.userTierName;
+  }
+
+  get paidTier(): GeminiUserTier | undefined {
+    return this.wrapped.paidTier;
+  }
+
+  private mapModel<T extends { model?: string }>(req: T): T {
+    if (req.model && this.mappings[req.model]) {
+      return {
+        ...req,
+        model: this.mappings[req.model],
+      };
+    }
+    return req;
+  }
+
+  generateContent(
+    request: GenerateContentParameters,
+    userPromptId: string,
+    role: LlmRole,
+  ): Promise<GenerateContentResponse> {
+    return this.wrapped.generateContent(
+      this.mapModel(request),
+      userPromptId,
+      role,
+    );
+  }
+
+  generateContentStream(
+    request: GenerateContentParameters,
+    userPromptId: string,
+    role: LlmRole,
+  ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    return this.wrapped.generateContentStream(
+      this.mapModel(request),
+      userPromptId,
+      role,
+    );
+  }
+
+  countTokens(request: CountTokensParameters): Promise<CountTokensResponse> {
+    return this.wrapped.countTokens(this.mapModel(request));
+  }
+
+  embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse> {
+    return this.wrapped.embedContent(this.mapModel(request));
+  }
+}