From 415070e504a86e16e630fcdcacef3adf629433e7 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 6 Jan 2026 23:00:13 +0000 Subject: [PATCH] tests --- .../core/src/agents/local-executor.test.ts | 101 +++++++++++++++++- packages/core/src/agents/registry.test.ts | 36 +++++++ .../strategies/classifierStrategy.test.ts | 26 +++++ .../strategies/fallbackStrategy.test.ts | 21 ++++ .../strategies/overrideStrategy.test.ts | 21 ++++ .../src/services/modelConfigService.test.ts | 75 +++++++++++++ 6 files changed, 279 insertions(+), 1 deletion(-) 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/registry.test.ts b/packages/core/src/agents/registry.test.ts index f369e59b21..5d6dfedd97 100644 --- a/packages/core/src/agents/registry.test.ts +++ b/packages/core/src/agents/registry.test.ts @@ -244,6 +244,42 @@ 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 two overrides: one for the alias and one for the agent name (scope) + expect(registerOverrideSpy).toHaveBeenCalledTimes(2); + + // Check alias override + expect(registerOverrideSpy).toHaveBeenCalledWith( + expect.objectContaining({ + match: { model: getModelConfigAlias(autoAgent) }, + modelConfig: expect.objectContaining({ model: 'auto' }), + }), + ); + + // 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/routing/strategies/classifierStrategy.test.ts b/packages/core/src/routing/strategies/classifierStrategy.test.ts index 21d324c1fb..e883b0be45 100644 --- a/packages/core/src/routing/strategies/classifierStrategy.test.ts +++ b/packages/core/src/routing/strategies/classifierStrategy.test.ts @@ -281,4 +281,30 @@ describe('ClassifierStrategy', () => { ); consoleWarnSpy.mockRestore(); }); + + it('should respect requestedModel from context in resolveClassifierModel', async () => { + const requestedModel = DEFAULT_GEMINI_MODEL; // Pro model + const mockApiResponse = { + reasoning: 'Choice is flash', + model_choice: 'flash', + }; + vi.mocked(mockBaseLlmClient.generateJson).mockResolvedValue( + mockApiResponse, + ); + + const contextWithRequestedModel = { + ...mockContext, + requestedModel, + } as RoutingContext; + + const decision = await strategy.route( + contextWithRequestedModel, + mockConfig, + mockBaseLlmClient, + ); + + expect(decision).not.toBeNull(); + // Since requestedModel is Pro, and choice is flash, it should resolve to Flash + expect(decision?.model).toBe(DEFAULT_GEMINI_FLASH_MODEL); + }); }); diff --git a/packages/core/src/routing/strategies/fallbackStrategy.test.ts b/packages/core/src/routing/strategies/fallbackStrategy.test.ts index 6196e59526..2d30b153e5 100644 --- a/packages/core/src/routing/strategies/fallbackStrategy.test.ts +++ b/packages/core/src/routing/strategies/fallbackStrategy.test.ts @@ -108,4 +108,25 @@ describe('FallbackStrategy', () => { // Important: check that it queried snapshot with the RESOLVED model, not 'auto' expect(mockService.snapshot).toHaveBeenCalledWith(DEFAULT_GEMINI_MODEL); }); + + it('should respect requestedModel from context', async () => { + const requestedModel = 'requested-model'; + const configModel = 'config-model'; + vi.mocked(mockConfig.getModel).mockReturnValue(configModel); + vi.mocked(mockService.snapshot).mockReturnValue({ available: true }); + + const contextWithRequestedModel = { + requestedModel, + } as RoutingContext; + + const decision = await strategy.route( + contextWithRequestedModel, + mockConfig, + mockClient, + ); + + expect(decision).toBeNull(); + // Should check availability of the requested model from context + expect(mockService.snapshot).toHaveBeenCalledWith(requestedModel); + }); }); diff --git a/packages/core/src/routing/strategies/overrideStrategy.test.ts b/packages/core/src/routing/strategies/overrideStrategy.test.ts index f1ec54098d..97e9f4915f 100644 --- a/packages/core/src/routing/strategies/overrideStrategy.test.ts +++ b/packages/core/src/routing/strategies/overrideStrategy.test.ts @@ -56,4 +56,25 @@ describe('OverrideStrategy', () => { expect(decision).not.toBeNull(); expect(decision?.model).toBe(overrideModel); }); + + it('should respect requestedModel from context', async () => { + const requestedModel = 'requested-model'; + const configModel = 'config-model'; + const mockConfig = { + getModel: () => configModel, + getPreviewFeatures: () => false, + } as Config; + const contextWithRequestedModel = { + requestedModel, + } as RoutingContext; + + const decision = await strategy.route( + contextWithRequestedModel, + mockConfig, + mockClient, + ); + + expect(decision).not.toBeNull(); + expect(decision?.model).toBe(requestedModel); + }); }); diff --git a/packages/core/src/services/modelConfigService.test.ts b/packages/core/src/services/modelConfigService.test.ts index 8d08e4f775..ee6cd09f40 100644 --- a/packages/core/src/services/modelConfigService.test.ts +++ b/packages/core/src/services/modelConfigService.test.ts @@ -577,6 +577,81 @@ describe('ModelConfigService', () => { }); }); + describe('runtime overrides', () => { + it('should resolve a simple runtime-registered override', () => { + const config: ModelConfigServiceConfig = { + aliases: {}, + overrides: [], + }; + const service = new ModelConfigService(config); + + service.registerRuntimeModelOverride({ + match: { model: 'gemini-pro' }, + modelConfig: { + generateContentConfig: { + temperature: 0.99, + }, + }, + }); + + const resolved = service.getResolvedConfig({ model: 'gemini-pro' }); + + expect(resolved.model).toBe('gemini-pro'); + expect(resolved.generateContentConfig.temperature).toBe(0.99); + }); + + it('should prioritize runtime overrides over default overrides when they have the same specificity', () => { + const config: ModelConfigServiceConfig = { + aliases: {}, + overrides: [ + { + match: { model: 'gemini-pro' }, + modelConfig: { generateContentConfig: { temperature: 0.1 } }, + }, + ], + }; + const service = new ModelConfigService(config); + + service.registerRuntimeModelOverride({ + match: { model: 'gemini-pro' }, + modelConfig: { generateContentConfig: { temperature: 0.9 } }, + }); + + const resolved = service.getResolvedConfig({ model: 'gemini-pro' }); + + // Runtime overrides are appended after overrides/customOverrides, so they should win. + expect(resolved.generateContentConfig.temperature).toBe(0.9); + }); + + it('should still respect specificity with runtime overrides', () => { + const config: ModelConfigServiceConfig = { + aliases: {}, + overrides: [], + }; + const service = new ModelConfigService(config); + + // Register a more specific runtime override + service.registerRuntimeModelOverride({ + match: { model: 'gemini-pro', overrideScope: 'my-agent' }, + modelConfig: { generateContentConfig: { temperature: 0.1 } }, + }); + + // Register a less specific runtime override later + service.registerRuntimeModelOverride({ + match: { model: 'gemini-pro' }, + modelConfig: { generateContentConfig: { temperature: 0.9 } }, + }); + + const resolved = service.getResolvedConfig({ + model: 'gemini-pro', + overrideScope: 'my-agent', + }); + + // Specificity should win over order + expect(resolved.generateContentConfig.temperature).toBe(0.1); + }); + }); + describe('custom aliases', () => { it('should resolve a custom alias', () => { const config: ModelConfigServiceConfig = {