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

Author: joshualitt

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

@@ -132,34 +132,46 @@ 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');
     });
 
     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({
         temperature: 0.1,
         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,16 +198,18 @@ 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,
       });
 
       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;
-
-  // If model changed, re-resolve config
-  if (finalModel !== requestedModel) {
-    const { generateContentConfig } =
-      config.modelConfigService.getResolvedConfig({
-        overrideScope,
-        model: finalModel,
-      });
+  const finalModel = selection.selectedModel ?? model;
+  let generateContentConfig = resolved.generateContentConfig;
 
-    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,
   };
 }

packages/core/src/core/baseLlmClient.test.ts

@@ -776,13 +776,15 @@ describe('BaseLlmClient', () => {
       const getResolvedConfigMock = vi.mocked(
         mockConfig.modelConfigService.getResolvedConfig,
       );
-      getResolvedConfigMock
-        .mockReturnValueOnce(
-          makeResolvedModelConfig(firstModel, { temperature: 0.1 }),
-        )
-        .mockReturnValueOnce(
-          makeResolvedModelConfig(fallbackModel, { temperature: 0.9 }),
-        );
+      getResolvedConfigMock.mockImplementation((key) => {
+        if (key.model === firstModel) {
+          return makeResolvedModelConfig(firstModel, { temperature: 0.1 });
+        }
+        if (key.model === fallbackModel) {
+          return makeResolvedModelConfig(fallbackModel, { temperature: 0.9 });
+        }
+        return makeResolvedModelConfig(key.model);
+      });
 
       // Availability selects the first model initially
       vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({

packages/core/src/core/baseLlmClient.ts

@@ -10,6 +10,7 @@ import type {
   EmbedContentParameters,
   GenerateContentResponse,
   GenerateContentParameters,
+  GenerateContentConfig,
 } from '@google/genai';
 import type { Config } from '../config/config.js';
 import type { ContentGenerator } from './contentGenerator.js';
@@ -81,6 +82,19 @@ export interface GenerateContentOptions {
   maxAttempts?: number;
 }
 
+interface _CommonGenerateOptions {
+  modelConfigKey: ModelConfigKey;
+  contents: Content[];
+  systemInstruction?: string | Part | Part[] | Content;
+  abortSignal: AbortSignal;
+  promptId: string;
+  maxAttempts?: number;
+  additionalProperties?: {
+    responseJsonSchema: Record<string, unknown>;
+    responseMimeType: string;
+  };
+}
+
 /**
  * A client dedicated to stateless, utility-focused LLM calls.
  */
@@ -104,7 +118,7 @@ export class BaseLlmClient {
       maxAttempts,
     } = options;
 
-    const { model, generateContentConfig } =
+    const { model } =
       this.config.modelConfigService.getResolvedConfig(modelConfigKey);
 
     const shouldRetryOnContent = (response: GenerateContentResponse) => {
@@ -123,18 +137,17 @@ export class BaseLlmClient {
 
     const result = await this._generateWithRetry(
       {
-        model,
+        modelConfigKey,
         contents,
-        config: {
-          ...generateContentConfig,
-          ...(systemInstruction && { systemInstruction }),
+        abortSignal,
+        promptId,
+        maxAttempts,
+        systemInstruction,
+        additionalProperties: {
           responseJsonSchema: schema,
           responseMimeType: 'application/json',
-          abortSignal,
         },
       },
-      promptId,
-      maxAttempts,
       shouldRetryOnContent,
       'generateJson',
     );
@@ -205,80 +218,81 @@ export class BaseLlmClient {
       maxAttempts,
     } = options;
 
-    const { model, generateContentConfig } =
-      this.config.modelConfigService.getResolvedConfig(modelConfigKey);
-
     const shouldRetryOnContent = (response: GenerateContentResponse) => {
       const text = getResponseText(response)?.trim();
       return !text; // Retry on empty response
     };
 
     return this._generateWithRetry(
       {
-        model,
+        modelConfigKey,
         contents,
-        config: {
-          ...generateContentConfig,
-          ...(systemInstruction && { systemInstruction }),
-          abortSignal,
-        },
+        systemInstruction,
+        abortSignal,
+        promptId,
+        maxAttempts,
       },
-      promptId,
-      maxAttempts,
       shouldRetryOnContent,
       'generateContent',
     );
   }
 
   private async _generateWithRetry(
-    requestParams: GenerateContentParameters,
-    promptId: string,
-    maxAttempts: number | undefined,
+    options: _CommonGenerateOptions,
     shouldRetryOnContent: (response: GenerateContentResponse) => boolean,
     errorContext: 'generateJson' | 'generateContent',
   ): Promise<GenerateContentResponse> {
-    const abortSignal = requestParams.config?.abortSignal;
-
-    // Define callback to fetch context dynamically since active model may get updated during retry loop
-    const getAvailabilityContext = createAvailabilityContextProvider(
-      this.config,
-      () => requestParams.model,
-    );
+    const {
+      modelConfigKey,
+      contents,
+      systemInstruction,
+      abortSignal,
+      promptId,
+      maxAttempts,
+      additionalProperties,
+    } = options;
 
     const {
       model,
-      config: newConfig,
+      config: generateContentConfig,
       maxAttempts: availabilityMaxAttempts,
-    } = applyModelSelection(
+    } = applyModelSelection(this.config, modelConfigKey);
+
+    let currentModel = model;
+    let currentGenerateContentConfig = generateContentConfig;
+
+    // Define callback to fetch context dynamically since active model may get updated during retry loop
+    const getAvailabilityContext = createAvailabilityContextProvider(
       this.config,
-      requestParams.model,
-      requestParams.config,
+      () => currentModel,
     );
-    requestParams.model = model;
-    if (newConfig) {
-      requestParams.config = newConfig;
-    }
-    if (abortSignal) {
-      requestParams.config = { ...requestParams.config, abortSignal };
-    }
 
     try {
       const apiCall = () => {
         // Ensure we use the current active model
         // in case a fallback occurred in a previous attempt.
         const activeModel = this.config.getActiveModel();
-        if (activeModel !== requestParams.model) {
-          requestParams.model = activeModel;
+        if (activeModel !== currentModel) {
+          currentModel = activeModel;
           // Re-resolve config if model changed during retry
           const { generateContentConfig } =
             this.config.modelConfigService.getResolvedConfig({
+              ...modelConfigKey,
               model: activeModel,
             });
-          requestParams.config = {
-            ...requestParams.config,
-            ...generateContentConfig,
-          };
+          currentGenerateContentConfig = generateContentConfig;
         }
+        const finalConfig: GenerateContentConfig = {
+          ...currentGenerateContentConfig,
+          ...(systemInstruction && { systemInstruction }),
+          ...additionalProperties,
+          abortSignal,
+        };
+        const requestParams: GenerateContentParameters = {
+          model: currentModel,
+          config: finalConfig,
+          contents,
+        };
         return this.contentGenerator.generateContent(requestParams, promptId);
       };
 
@@ -289,7 +303,7 @@ export class BaseLlmClient {
         getAvailabilityContext,
         onPersistent429: this.config.isInteractive()
           ? (authType, error) =>
-              handleFallback(this.config, requestParams.model, authType, error)
+              handleFallback(this.config, currentModel, authType, error)
           : undefined,
         authType:
           this.authType ?? this.config.getContentGeneratorConfig()?.authType,
@@ -307,14 +321,14 @@ export class BaseLlmClient {
         await reportError(
           error,
           `API returned invalid content after all retries.`,
-          requestParams.contents as Content[],
+          contents,
           `${errorContext}-invalid-content`,
         );
       } else {
         await reportError(
           error,
           `Error generating content via API.`,
-          requestParams.contents as Content[],
+          contents,
           `${errorContext}-api`,
         );
       }