From 41a8809280f844870600f63a4beea4c0bd585d8e Mon Sep 17 00:00:00 2001 From: joshualitt Date: Thu, 8 Jan 2026 12:39:40 -0800 Subject: [PATCH] feat(core): Wire up model routing to subagents. (#16043) --- .../core/src/agents/local-executor.test.ts | 101 +++++++++++++++++- packages/core/src/agents/local-executor.ts | 39 ++++++- packages/core/src/agents/registry.test.ts | 45 ++++++++ packages/core/src/agents/registry.ts | 42 +++++--- 4 files changed, 211 insertions(+), 16 deletions(-) diff --git a/packages/core/src/agents/local-executor.test.ts b/packages/core/src/agents/local-executor.test.ts index 98d017c864..a0a8a513f2 100644 --- a/packages/core/src/agents/local-executor.test.ts +++ b/packages/core/src/agents/local-executor.test.ts @@ -57,8 +57,12 @@ import { AgentTerminateMode } from './types.js'; import type { AnyDeclarativeTool, AnyToolInvocation } from '../tools/tools.js'; import { CompressionStatus } from '../core/turn.js'; import { ChatCompressionService } from '../services/chatCompressionService.js'; -import type { ModelConfigKey } from '../services/modelConfigService.js'; +import type { + ModelConfigKey, + ResolvedModelConfig, +} from '../services/modelConfigService.js'; import { getModelConfigAlias } from './registry.js'; +import type { ModelRouterService } from '../routing/modelRouterService.js'; const { mockSendMessageStream, @@ -1192,6 +1196,101 @@ describe('LocalAgentExecutor', () => { }); }); + describe('Model Routing', () => { + it('should use model routing when the agent model is "auto"', async () => { + const definition = createTestDefinition(); + definition.modelConfig.model = 'auto'; + + const mockRouter = { + route: vi.fn().mockResolvedValue({ + model: 'routed-model', + metadata: { source: 'test', reasoning: 'test' }, + }), + }; + vi.spyOn(mockConfig, 'getModelRouterService').mockReturnValue( + mockRouter as unknown as ModelRouterService, + ); + + // Mock resolved config to return 'auto' + vi.spyOn( + mockConfig.modelConfigService, + 'getResolvedConfig', + ).mockReturnValue({ + model: 'auto', + generateContentConfig: {}, + } as unknown as ResolvedModelConfig); + + const executor = await LocalAgentExecutor.create( + definition, + mockConfig, + onActivity, + ); + + mockModelResponse([ + { + name: TASK_COMPLETE_TOOL_NAME, + args: { finalResult: 'done' }, + id: 'call1', + }, + ]); + + await executor.run({ goal: 'test' }, signal); + + expect(mockRouter.route).toHaveBeenCalled(); + expect(mockSendMessageStream).toHaveBeenCalledWith( + expect.objectContaining({ model: 'routed-model' }), + expect.any(Array), + expect.any(String), + expect.any(AbortSignal), + ); + }); + + it('should NOT use model routing when the agent model is NOT "auto"', async () => { + const definition = createTestDefinition(); + definition.modelConfig.model = 'concrete-model'; + + const mockRouter = { + route: vi.fn(), + }; + vi.spyOn(mockConfig, 'getModelRouterService').mockReturnValue( + mockRouter as unknown as ModelRouterService, + ); + + // Mock resolved config to return 'concrete-model' + vi.spyOn( + mockConfig.modelConfigService, + 'getResolvedConfig', + ).mockReturnValue({ + model: 'concrete-model', + generateContentConfig: {}, + } as unknown as ResolvedModelConfig); + + const executor = await LocalAgentExecutor.create( + definition, + mockConfig, + onActivity, + ); + + mockModelResponse([ + { + name: TASK_COMPLETE_TOOL_NAME, + args: { finalResult: 'done' }, + id: 'call1', + }, + ]); + + await executor.run({ goal: 'test' }, signal); + + expect(mockRouter.route).not.toHaveBeenCalled(); + expect(mockSendMessageStream).toHaveBeenCalledWith( + expect.objectContaining({ model: 'concrete-model' }), + expect.any(Array), + expect.any(String), + expect.any(AbortSignal), + ); + }); + }); + describe('run (Termination Conditions)', () => { const mockWorkResponse = (id: string) => { mockModelResponse([{ name: LS_TOOL_NAME, args: { path: '.' }, id }]); diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 994c616594..fc866c97b5 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -40,6 +40,8 @@ import type { } from './types.js'; import { AgentTerminateMode } from './types.js'; import { templateString } from './utils.js'; +import { DEFAULT_GEMINI_MODEL, isAutoModel } from '../config/models.js'; +import type { RoutingContext } from '../routing/routingStrategy.js'; import { parseThought } from '../utils/thoughtUtils.js'; import { type z } from 'zod'; import { zodToJsonSchema } from 'zod-to-json-schema'; @@ -589,9 +591,44 @@ export class LocalAgentExecutor { signal: AbortSignal, promptId: string, ): Promise<{ functionCalls: FunctionCall[]; textResponse: string }> { + const modelConfigAlias = getModelConfigAlias(this.definition); + + // Resolve the model config early to get the concrete model string (which may be `auto`). + const resolvedConfig = + this.runtimeContext.modelConfigService.getResolvedConfig({ + model: modelConfigAlias, + overrideScope: this.definition.name, + }); + const requestedModel = resolvedConfig.model; + + let modelToUse: string; + if (isAutoModel(requestedModel)) { + // TODO(joshualitt): This try / catch is inconsistent with the routing + // behavior for the main agent. Ideally, we would have a universal + // policy for routing failure. Given routing failure does not necessarily + // mean generation will fail, we may want to share this logic with + // other places we use model routing. + try { + const routingContext: RoutingContext = { + history: chat.getHistory(/*curated=*/ true), + request: message.parts || [], + signal, + requestedModel, + }; + const router = this.runtimeContext.getModelRouterService(); + const decision = await router.route(routingContext); + modelToUse = decision.model; + } catch (error) { + debugLogger.warn(`Error during model routing: ${error}`); + modelToUse = DEFAULT_GEMINI_MODEL; + } + } else { + modelToUse = requestedModel; + } + const responseStream = await chat.sendMessageStream( { - model: getModelConfigAlias(this.definition), + model: modelToUse, overrideScope: this.definition.name, }, message.parts || [], diff --git a/packages/core/src/agents/registry.test.ts b/packages/core/src/agents/registry.test.ts index 4864566bc0..84a9001a03 100644 --- a/packages/core/src/agents/registry.test.ts +++ b/packages/core/src/agents/registry.test.ts @@ -243,6 +243,51 @@ describe('AgentRegistry', () => { }); describe('registration logic', () => { + it('should register runtime overrides when the model is "auto"', async () => { + const autoAgent: LocalAgentDefinition = { + ...MOCK_AGENT_V1, + name: 'AutoAgent', + modelConfig: { ...MOCK_AGENT_V1.modelConfig, model: 'auto' }, + }; + + const registerOverrideSpy = vi.spyOn( + mockConfig.modelConfigService, + 'registerRuntimeModelOverride', + ); + + await registry.testRegisterAgent(autoAgent); + + // Should register one alias for the custom model config. + expect( + mockConfig.modelConfigService.getResolvedConfig({ + model: getModelConfigAlias(autoAgent), + }), + ).toStrictEqual({ + model: 'auto', + generateContentConfig: { + temperature: autoAgent.modelConfig.temp, + topP: autoAgent.modelConfig.top_p, + thinkingConfig: { + includeThoughts: true, + thinkingBudget: -1, + }, + }, + }); + + // Should register one override for the agent name (scope) + expect(registerOverrideSpy).toHaveBeenCalledTimes(1); + + // Check scope override + expect(registerOverrideSpy).toHaveBeenCalledWith( + expect.objectContaining({ + match: { overrideScope: autoAgent.name }, + modelConfig: expect.objectContaining({ + generateContentConfig: expect.any(Object), + }), + }), + ); + }); + it('should register a valid agent definition', async () => { await registry.testRegisterAgent(MOCK_AGENT_V1); expect(registry.getDefinition('MockAgent')).toEqual(MOCK_AGENT_V1); diff --git a/packages/core/src/agents/registry.ts b/packages/core/src/agents/registry.ts index 13f203d7d1..ee42795a66 100644 --- a/packages/core/src/agents/registry.ts +++ b/packages/core/src/agents/registry.ts @@ -20,8 +20,8 @@ import { GEMINI_MODEL_ALIAS_AUTO, PREVIEW_GEMINI_FLASH_MODEL, isPreviewModel, + isAutoModel, } from '../config/models.js'; -import type { ModelConfigAlias } from '../services/modelConfigService.js'; /** * Returns the model config alias for a given agent definition. @@ -199,7 +199,10 @@ export class AgentRegistry { this.agents.set(definition.name, definition); - // Register model config. + // Register model config. We always create a runtime alias. However, + // if the user is using `auto` as a model string then we also create + // runtime overrides to ensure the subagent generation settings are + // respected regardless of the final model string from routing. // TODO(12916): Migrate sub-agents where possible to static configs. const modelConfig = definition.modelConfig; let model = modelConfig.model; @@ -207,24 +210,35 @@ export class AgentRegistry { model = this.config.getModel(); } - const runtimeAlias: ModelConfigAlias = { - modelConfig: { - model, - generateContentConfig: { - temperature: modelConfig.temp, - topP: modelConfig.top_p, - thinkingConfig: { - includeThoughts: true, - thinkingBudget: modelConfig.thinkingBudget ?? -1, - }, - }, + const generateContentConfig = { + temperature: modelConfig.temp, + topP: modelConfig.top_p, + thinkingConfig: { + includeThoughts: true, + thinkingBudget: modelConfig.thinkingBudget ?? -1, }, }; this.config.modelConfigService.registerRuntimeModelConfig( getModelConfigAlias(definition), - runtimeAlias, + { + modelConfig: { + model, + generateContentConfig, + }, + }, ); + + if (isAutoModel(model)) { + this.config.modelConfigService.registerRuntimeModelOverride({ + match: { + overrideScope: definition.name, + }, + modelConfig: { + generateContentConfig, + }, + }); + } } /**