feat(core): Late resolve GenerateContentConfigs and reduce mutation. (#14920)

packages/core/src/availability/policyHelpers.test.ts

@@ -132,11 +132,15 @@ describe('policyHelpers', () => {
 
     it('returns requested model if it is available', () => {
       const config = createExtendedMockConfig();
+      mockModelConfigService.getResolvedConfig.mockReturnValue({
+        model: 'gemini-pro',
+        generateContentConfig: {},
+      });
       mockAvailabilityService.selectFirstAvailable.mockReturnValue({
         selectedModel: 'gemini-pro',
       });
 
-      const result = applyModelSelection(config, 'gemini-pro');
+      const result = applyModelSelection(config, { model: 'gemini-pro' });
       expect(result.model).toBe('gemini-pro');
       expect(result.maxAttempts).toBeUndefined();
       expect(config.setActiveModel).toHaveBeenCalledWith('gemini-pro');
@@ -144,15 +148,20 @@ describe('policyHelpers', () => {
 
     it('switches to backup model and updates config if requested is unavailable', () => {
       const config = createExtendedMockConfig();
+      mockModelConfigService.getResolvedConfig
+        .mockReturnValueOnce({
+          model: 'gemini-pro',
+          generateContentConfig: { temperature: 0.9, topP: 1 },
+        })
+        .mockReturnValueOnce({
+          model: 'gemini-flash',
+          generateContentConfig: { temperature: 0.1, topP: 1 },
+        });
       mockAvailabilityService.selectFirstAvailable.mockReturnValue({
         selectedModel: 'gemini-flash',
       });
-      mockModelConfigService.getResolvedConfig.mockReturnValue({
-        generateContentConfig: { temperature: 0.1 },
-      });
 
-      const currentConfig = { temperature: 0.9, topP: 1 };
-      const result = applyModelSelection(config, 'gemini-pro', currentConfig);
+      const result = applyModelSelection(config, { model: 'gemini-pro' });
 
       expect(result.model).toBe('gemini-flash');
       expect(result.config).toEqual({
@@ -160,6 +169,9 @@ describe('policyHelpers', () => {
         topP: 1,
       });
 
+      expect(mockModelConfigService.getResolvedConfig).toHaveBeenCalledWith({
+        model: 'gemini-pro',
+      });
       expect(mockModelConfigService.getResolvedConfig).toHaveBeenCalledWith({
         model: 'gemini-flash',
       });
@@ -168,12 +180,16 @@ describe('policyHelpers', () => {
 
     it('consumes sticky attempt if indicated', () => {
       const config = createExtendedMockConfig();
+      mockModelConfigService.getResolvedConfig.mockReturnValue({
+        model: 'gemini-pro',
+        generateContentConfig: {},
+      });
       mockAvailabilityService.selectFirstAvailable.mockReturnValue({
         selectedModel: 'gemini-pro',
         attempts: 1,
       });
 
-      const result = applyModelSelection(config, 'gemini-pro');
+      const result = applyModelSelection(config, { model: 'gemini-pro' });
       expect(mockAvailabilityService.consumeStickyAttempt).toHaveBeenCalledWith(
         'gemini-pro',
       );
@@ -182,6 +198,10 @@ describe('policyHelpers', () => {
 
     it('does not consume sticky attempt if consumeAttempt is false', () => {
       const config = createExtendedMockConfig();
+      mockModelConfigService.getResolvedConfig.mockReturnValue({
+        model: 'gemini-pro',
+        generateContentConfig: {},
+      });
       mockAvailabilityService.selectFirstAvailable.mockReturnValue({
         selectedModel: 'gemini-pro',
         attempts: 1,
@@ -189,9 +209,7 @@ describe('policyHelpers', () => {
 
       const result = applyModelSelection(
         config,
-        'gemini-pro',
-        undefined,
-        undefined,
+        { model: 'gemini-pro' },
         {
           consumeAttempt: false,
         },

packages/core/src/availability/policyHelpers.ts

@@ -25,6 +25,7 @@ import {
   resolveModel,
 } from '../config/models.js';
 import type { ModelSelectionResult } from './modelAvailabilityService.js';
+import type { ModelConfigKey } from '../services/modelConfigService.js';
 
 /**
  * Resolves the active policy chain for the given config, ensuring the
@@ -155,31 +156,26 @@ export function selectModelForAvailability(
  */
 export function applyModelSelection(
   config: Config,
-  requestedModel: string,
-  currentConfig?: GenerateContentConfig,
-  overrideScope?: string,
+  modelConfigKey: ModelConfigKey,
   options: { consumeAttempt?: boolean } = {},
-): { model: string; config?: GenerateContentConfig; maxAttempts?: number } {
-  const selection = selectModelForAvailability(config, requestedModel);
+): { model: string; config: GenerateContentConfig; maxAttempts?: number } {
+  const resolved = config.modelConfigService.getResolvedConfig(modelConfigKey);
+  const model = resolved.model;
+  const selection = selectModelForAvailability(config, model);
 
-  if (!selection?.selectedModel) {
-    return { model: requestedModel, config: currentConfig };
+  if (!selection) {
+    return { model, config: resolved.generateContentConfig };
   }
 
-  const finalModel = selection.selectedModel;
-  let finalConfig = currentConfig;
+  const finalModel = selection.selectedModel ?? model;
+  let generateContentConfig = resolved.generateContentConfig;
 
-  // If model changed, re-resolve config
-  if (finalModel !== requestedModel) {
-    const { generateContentConfig } =
-      config.modelConfigService.getResolvedConfig({
-        overrideScope,
-        model: finalModel,
-      });
-
-    finalConfig = currentConfig
-      ? { ...currentConfig, ...generateContentConfig }
-      : generateContentConfig;
+  if (finalModel !== model) {
+    const fallbackResolved = config.modelConfigService.getResolvedConfig({
+      ...modelConfigKey,
+      model: finalModel,
+    });
+    generateContentConfig = fallbackResolved.generateContentConfig;
   }
 
   config.setActiveModel(finalModel);
@@ -190,7 +186,7 @@ export function applyModelSelection(
 
   return {
     model: finalModel,
-    config: finalConfig,
+    config: generateContentConfig,
     maxAttempts: selection.attempts,
   };
 }