feat(core): Migrate generateJson to resolved model configs. (#12626)

packages/core/src/utils/editCorrector.test.ts

@@ -236,6 +236,14 @@ describe('editCorrector', () => {
       mockGeminiClientInstance.getHistory = vi.fn().mockResolvedValue([]);
       mockBaseLlmClientInstance = {
         generateJson: mockGenerateJson,
+        config: {
+          generationConfigService: {
+            getResolvedConfig: vi.fn().mockReturnValue({
+              model: 'edit-corrector',
+              generateContentConfig: {},
+            }),
+          },
+        },
       } as unknown as Mocked<BaseLlmClient>;
       resetEditCorrectorCaches_TEST_ONLY();
     });
@@ -634,6 +642,14 @@ describe('editCorrector', () => {
 
       mockBaseLlmClientInstance = {
         generateJson: mockGenerateJson,
+        config: {
+          generationConfigService: {
+            getResolvedConfig: vi.fn().mockReturnValue({
+              model: 'edit-corrector',
+              generateContentConfig: {},
+            }),
+          },
+        },
       } as unknown as Mocked<BaseLlmClient>;
       resetEditCorrectorCaches_TEST_ONLY();
     });

packages/core/src/utils/editCorrector.ts

@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { Content, GenerateContentConfig } from '@google/genai';
+import type { Content } from '@google/genai';
 import type { GeminiClient } from '../core/client.js';
 import type { BaseLlmClient } from '../core/baseLlmClient.js';
 import type { EditToolParams } from '../tools/edit.js';
@@ -16,7 +16,6 @@ import {
   WRITE_FILE_TOOL_NAME,
 } from '../tools/tool-names.js';
 import { LruCache } from './LruCache.js';
-import { DEFAULT_GEMINI_FLASH_LITE_MODEL } from '../config/models.js';
 import {
   isFunctionResponse,
   isFunctionCall,
@@ -24,13 +23,6 @@ import {
 import * as fs from 'node:fs';
 import { promptIdContext } from './promptIdContext.js';
 
-const EDIT_MODEL = DEFAULT_GEMINI_FLASH_LITE_MODEL;
-const EDIT_CONFIG: GenerateContentConfig = {
-  thinkingConfig: {
-    thinkingBudget: 0,
-  },
-};
-
 const CODE_CORRECTION_SYSTEM_PROMPT = `
 You are an expert code-editing assistant. Your task is to analyze a failed edit attempt and provide a corrected version of the text snippets.
 The correction should be as minimal as possible, staying very close to the original.
@@ -420,11 +412,10 @@ Return ONLY the corrected target snippet in the specified JSON format with the k
 
   try {
     const result = await baseLlmClient.generateJson({
+      modelConfigKey: { model: 'edit-corrector' },
       contents,
       schema: OLD_STRING_CORRECTION_SCHEMA,
       abortSignal,
-      model: EDIT_MODEL,
-      config: EDIT_CONFIG,
       systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT,
       promptId: getPromptId(),
     });
@@ -510,11 +501,10 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr
 
   try {
     const result = await baseLlmClient.generateJson({
+      modelConfigKey: { model: 'edit-corrector' },
       contents,
       schema: NEW_STRING_CORRECTION_SCHEMA,
       abortSignal,
-      model: EDIT_MODEL,
-      config: EDIT_CONFIG,
       systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT,
       promptId: getPromptId(),
     });
@@ -581,11 +571,10 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr
 
   try {
     const result = await baseLlmClient.generateJson({
+      modelConfigKey: { model: 'edit-corrector' },
       contents,
       schema: CORRECT_NEW_STRING_ESCAPING_SCHEMA,
       abortSignal,
-      model: EDIT_MODEL,
-      config: EDIT_CONFIG,
       systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT,
       promptId: getPromptId(),
     });
@@ -649,11 +638,10 @@ Return ONLY the corrected string in the specified JSON format with the key 'corr
 
   try {
     const result = await baseLlmClient.generateJson({
+      modelConfigKey: { model: 'edit-corrector' },
       contents,
       schema: CORRECT_STRING_ESCAPING_SCHEMA,
       abortSignal,
-      model: EDIT_MODEL,
-      config: EDIT_CONFIG,
       systemInstruction: CODE_CORRECTION_SYSTEM_PROMPT,
       promptId: getPromptId(),
     });

packages/core/src/utils/llm-edit-fixer.test.ts

@@ -17,6 +17,14 @@ import type { BaseLlmClient } from '../core/baseLlmClient.js';
 const mockGenerateJson = vi.fn();
 const mockBaseLlmClient = {
   generateJson: mockGenerateJson,
+  config: {
+    generationConfigService: {
+      getResolvedConfig: vi.fn().mockReturnValue({
+        model: 'edit-corrector',
+        generateContentConfig: {},
+      }),
+    },
+  },
 } as unknown as BaseLlmClient;
 
 describe('FixLLMEditWithInstruction', () => {

packages/core/src/utils/llm-edit-fixer.ts

@@ -8,7 +8,6 @@ import { createHash } from 'node:crypto';
 import { type Content, Type } from '@google/genai';
 import { type BaseLlmClient } from '../core/baseLlmClient.js';
 import { LruCache } from './LruCache.js';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { promptIdContext } from './promptIdContext.js';
 import { debugLogger } from './debugLogger.js';
 
@@ -176,10 +175,10 @@ export async function FixLLMEditWithInstruction(
   const result = await generateJsonWithTimeout<SearchReplaceEdit>(
     baseLlmClient,
     {
+      modelConfigKey: { model: 'llm-edit-fixer' },
       contents,
       schema: SearchReplaceEditSchema,
       abortSignal,
-      model: DEFAULT_GEMINI_FLASH_MODEL,
       systemInstruction: EDIT_SYS_PROMPT,
       promptId,
       maxAttempts: 1,

packages/core/src/utils/nextSpeakerChecker.test.ts

@@ -7,7 +7,6 @@
 import type { Mock } from 'vitest';
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import type { Content } from '@google/genai';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { BaseLlmClient } from '../core/baseLlmClient.js';
 import type { ContentGenerator } from '../core/contentGenerator.js';
 import type { Config } from '../config/config.js';
@@ -54,6 +53,10 @@ describe('checkNextSpeaker', () => {
 
   beforeEach(() => {
     vi.resetAllMocks();
+    const mockResolvedConfig = {
+      model: 'next-speaker-v1',
+      generateContentConfig: {},
+    };
     mockConfig = {
       getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
       getSessionId: vi.fn().mockReturnValue('test-session-id'),
@@ -61,6 +64,9 @@ describe('checkNextSpeaker', () => {
       storage: {
         getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
       },
+      modelConfigService: {
+        getResolvedConfig: vi.fn().mockReturnValue(mockResolvedConfig),
+      },
     } as unknown as Config;
 
     mockBaseLlmClient = new BaseLlmClient(
@@ -265,8 +271,8 @@ describe('checkNextSpeaker', () => {
 
     expect(mockBaseLlmClient.generateJson).toHaveBeenCalled();
     const generateJsonCall = (mockBaseLlmClient.generateJson as Mock).mock
-      .calls[0];
-    expect(generateJsonCall[0].model).toBe(DEFAULT_GEMINI_FLASH_MODEL);
-    expect(generateJsonCall[0].promptId).toBe(promptId);
+      .calls[0][0];
+    expect(generateJsonCall.modelConfigKey.model).toBe('next-speaker-checker');
+    expect(generateJsonCall.promptId).toBe(promptId);
   });
 });

packages/core/src/utils/nextSpeakerChecker.ts

@@ -5,7 +5,6 @@
  */
 
 import type { Content } from '@google/genai';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import type { BaseLlmClient } from '../core/baseLlmClient.js';
 import type { GeminiChat } from '../core/geminiChat.js';
 import { isFunctionResponse } from './messageInspectors.js';
@@ -111,9 +110,9 @@ export async function checkNextSpeaker(
 
   try {
     const parsedResponse = (await baseLlmClient.generateJson({
+      modelConfigKey: { model: 'next-speaker-checker' },
       contents,
       schema: RESPONSE_SCHEMA,
-      model: DEFAULT_GEMINI_FLASH_MODEL,
       abortSignal,
       promptId,
     })) as unknown as NextSpeakerResponse;