From 5200c4c5cf4c1b7ad46e1c8b7768f197e1cf65bc Mon Sep 17 00:00:00 2001 From: davidapierce Date: Mon, 8 Jun 2026 21:20:51 +0000 Subject: [PATCH 1/9] add a model mapping that intercepts set models and provides their expected model names to the proper endpoints in contentGenerator.ts --- packages/core/src/config/config.ts | 2 +- packages/core/src/config/models.ts | 10 ++ .../core/src/core/contentGenerator.test.ts | 98 ++++++++++++++++++- packages/core/src/core/contentGenerator.ts | 19 +++- .../src/core/modelMappingContentGenerator.ts | 78 +++++++++++++++ 5 files changed, 203 insertions(+), 4 deletions(-) create mode 100644 packages/core/src/core/modelMappingContentGenerator.ts diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index f7895931d5f..85e93a89da4 100644 --- a/packages/core/src/config/config.ts +++ b/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'); } diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index 1fd8047f02d..1d9e3c25db6 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -574,3 +574,13 @@ export function isActiveModel( ); } } + +export const VERTEX_AI_MODEL_MAPPINGS: Record = { + 'gemini-3-flash': 'gemini-3.5-flash', + 'gemini-3-pro-preview': 'gemini-3.1-pro-preview', +}; + +export const GEMINI_API_MODEL_MAPPINGS: Record = { + 'gemini-3-flash': 'gemini-3.5-flash', + 'gemini-3-pro-preview': 'gemini-3.1-pro-preview', +}; diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 72e9c5b514f..893451d4c9c 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -22,6 +22,9 @@ import { loadApiKey } from './apiKeyCredentialStorage.js'; 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 @@ const mockConfig = { 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 @@ describe('createContentGenerator', () => { }), }); expect(generator).toEqual( - new LoggingContentGenerator(mockGenerator.models, mockConfig), + new LoggingContentGenerator( + new ModelMappingContentGenerator( + mockGenerator.models, + GEMINI_API_MODEL_MAPPINGS, + ), + mockConfig, + ), ); }); @@ -738,7 +755,13 @@ describe('createContentGenerator', () => { }), }); expect(generator).toEqual( - new LoggingContentGenerator(mockGenerator.models, mockConfig), + new LoggingContentGenerator( + new ModelMappingContentGenerator( + mockGenerator.models, + GEMINI_API_MODEL_MAPPINGS, + ), + mockConfig, + ), ); }); @@ -1095,6 +1118,77 @@ describe('createContentGenerator', () => { }), ); }); + + 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', () => { diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 04493c6d73b..a1940efa254 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/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}`, diff --git a/packages/core/src/core/modelMappingContentGenerator.ts b/packages/core/src/core/modelMappingContentGenerator.ts new file mode 100644 index 00000000000..aac110142e7 --- /dev/null +++ b/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, + ) {} + + 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(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 { + return this.wrapped.generateContent( + this.mapModel(request), + userPromptId, + role, + ); + } + + generateContentStream( + request: GenerateContentParameters, + userPromptId: string, + role: LlmRole, + ): Promise> { + return this.wrapped.generateContentStream( + this.mapModel(request), + userPromptId, + role, + ); + } + + countTokens(request: CountTokensParameters): Promise { + return this.wrapped.countTokens(this.mapModel(request)); + } + + embedContent(request: EmbedContentParameters): Promise { + return this.wrapped.embedContent(this.mapModel(request)); + } +} From d257f8e10212fd9a2361cd9d8dfab33caf184289 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Mon, 8 Jun 2026 21:35:43 +0000 Subject: [PATCH 2/9] use existing consts in new mapping for consistency --- packages/core/src/config/models.ts | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index 1d9e3c25db6..2708c0fd4b7 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -576,11 +576,9 @@ export function isActiveModel( } export const VERTEX_AI_MODEL_MAPPINGS: Record = { - 'gemini-3-flash': 'gemini-3.5-flash', - 'gemini-3-pro-preview': 'gemini-3.1-pro-preview', + SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL, }; export const GEMINI_API_MODEL_MAPPINGS: Record = { - 'gemini-3-flash': 'gemini-3.5-flash', - 'gemini-3-pro-preview': 'gemini-3.1-pro-preview', + SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL, }; From 4bc56b7e7ad218b19c44c95cf4a6f7c4e67f87c9 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Mon, 8 Jun 2026 21:59:08 +0000 Subject: [PATCH 3/9] fix mapping so that it doesn't try for literal strings. --- packages/core/src/config/models.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index 2708c0fd4b7..a6639afd088 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -576,9 +576,9 @@ export function isActiveModel( } export const VERTEX_AI_MODEL_MAPPINGS: Record = { - SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL, + [SECONDARY_GEMINI_3_5_FLASH_MODEL]: DEFAULT_GEMINI_3_5_FLASH_MODEL, }; export const GEMINI_API_MODEL_MAPPINGS: Record = { - SECONDARY_GEMINI_3_5_FLASH_MODEL: DEFAULT_GEMINI_3_5_FLASH_MODEL, + [SECONDARY_GEMINI_3_5_FLASH_MODEL]: DEFAULT_GEMINI_3_5_FLASH_MODEL, }; From 3cb79266532cd8f459f7d304ea2fa62ee81b0c93 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Mon, 8 Jun 2026 22:05:12 +0000 Subject: [PATCH 4/9] fix failing test that referenced 3.5 flash now that the override has been moved. --- packages/core/src/config/config.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 48b15253bcf..e48536181f0 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -4357,7 +4357,7 @@ describe('hasGemini35FlashGAAccess model setting', () => { cwd: '.', }; - it('should set DEFAULT_GEMINI_FLASH_MODEL to gemini-3.5-flash and PREVIEW_GEMINI_FLASH_MODEL to gemini-3-flash-preview if hasGemini35FlashGAAccess returns true and authType is USE_GEMINI', () => { + it('should set DEFAULT_GEMINI_FLASH_MODEL to gemini-3-flash and PREVIEW_GEMINI_FLASH_MODEL to gemini-3-flash-preview if hasGemini35FlashGAAccess returns true and authType is USE_GEMINI', () => { const config = new Config(baseParams); config['contentGeneratorConfig'] = { authType: AuthType.USE_GEMINI }; @@ -4375,7 +4375,7 @@ describe('hasGemini35FlashGAAccess model setting', () => { const result = config.hasGemini35FlashGAAccess(); expect(result).toBe(true); - expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3.5-flash'); + expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3-flash'); expect(PREVIEW_GEMINI_FLASH_MODEL).toBe('gemini-3-flash-preview'); }); From b79f67a833017c4733b438ddbcfebe5b07343257 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Tue, 9 Jun 2026 01:45:34 +0000 Subject: [PATCH 5/9] inverted overwriting so that the special case of gemin-3-flash is treated as the exception not the rule. --- packages/core/src/config/config.ts | 4 +- packages/core/src/config/models.ts | 8 +-- .../core/src/core/contentGenerator.test.ts | 71 +++++++++++++------ packages/core/src/core/contentGenerator.ts | 30 ++++---- 4 files changed, 68 insertions(+), 45 deletions(-) diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 85e93a89da4..5196e5cd63d 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -3564,9 +3564,9 @@ 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-flash'); + setFlashModels('gemini-3-flash-preview', 'gemini-3.5-flash'); } else { - setFlashModels('gemini-3-flash', 'gemini-3-flash'); + setFlashModels('gemini-3.5-flash', 'gemini-3.5-flash'); } } else { setFlashModels('gemini-3-flash-preview', 'gemini-2.5-flash'); diff --git a/packages/core/src/config/models.ts b/packages/core/src/config/models.ts index a6639afd088..dd6506b837c 100644 --- a/packages/core/src/config/models.ts +++ b/packages/core/src/config/models.ts @@ -575,10 +575,6 @@ export function isActiveModel( } } -export const VERTEX_AI_MODEL_MAPPINGS: Record = { - [SECONDARY_GEMINI_3_5_FLASH_MODEL]: DEFAULT_GEMINI_3_5_FLASH_MODEL, -}; - -export const GEMINI_API_MODEL_MAPPINGS: Record = { - [SECONDARY_GEMINI_3_5_FLASH_MODEL]: DEFAULT_GEMINI_3_5_FLASH_MODEL, +export const CCPA_AI_MODEL_MAPPINGS: Record = { + [DEFAULT_GEMINI_3_5_FLASH_MODEL]: SECONDARY_GEMINI_3_5_FLASH_MODEL, }; diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 893451d4c9c..1a2bb022e08 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -18,13 +18,13 @@ import { HttpProxyAgent } from 'http-proxy-agent'; import { HttpsProxyAgent } from 'https-proxy-agent'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from './loggingContentGenerator.js'; +import { ModelMappingContentGenerator } from './modelMappingContentGenerator.js'; +import { CCPA_AI_MODEL_MAPPINGS } from '../config/models.js'; import { loadApiKey } from './apiKeyCredentialStorage.js'; 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'); @@ -153,7 +153,10 @@ describe('createContentGenerator', () => { ); expect(createCodeAssistContentGenerator).toHaveBeenCalled(); expect(generator).toEqual( - new LoggingContentGenerator(mockGenerator, mockConfig), + new LoggingContentGenerator( + new ModelMappingContentGenerator(mockGenerator, CCPA_AI_MODEL_MAPPINGS), + mockConfig, + ), ); }); @@ -170,7 +173,10 @@ describe('createContentGenerator', () => { ); expect(createCodeAssistContentGenerator).toHaveBeenCalled(); expect(generator).toEqual( - new LoggingContentGenerator(mockGenerator, mockConfig), + new LoggingContentGenerator( + new ModelMappingContentGenerator(mockGenerator, CCPA_AI_MODEL_MAPPINGS), + mockConfig, + ), ); }); @@ -212,13 +218,7 @@ describe('createContentGenerator', () => { }), }); expect(generator).toEqual( - new LoggingContentGenerator( - new ModelMappingContentGenerator( - mockGenerator.models, - GEMINI_API_MODEL_MAPPINGS, - ), - mockConfig, - ), + new LoggingContentGenerator(mockGenerator.models, mockConfig), ); }); @@ -755,13 +755,7 @@ describe('createContentGenerator', () => { }), }); expect(generator).toEqual( - new LoggingContentGenerator( - new ModelMappingContentGenerator( - mockGenerator.models, - GEMINI_API_MODEL_MAPPINGS, - ), - mockConfig, - ), + new LoggingContentGenerator(mockGenerator.models, mockConfig), ); }); @@ -1119,7 +1113,7 @@ describe('createContentGenerator', () => { ); }); - it('should apply model mapping for Vertex AI', async () => { + it('should not apply model mapping for Vertex AI', async () => { const mockModels = { generateContent: vi.fn().mockResolvedValue({}), }; @@ -1148,14 +1142,14 @@ describe('createContentGenerator', () => { expect(mockModels.generateContent).toHaveBeenCalledWith( expect.objectContaining({ - model: 'gemini-3.5-flash', + model: 'gemini-3-flash', }), 'prompt-id', 'user', ); }); - it('should apply model mapping for Gemini API', async () => { + it('should not apply model mapping for Gemini API', async () => { const mockModels = { generateContent: vi.fn().mockResolvedValue({}), }; @@ -1183,7 +1177,42 @@ describe('createContentGenerator', () => { expect(mockModels.generateContent).toHaveBeenCalledWith( expect.objectContaining({ + model: 'gemini-3-flash', + }), + 'prompt-id', + 'user', + ); + }); + + it('should apply model mapping for inverse cases (e.g. GATEWAY)', 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.GATEWAY, + }, + mockConfig, + ); + + await generator.generateContent( + { model: 'gemini-3.5-flash', + contents: [], + }, + 'prompt-id', + 'user' as LlmRole, + ); + + expect(mockModels.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'gemini-3-flash', }), 'prompt-id', 'user', diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index a1940efa254..68387e94f12 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -31,10 +31,7 @@ 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'; +import { CCPA_AI_MODEL_MAPPINGS } from '../config/models.js'; /** * Interface abstracting the core functionalities for generating content and counting tokens. @@ -287,11 +284,14 @@ export async function createContentGenerator( ) { const httpOptions = { headers: baseHeaders }; return new LoggingContentGenerator( - await createCodeAssistContentGenerator( - httpOptions, - config.authType, - gcConfig, - sessionId, + new ModelMappingContentGenerator( + await createCodeAssistContentGenerator( + httpOptions, + config.authType, + gcConfig, + sessionId, + ), + CCPA_AI_MODEL_MAPPINGS, ), gcConfig, ); @@ -381,15 +381,13 @@ export async function createContentGenerator( }), }); 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) { + if ( + config.authType !== AuthType.USE_VERTEX_AI && + config.authType !== AuthType.USE_GEMINI + ) { generator = new ModelMappingContentGenerator( generator, - GEMINI_API_MODEL_MAPPINGS, + CCPA_AI_MODEL_MAPPINGS, ); } return new LoggingContentGenerator(generator, gcConfig); From fc6e81d7c50e4aa86da1a28db41875e3ca3a4a34 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Tue, 9 Jun 2026 01:56:01 +0000 Subject: [PATCH 6/9] update config tests to accept 3.5-flash as the default model name for 3.5 flash --- packages/core/src/config/config.test.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index e48536181f0..69a1dc5c292 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -4357,7 +4357,7 @@ describe('hasGemini35FlashGAAccess model setting', () => { cwd: '.', }; - it('should set DEFAULT_GEMINI_FLASH_MODEL to gemini-3-flash and PREVIEW_GEMINI_FLASH_MODEL to gemini-3-flash-preview if hasGemini35FlashGAAccess returns true and authType is USE_GEMINI', () => { + it('should set DEFAULT_GEMINI_FLASH_MODEL to gemini-3.5-flash and PREVIEW_GEMINI_FLASH_MODEL to gemini-3-flash-preview if hasGemini35FlashGAAccess returns true and authType is USE_GEMINI', () => { const config = new Config(baseParams); config['contentGeneratorConfig'] = { authType: AuthType.USE_GEMINI }; @@ -4375,11 +4375,11 @@ describe('hasGemini35FlashGAAccess model setting', () => { const result = config.hasGemini35FlashGAAccess(); expect(result).toBe(true); - expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3-flash'); + expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3.5-flash'); expect(PREVIEW_GEMINI_FLASH_MODEL).toBe('gemini-3-flash-preview'); }); - it('should set DEFAULT_GEMINI_FLASH_MODEL and PREVIEW_GEMINI_FLASH_MODEL to gemini-3-flash if hasGemini35FlashGAAccess returns true and authType is not USE_GEMINI', () => { + it('should set DEFAULT_GEMINI_FLASH_MODEL and PREVIEW_GEMINI_FLASH_MODEL to gemini-3.5-flash if hasGemini35FlashGAAccess returns true and authType is not USE_GEMINI', () => { const config = new Config(baseParams); config['contentGeneratorConfig'] = { authType: AuthType.LOGIN_WITH_GOOGLE }; @@ -4397,7 +4397,7 @@ describe('hasGemini35FlashGAAccess model setting', () => { const result = config.hasGemini35FlashGAAccess(); expect(result).toBe(true); - expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3-flash'); - expect(PREVIEW_GEMINI_FLASH_MODEL).toBe('gemini-3-flash'); + expect(DEFAULT_GEMINI_FLASH_MODEL).toBe('gemini-3.5-flash'); + expect(PREVIEW_GEMINI_FLASH_MODEL).toBe('gemini-3.5-flash'); }); }); From b43621c5604f515ca5b5be8efaac7f394e0f036d Mon Sep 17 00:00:00 2001 From: davidapierce Date: Tue, 9 Jun 2026 03:31:16 +0000 Subject: [PATCH 7/9] update wrapping method to work with prefix. Also update unwrapping for loggingContnentGenerator so that it doesn't ignore experiment flag checks when wrapped. --- .../core/src/code_assist/codeAssist.test.ts | 47 ++++++ packages/core/src/code_assist/codeAssist.ts | 13 +- .../core/modelMappingContentGenerator.test.ts | 135 ++++++++++++++++++ .../src/core/modelMappingContentGenerator.ts | 20 ++- 4 files changed, 207 insertions(+), 8 deletions(-) create mode 100644 packages/core/src/core/modelMappingContentGenerator.test.ts diff --git a/packages/core/src/code_assist/codeAssist.test.ts b/packages/core/src/code_assist/codeAssist.test.ts index 1a4ba66f27a..0a20be1d43d 100644 --- a/packages/core/src/code_assist/codeAssist.test.ts +++ b/packages/core/src/code_assist/codeAssist.test.ts @@ -15,6 +15,7 @@ import { } from './codeAssist.js'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from '../core/loggingContentGenerator.js'; +import { ModelMappingContentGenerator } from '../core/modelMappingContentGenerator.js'; import { UserTierId } from './types.js'; // Mock dependencies @@ -22,11 +23,15 @@ vi.mock('./oauth2.js'); vi.mock('./setup.js'); vi.mock('./server.js'); vi.mock('../core/loggingContentGenerator.js'); +vi.mock('../core/modelMappingContentGenerator.js'); const mockedGetOauthClient = vi.mocked(getOauthClient); const mockedSetupUser = vi.mocked(setupUser); const MockedCodeAssistServer = vi.mocked(CodeAssistServer); const MockedLoggingContentGenerator = vi.mocked(LoggingContentGenerator); +const MockedModelMappingContentGenerator = vi.mocked( + ModelMappingContentGenerator, +); describe('codeAssist', () => { beforeEach(() => { @@ -178,5 +183,47 @@ describe('codeAssist', () => { const server = getCodeAssistServer(mockConfig); expect(server).toBeUndefined(); }); + + it('should unwrap and return the server if it is wrapped in a ModelMappingContentGenerator', () => { + const mockServer = new MockedCodeAssistServer({} as never, '', {}); + const mockMapper = new MockedModelMappingContentGenerator( + {} as never, + {}, + ); + vi.spyOn(mockMapper, 'getWrapped').mockReturnValue(mockServer); + + const mockConfig = { + getContentGenerator: () => mockMapper, + } as unknown as Config; + + const server = getCodeAssistServer(mockConfig); + expect(server).toBe(mockServer); + expect(mockMapper.getWrapped).toHaveBeenCalled(); + }); + + it('should recursively unwrap multiple layers of LoggingContentGenerator and ModelMappingContentGenerator', () => { + const mockServer = new MockedCodeAssistServer({} as never, '', {}); + const mockLogger = new MockedLoggingContentGenerator( + {} as never, + {} as never, + ); + const mockMapper = new MockedModelMappingContentGenerator( + {} as never, + {}, + ); + + // Mapper wraps Logger wraps Server + vi.spyOn(mockMapper, 'getWrapped').mockReturnValue(mockLogger); + vi.spyOn(mockLogger, 'getWrapped').mockReturnValue(mockServer); + + const mockConfig = { + getContentGenerator: () => mockMapper, + } as unknown as Config; + + const server = getCodeAssistServer(mockConfig); + expect(server).toBe(mockServer); + expect(mockMapper.getWrapped).toHaveBeenCalled(); + expect(mockLogger.getWrapped).toHaveBeenCalled(); + }); }); }); diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 4fcbea7853a..b6c28c44a7f 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/packages/core/src/code_assist/codeAssist.ts @@ -10,6 +10,7 @@ import { setupUser } from './setup.js'; import { CodeAssistServer, type HttpOptions } from './server.js'; import type { Config } from '../config/config.js'; import { LoggingContentGenerator } from '../core/loggingContentGenerator.js'; +import { ModelMappingContentGenerator } from '../core/modelMappingContentGenerator.js'; export async function createCodeAssistContentGenerator( httpOptions: HttpOptions, @@ -43,9 +44,15 @@ export function getCodeAssistServer( ): CodeAssistServer | undefined { let server = config.getContentGenerator(); - // Unwrap LoggingContentGenerator if present - if (server instanceof LoggingContentGenerator) { - server = server.getWrapped(); + // Recursively unwrap LoggingContentGenerator and ModelMappingContentGenerator + while (true) { + if (server instanceof LoggingContentGenerator) { + server = server.getWrapped(); + } else if (server instanceof ModelMappingContentGenerator) { + server = server.getWrapped(); + } else { + break; + } } if (!(server instanceof CodeAssistServer)) { diff --git a/packages/core/src/core/modelMappingContentGenerator.test.ts b/packages/core/src/core/modelMappingContentGenerator.test.ts new file mode 100644 index 00000000000..926f55fb368 --- /dev/null +++ b/packages/core/src/core/modelMappingContentGenerator.test.ts @@ -0,0 +1,135 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi } from 'vitest'; +import { ModelMappingContentGenerator } from './modelMappingContentGenerator.js'; +import type { ContentGenerator } from './contentGenerator.js'; +import { LlmRole } from '../telemetry/llmRole.js'; +import type { GenerateContentParameters } from '@google/genai'; + +describe('ModelMappingContentGenerator', () => { + const mockMappings = { + 'gemini-3.5-flash': 'gemini-3-flash', + 'gemini-pro': 'gemini-1.5-pro', + }; + + it('delegates userTier, userTierName, and paidTier properties', () => { + const mockWrapped = { + userTier: 'free', + userTierName: 'Free Tier', + paidTier: { id: 'paid' }, + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + + expect(generator.userTier).toBe('free'); + expect(generator.userTierName).toBe('Free Tier'); + expect(generator.paidTier).toEqual({ id: 'paid' }); + }); + + it('maps matching model without prefix', async () => { + const mockWrapped = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + const req = { model: 'gemini-3.5-flash', contents: [] }; + + await generator.generateContent(req, 'prompt-id', LlmRole.MAIN); + + expect(mockWrapped.generateContent).toHaveBeenCalledWith( + { model: 'gemini-3-flash', contents: [] }, + 'prompt-id', + LlmRole.MAIN, + ); + }); + + it('maps matching model with models/ prefix', async () => { + const mockWrapped = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + const req = { model: 'models/gemini-3.5-flash', contents: [] }; + + await generator.generateContent(req, 'prompt-id', LlmRole.MAIN); + + expect(mockWrapped.generateContent).toHaveBeenCalledWith( + { model: 'models/gemini-3-flash', contents: [] }, + 'prompt-id', + LlmRole.MAIN, + ); + }); + + it('leaves unmapped model unchanged', async () => { + const mockWrapped = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + const req = { model: 'unknown-model', contents: [] }; + + await generator.generateContent(req, 'prompt-id', LlmRole.MAIN); + + expect(mockWrapped.generateContent).toHaveBeenCalledWith( + { model: 'unknown-model', contents: [] }, + 'prompt-id', + LlmRole.MAIN, + ); + }); + + it('leaves model with prefix unchanged if no match after normalization', async () => { + const mockWrapped = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + const req = { model: 'models/unknown-model', contents: [] }; + + await generator.generateContent(req, 'prompt-id', LlmRole.MAIN); + + expect(mockWrapped.generateContent).toHaveBeenCalledWith( + { model: 'models/unknown-model', contents: [] }, + 'prompt-id', + LlmRole.MAIN, + ); + }); + + it('handles missing/undefined model property safely', async () => { + const mockWrapped = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + + const generator = new ModelMappingContentGenerator( + mockWrapped, + mockMappings, + ); + const req = { contents: [] } as unknown as GenerateContentParameters; + + await generator.generateContent(req, 'prompt-id', LlmRole.MAIN); + + expect(mockWrapped.generateContent).toHaveBeenCalledWith( + { contents: [] }, + 'prompt-id', + LlmRole.MAIN, + ); + }); +}); diff --git a/packages/core/src/core/modelMappingContentGenerator.ts b/packages/core/src/core/modelMappingContentGenerator.ts index aac110142e7..ef07f614ae8 100644 --- a/packages/core/src/core/modelMappingContentGenerator.ts +++ b/packages/core/src/core/modelMappingContentGenerator.ts @@ -15,6 +15,7 @@ import { import { type ContentGenerator } from './contentGenerator.js'; import type { LlmRole } from '../telemetry/llmRole.js'; import type { UserTierId, GeminiUserTier } from '../code_assist/types.js'; +import { normalizeModelId } from '../utils/modelUtils.js'; export class ModelMappingContentGenerator implements ContentGenerator { constructor( @@ -22,6 +23,10 @@ export class ModelMappingContentGenerator implements ContentGenerator { private readonly mappings: Record, ) {} + getWrapped(): ContentGenerator { + return this.wrapped; + } + get userTier(): UserTierId | undefined { return this.wrapped.userTier; } @@ -35,11 +40,16 @@ export class ModelMappingContentGenerator implements ContentGenerator { } private mapModel(req: T): T { - if (req.model && this.mappings[req.model]) { - return { - ...req, - model: this.mappings[req.model], - }; + if (req.model) { + const normalizedModel = normalizeModelId(req.model); + if (this.mappings[normalizedModel]) { + return { + ...req, + model: req.model.startsWith('models/') + ? `models/${this.mappings[normalizedModel]}` + : this.mappings[normalizedModel], + }; + } } return req; } From 2f3d1a6d92ad76a78fd3ee735c621c9fe8b599d6 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Tue, 9 Jun 2026 03:40:15 +0000 Subject: [PATCH 8/9] revert changes to use vertex ai and use gemini flow in content generator now that their model names are assumed correct by default. --- packages/core/src/core/contentGenerator.ts | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 68387e94f12..c893860d4ce 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -380,17 +380,7 @@ export async function createContentGenerator( }, }), }); - let generator: ContentGenerator = googleGenAI.models; - if ( - config.authType !== AuthType.USE_VERTEX_AI && - config.authType !== AuthType.USE_GEMINI - ) { - generator = new ModelMappingContentGenerator( - generator, - CCPA_AI_MODEL_MAPPINGS, - ); - } - return new LoggingContentGenerator(generator, gcConfig); + return new LoggingContentGenerator(googleGenAI.models, gcConfig); } throw new Error( `Error creating contentGenerator: Unsupported authType: ${config.authType}`, From 817fe9bfb32c9b37816293d14cdac3503179d943 Mon Sep 17 00:00:00 2001 From: davidapierce Date: Tue, 9 Jun 2026 03:59:21 +0000 Subject: [PATCH 9/9] update tests in content Generator to properly validate model mappings on LOGIN_WITH_GOOGLE and COMPUTE_ADC --- .../core/src/core/contentGenerator.test.ts | 68 ++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 1a2bb022e08..60cb4c563ef 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -1184,7 +1184,7 @@ describe('createContentGenerator', () => { ); }); - it('should apply model mapping for inverse cases (e.g. GATEWAY)', async () => { + it('should not apply model mapping for GATEWAY', async () => { const mockModels = { generateContent: vi.fn().mockResolvedValue({}), }; @@ -1211,6 +1211,72 @@ describe('createContentGenerator', () => { ); expect(mockModels.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'gemini-3.5-flash', + }), + 'prompt-id', + 'user', + ); + }); + + it('should apply model mapping for LOGIN_WITH_GOOGLE', async () => { + const mockInnerGenerator = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( + mockInnerGenerator as never, + ); + + const generator = await createContentGenerator( + { + authType: AuthType.LOGIN_WITH_GOOGLE, + }, + mockConfig, + ); + + await generator.generateContent( + { + model: 'gemini-3.5-flash', + contents: [], + }, + 'prompt-id', + 'user' as LlmRole, + ); + + expect(mockInnerGenerator.generateContent).toHaveBeenCalledWith( + expect.objectContaining({ + model: 'gemini-3-flash', + }), + 'prompt-id', + 'user', + ); + }); + + it('should apply model mapping for COMPUTE_ADC', async () => { + const mockInnerGenerator = { + generateContent: vi.fn().mockResolvedValue({}), + } as unknown as ContentGenerator; + vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( + mockInnerGenerator as never, + ); + + const generator = await createContentGenerator( + { + authType: AuthType.COMPUTE_ADC, + }, + mockConfig, + ); + + await generator.generateContent( + { + model: 'gemini-3.5-flash', + contents: [], + }, + 'prompt-id', + 'user' as LlmRole, + ); + + expect(mockInnerGenerator.generateContent).toHaveBeenCalledWith( expect.objectContaining({ model: 'gemini-3-flash', }),