From e377f87880eaea7c0665096bc090120ad8d080f2 Mon Sep 17 00:00:00 2001 From: Johnny Eric Amancio Date: Sun, 5 Jul 2026 21:11:22 +0200 Subject: [PATCH] feat(models): expose auto-routing choices in models API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Annotate the virtual kilo-auto/efficient and kilo-auto/free models with their concrete candidate ids (autoRouting.models) in the models API responses, so the model selector can show what auto-routing picks among. Applied on /api/openrouter/models (org and public/BYOK paths) and on /api/organizations/[id]/models — the endpoint the extension uses for organization catalogs, where the choices were originally missing. Candidates are filtered to models available to the caller: enrichment runs after org restrictions and feature filtering, so a deny-listed or feature-excluded model never surfaces as a routing choice. The benchmark routing table is cached in-process (5-min, stale-on-error) so it is not a per-request round-trip and a transient worker outage does not blank the choices. Adds unit coverage for addAutoRoutingModels and first direct test coverage for getAvailableModelsForOrganization. --- .../app/api/openrouter/models/route.test.ts | 139 ++++++++++++++---- .../src/app/api/openrouter/models/route.ts | 53 +------ .../organizations/[id]/models/route.test.ts | 129 ++++++++++++++++ .../api/organizations/[id]/models/route.ts | 10 +- .../ai-gateway/auto-routing-models.test.ts | 127 ++++++++++++++++ .../src/lib/ai-gateway/auto-routing-models.ts | 45 ++++++ .../ai-gateway/auto-routing-table-cache.ts | 24 +++ 7 files changed, 447 insertions(+), 80 deletions(-) create mode 100644 apps/web/src/app/api/organizations/[id]/models/route.test.ts create mode 100644 apps/web/src/lib/ai-gateway/auto-routing-models.test.ts create mode 100644 apps/web/src/lib/ai-gateway/auto-routing-models.ts create mode 100644 apps/web/src/lib/ai-gateway/auto-routing-table-cache.ts diff --git a/apps/web/src/app/api/openrouter/models/route.test.ts b/apps/web/src/app/api/openrouter/models/route.test.ts index 322a9ce5c4..547d41f91c 100644 --- a/apps/web/src/app/api/openrouter/models/route.test.ts +++ b/apps/web/src/app/api/openrouter/models/route.test.ts @@ -1,15 +1,7 @@ import { beforeEach, describe, expect, test } from '@jest/globals'; import { NextRequest } from 'next/server'; import type { OpenRouterModel } from '@/lib/organizations/organization-types'; -import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter'; -import { getUserFromAuth } from '@/lib/user/server'; -import { getDirectByokModelsForUser } from '@/lib/ai-gateway/providers/direct-byok'; -import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; -import { addUserByokAvailability, getUserByokProviderIds } from '@/lib/ai-gateway/byok'; -import { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models'; import { GET } from './route'; -import { getBenchmarkRoutingTable } from '@/lib/ai-gateway/auto-routing-benchmark-admin-client'; -import { getAutoFreeCandidates } from '@/lib/ai-gateway/auto-model/resolution'; jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() })); @@ -26,17 +18,36 @@ jest.mock('@/lib/ai-gateway/byok', () => ({ addUserByokAvailability: jest.fn(), getUserByokProviderIds: jest.fn(), })); +jest.mock('@/lib/ai-gateway/models', () => ({ + ...jest.requireActual('@/lib/ai-gateway/models'), + filterByFeature: jest.fn(), +})); jest.mock('@/lib/organizations/organization-models', () => ({ getAvailableModelsForOrganization: jest.fn(), })); -jest.mock('@/lib/ai-gateway/auto-routing-benchmark-admin-client', () => ({ - getBenchmarkRoutingTable: jest.fn(), +jest.mock('@/lib/ai-gateway/auto-routing-table-cache', () => ({ + getCachedRoutingTable: jest.fn(), })); jest.mock('@/lib/ai-gateway/auto-model/resolution', () => ({ getAutoFreeCandidates: jest.fn(), })); jest.mock('@/lib/drizzle', () => ({ readDb: {} })); +const { getUserFromAuth } = jest.requireMock('@/lib/user/server'); +const { getEnhancedOpenRouterModels } = jest.requireMock('@/lib/ai-gateway/providers/openrouter'); +const { getDirectByokModelsForUser } = jest.requireMock('@/lib/ai-gateway/providers/direct-byok'); +const { listAvailableExperimentModels } = jest.requireMock( + '@/lib/ai-gateway/experiments/list-available-experiment-models' +); +const { addUserByokAvailability, getUserByokProviderIds } = + jest.requireMock('@/lib/ai-gateway/byok'); +const { getAvailableModelsForOrganization } = jest.requireMock( + '@/lib/organizations/organization-models' +); +const { getCachedRoutingTable } = jest.requireMock('@/lib/ai-gateway/auto-routing-table-cache'); +const { getAutoFreeCandidates } = jest.requireMock('@/lib/ai-gateway/auto-model/resolution'); +const { filterByFeature } = jest.requireMock('@/lib/ai-gateway/models'); + const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); const mockedGetEnhancedOpenRouterModels = jest.mocked(getEnhancedOpenRouterModels); const mockedGetDirectByokModelsForUser = jest.mocked(getDirectByokModelsForUser); @@ -44,8 +55,9 @@ const mockedListAvailableExperimentModels = jest.mocked(listAvailableExperimentM const mockedAddUserByokAvailability = jest.mocked(addUserByokAvailability); const mockedGetUserByokProviderIds = jest.mocked(getUserByokProviderIds); const mockedGetAvailableModelsForOrganization = jest.mocked(getAvailableModelsForOrganization); -const mockedGetBenchmarkRoutingTable = jest.mocked(getBenchmarkRoutingTable); +const mockedGetCachedRoutingTable = jest.mocked(getCachedRoutingTable); const mockedGetAutoFreeCandidates = jest.mocked(getAutoFreeCandidates); +const mockedFilterByFeature = jest.mocked(filterByFeature); function makeModel(id: string): OpenRouterModel { return { @@ -64,13 +76,17 @@ function makeModel(id: string): OpenRouterModel { }; } -function request() { - return new NextRequest('http://localhost:3000/api/openrouter/models'); +function request(headers?: Record) { + return new NextRequest('http://localhost:3000/api/openrouter/models', { headers }); } describe('GET /api/openrouter/models', () => { beforeEach(() => { jest.resetAllMocks(); + mockedFilterByFeature.mockImplementation( + (models: Array<{ id: string }>, feature: string | null) => + feature === 'code-review' ? models.filter(model => model.id !== 'feature/excluded') : models + ); mockedGetUserFromAuth.mockResolvedValue({ user: null, organizationId: null, @@ -81,10 +97,7 @@ describe('GET /api/openrouter/models', () => { mockedListAvailableExperimentModels.mockResolvedValue([]); mockedGetUserByokProviderIds.mockResolvedValue([]); mockedGetAvailableModelsForOrganization.mockResolvedValue(null); - mockedGetBenchmarkRoutingTable.mockResolvedValue({ - status: 200, - body: { table: null, publishedAt: null }, - }); + mockedGetCachedRoutingTable.mockResolvedValue(null); mockedGetAutoFreeCandidates.mockResolvedValue([]); }); @@ -95,6 +108,8 @@ describe('GET /api/openrouter/models', () => { await expect(response.json()).resolves.toEqual({ data: [makeModel('public/model')] }); expect(mockedGetUserByokProviderIds).not.toHaveBeenCalled(); expect(mockedAddUserByokAvailability).not.toHaveBeenCalled(); + expect(mockedGetCachedRoutingTable).not.toHaveBeenCalled(); + expect(mockedGetAutoFreeCandidates).not.toHaveBeenCalled(); }); test('returns BYOK availability for regular and direct authenticated models', async () => { @@ -135,21 +150,16 @@ describe('GET /api/openrouter/models', () => { 'poolside/laguna-m.1:free', 'missing/free-model', ]); - mockedGetBenchmarkRoutingTable.mockResolvedValue({ - status: 200, - body: { - table: { - routes: { - 'implementation/code_generation': [ - { model: 'google/gemini-2.5-flash' }, - { model: 'openai/gpt-5.4-mini' }, - ], - 'analysis/debugging': [ - { model: 'kilo-auto/balanced' }, - { model: 'google/gemini-2.5-flash' }, - ], - }, - }, + mockedGetCachedRoutingTable.mockResolvedValue({ + routes: { + 'implementation/code_generation': [ + { model: 'google/gemini-2.5-flash' }, + { model: 'openai/gpt-5.4-mini' }, + ], + 'analysis/debugging': [ + { model: 'kilo-auto/balanced' }, + { model: 'google/gemini-2.5-flash' }, + ], }, } as never); @@ -177,4 +187,69 @@ describe('GET /api/openrouter/models', () => { ], }); }); + + test('builds auto-routing choices from feature-filtered models', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + const allowedModel = makeModel('openai/gpt-5.4-mini'); + const excludedModel = makeModel('feature/excluded'); + mockedGetEnhancedOpenRouterModels.mockResolvedValue({ + data: [efficientModel, allowedModel, excludedModel], + }); + mockedGetCachedRoutingTable.mockResolvedValue({ + routes: { + 'implementation/code_generation': [{ model: allowedModel.id }, { model: excludedModel.id }], + }, + } as never); + + const response = await GET(request({ 'x-kilocode-feature': 'code-review' })); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + data: [ + { + ...efficientModel, + autoRouting: { + models: [allowedModel.id], + }, + }, + allowedModel, + ], + }); + }); + + test('adds auto-routing to organization models after feature filtering', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + const allowedModel = makeModel('openai/gpt-5.4-mini'); + const excludedModel = makeModel('feature/excluded'); + mockedGetUserFromAuth.mockResolvedValue({ + user: { id: 'user-id' }, + organizationId: 'org-1', + authFailedResponse: null, + } as never); + mockedGetAvailableModelsForOrganization.mockResolvedValue({ + data: [efficientModel, allowedModel, excludedModel], + } as never); + mockedGetCachedRoutingTable.mockResolvedValue({ + routes: { + 'implementation/code_generation': [{ model: allowedModel.id }, { model: excludedModel.id }], + }, + } as never); + + const response = await GET(request({ 'x-kilocode-feature': 'code-review' })); + + // The feature-excluded model must appear in neither the top-level list nor the + // Auto Efficient choices — enrichment runs after filterByFeature. + expect(response.status).toBe(200); + await expect(response.json()).resolves.toEqual({ + data: [ + { + ...efficientModel, + autoRouting: { + models: [allowedModel.id], + }, + }, + allowedModel, + ], + }); + }); }); diff --git a/apps/web/src/app/api/openrouter/models/route.ts b/apps/web/src/app/api/openrouter/models/route.ts index 99960288bb..c9fd71e439 100644 --- a/apps/web/src/app/api/openrouter/models/route.ts +++ b/apps/web/src/app/api/openrouter/models/route.ts @@ -11,10 +11,7 @@ import { filterByFeature } from '@/lib/ai-gateway/models'; import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models'; import { addUserByokAvailability, getUserByokProviderIds } from '@/lib/ai-gateway/byok'; import { readDb } from '@/lib/drizzle'; -import { getBenchmarkRoutingTable } from '@/lib/ai-gateway/auto-routing-benchmark-admin-client'; -import { KILO_AUTO_EFFICIENT_MODEL, KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; -import { getAutoFreeCandidates } from '@/lib/ai-gateway/auto-model/resolution'; -import { isVirtualAutoModelId } from '@kilocode/auto-routing-contracts'; +import { addAutoRoutingModels } from '@/lib/ai-gateway/auto-routing-models'; async function tryGetUserFromAuth() { try { @@ -25,41 +22,6 @@ async function tryGetUserFromAuth() { } } -function visibleConcreteModelIds(models: Iterable, availableModelIds: ReadonlySet) { - return [ - ...new Set([...models].filter(id => availableModelIds.has(id) && !isVirtualAutoModelId(id))), - ].sort((left, right) => left.localeCompare(right)); -} - -async function addAutoRoutingModels(models: OpenRouterModelsResponse['data']) { - const availableModelIds = new Set(models.map(model => model.id)); - const [routingTableResult, autoFreeCandidates] = await Promise.all([ - getBenchmarkRoutingTable().catch(() => null), - getAutoFreeCandidates(null).catch(() => []), - ]); - - const table = - routingTableResult?.status === 200 && 'table' in routingTableResult.body - ? routingTableResult.body.table - : null; - const efficientModelIds = visibleConcreteModelIds( - Object.values(table?.routes ?? {}) - .flat() - .map(candidate => candidate.model), - availableModelIds - ); - const freeModelIds = visibleConcreteModelIds(autoFreeCandidates, availableModelIds); - const autoRoutingChoices = new Map([ - [KILO_AUTO_EFFICIENT_MODEL.id, efficientModelIds], - [KILO_AUTO_FREE_MODEL.id, freeModelIds], - ]); - - return models.map(model => { - const modelIds = autoRoutingChoices.get(model.id); - return modelIds?.length ? { ...model, autoRouting: { models: modelIds } } : model; - }); -} - /** * Test using: * curl -vvv 'http://localhost:3000/api/openrouter/models' @@ -74,10 +36,9 @@ export async function GET( ? await getAvailableModelsForOrganization(auth.organizationId) : null; if (result) { - const models = await addAutoRoutingModels(result.data); return NextResponse.json({ ...result, - data: filterByFeature(models, feature), + data: await addAutoRoutingModels(filterByFeature(result.data, feature)), }); } @@ -85,11 +46,11 @@ export async function GET( if (!Array.isArray(data.data)) { return NextResponse.json(data); } - const models = await addAutoRoutingModels(data.data); + const models = await addAutoRoutingModels(filterByFeature(data.data, feature)); if (!auth?.user) { const experimentModels = await listAvailableExperimentModels(); return NextResponse.json({ - data: filterByFeature(models.concat(experimentModels), feature), + data: models.concat(filterByFeature(experimentModels, feature)), }); } @@ -103,9 +64,9 @@ export async function GET( enabledByokProviderIds ); return NextResponse.json({ - data: filterByFeature( - modelsWithByokAvailability.concat(byokModels, experimentModels), - feature + data: modelsWithByokAvailability.concat( + filterByFeature(byokModels, feature), + filterByFeature(experimentModels, feature) ), }); } catch (error) { diff --git a/apps/web/src/app/api/organizations/[id]/models/route.test.ts b/apps/web/src/app/api/organizations/[id]/models/route.test.ts new file mode 100644 index 0000000000..2324530655 --- /dev/null +++ b/apps/web/src/app/api/organizations/[id]/models/route.test.ts @@ -0,0 +1,129 @@ +import { beforeEach, describe, expect, test } from '@jest/globals'; +import { NextRequest, NextResponse } from 'next/server'; +import type { OpenRouterModel } from '@/lib/organizations/organization-types'; +import { handleTRPCRequest } from '@/lib/trpc-route-handler'; +import { GET } from './route'; + +jest.mock('@/lib/trpc-route-handler', () => ({ handleTRPCRequest: jest.fn() })); +jest.mock('@/lib/ai-gateway/auto-routing-table-cache', () => ({ + getCachedRoutingTable: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/auto-model/resolution', () => ({ + getAutoFreeCandidates: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/models', () => ({ + ...jest.requireActual('@/lib/ai-gateway/models'), + filterByFeature: jest.fn(), +})); + +const mockedHandleTRPCRequest = jest.mocked(handleTRPCRequest); +const { getCachedRoutingTable } = jest.requireMock('@/lib/ai-gateway/auto-routing-table-cache'); +const { getAutoFreeCandidates } = jest.requireMock('@/lib/ai-gateway/auto-model/resolution'); +const { filterByFeature } = jest.requireMock('@/lib/ai-gateway/models'); +const mockedGetCachedRoutingTable = jest.mocked(getCachedRoutingTable); +const mockedGetAutoFreeCandidates = jest.mocked(getAutoFreeCandidates); +const mockedFilterByFeature = jest.mocked(filterByFeature); +const listAvailableModels = jest.fn(); + +function makeModel(id: string): OpenRouterModel { + return { + id, + name: id, + created: 0, + description: '', + architecture: { + input_modalities: ['text'], + output_modalities: ['text'], + tokenizer: 'test', + }, + top_provider: { is_moderated: false }, + pricing: { prompt: '0', completion: '0' }, + context_length: 0, + supported_parameters: ['tools'], + }; +} + +function request(headers?: Record) { + return new NextRequest('http://localhost:3000/api/organizations/org-1/models', { headers }); +} + +describe('GET /api/organizations/[id]/models', () => { + beforeEach(() => { + jest.resetAllMocks(); + mockedFilterByFeature.mockImplementation( + (models: Array<{ id: string }>, feature: string | null) => + feature === 'code-review' ? models.filter(model => model.id !== 'feature/excluded') : models + ); + mockedGetCachedRoutingTable.mockResolvedValue(null); + mockedGetAutoFreeCandidates.mockResolvedValue([]); + mockedHandleTRPCRequest.mockImplementation(async (request, handler) => { + const result = await handler({ + organizations: { settings: { listAvailableModels } }, + } as never); + return NextResponse.json(result); + }); + }); + + test('annotates the organization catalog with auto-routing choices', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + const allowedModel = makeModel('openai/gpt-5.4-mini'); + const deniedModel = 'anthropic/claude-opus-4.8'; + listAvailableModels.mockResolvedValue({ + data: [efficientModel, allowedModel], + }); + mockedGetCachedRoutingTable.mockResolvedValue({ + routes: { + 'implementation/code_generation': [{ model: allowedModel.id }, { model: deniedModel }], + }, + } as never); + + const response = await GET(request(), { + params: Promise.resolve({ id: 'org-1' }), + }); + + // Candidates outside the organization catalog (e.g. deny-listed) must not appear. + expect(listAvailableModels).toHaveBeenCalledWith({ + organizationId: 'org-1', + }); + await expect(response.json()).resolves.toEqual({ + data: [{ ...efficientModel, autoRouting: { models: [allowedModel.id] } }, allowedModel], + }); + }); + + test('builds auto-routing choices from the feature-filtered catalog', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + const allowedModel = makeModel('openai/gpt-5.4-mini'); + const excludedModel = makeModel('feature/excluded'); + listAvailableModels.mockResolvedValue({ + data: [efficientModel, allowedModel, excludedModel], + }); + mockedGetCachedRoutingTable.mockResolvedValue({ + routes: { + 'implementation/code_generation': [{ model: allowedModel.id }, { model: excludedModel.id }], + }, + } as never); + + const response = await GET(request({ 'x-kilocode-feature': 'code-review' }), { + params: Promise.resolve({ id: 'org-1' }), + }); + + // The feature-excluded model must appear in neither the top-level list nor + // the Auto Efficient choices — enrichment runs after filterByFeature. + await expect(response.json()).resolves.toEqual({ + data: [{ ...efficientModel, autoRouting: { models: [allowedModel.id] } }, allowedModel], + }); + }); + + test('returns the catalog unchanged when auto models are absent', async () => { + const model = makeModel('openai/gpt-5.4-mini'); + listAvailableModels.mockResolvedValue({ data: [model] }); + + const response = await GET(request(), { + params: Promise.resolve({ id: 'org-1' }), + }); + + await expect(response.json()).resolves.toEqual({ data: [model] }); + expect(mockedGetCachedRoutingTable).not.toHaveBeenCalled(); + expect(mockedGetAutoFreeCandidates).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/app/api/organizations/[id]/models/route.ts b/apps/web/src/app/api/organizations/[id]/models/route.ts index 1a47cf6ba5..eb00852863 100644 --- a/apps/web/src/app/api/organizations/[id]/models/route.ts +++ b/apps/web/src/app/api/organizations/[id]/models/route.ts @@ -3,13 +3,19 @@ import type { OpenRouterModelsResponse } from '@/lib/organizations/organization- import { handleTRPCRequest } from '@/lib/trpc-route-handler'; import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection'; import { filterByFeature } from '@/lib/ai-gateway/models'; +import { addAutoRoutingModels } from '@/lib/ai-gateway/auto-routing-models'; export async function GET(request: NextRequest, { params }: { params: Promise<{ id: string }> }) { const organizationId = (await params).id; const feature = validateFeatureHeader(request.headers.get(FEATURE_HEADER)); return handleTRPCRequest(request, async caller => { - const result = await caller.organizations.settings.listAvailableModels({ organizationId }); - return { ...result, data: filterByFeature(result.data, feature) }; + const result = await caller.organizations.settings.listAvailableModels({ + organizationId, + }); + return { + ...result, + data: await addAutoRoutingModels(filterByFeature(result.data, feature)), + }; }); } diff --git a/apps/web/src/lib/ai-gateway/auto-routing-models.test.ts b/apps/web/src/lib/ai-gateway/auto-routing-models.test.ts new file mode 100644 index 0000000000..18ad9a5d39 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/auto-routing-models.test.ts @@ -0,0 +1,127 @@ +import { beforeEach, describe, expect, test } from '@jest/globals'; +import type { OpenRouterModel } from '@/lib/organizations/organization-types'; +import { addAutoRoutingModels } from './auto-routing-models'; + +jest.mock('@/lib/ai-gateway/auto-routing-table-cache', () => ({ + getCachedRoutingTable: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/auto-model/resolution', () => ({ + getAutoFreeCandidates: jest.fn(), +})); + +const { getCachedRoutingTable } = jest.requireMock('@/lib/ai-gateway/auto-routing-table-cache'); +const { getAutoFreeCandidates } = jest.requireMock('@/lib/ai-gateway/auto-model/resolution'); +const mockedGetCachedRoutingTable = jest.mocked(getCachedRoutingTable); +const mockedGetAutoFreeCandidates = jest.mocked(getAutoFreeCandidates); + +function makeModel(id: string): OpenRouterModel { + return { + id, + name: id, + created: 0, + description: '', + architecture: { + input_modalities: ['text'], + output_modalities: ['text'], + tokenizer: 'test', + }, + top_provider: { is_moderated: false }, + pricing: { prompt: '0', completion: '0' }, + context_length: 0, + }; +} + +function routingTable(candidateIds: string[]) { + return { + routes: { + 'implementation/code_generation': candidateIds.map(model => ({ model })), + }, + } as never; +} + +describe('addAutoRoutingModels', () => { + beforeEach(() => { + jest.resetAllMocks(); + mockedGetCachedRoutingTable.mockResolvedValue(null); + mockedGetAutoFreeCandidates.mockResolvedValue([]); + }); + + test('only lists candidates present in the provided catalog', async () => { + // The catalog passed in is the caller's already-filtered view (org deny + // lists, feature filtering). A routing-table candidate missing from it — + // e.g. deny-listed for an enterprise org — must not surface as a choice, + // even though inference-time routing knows about it. + const efficientModel = makeModel('kilo-auto/efficient'); + const visibleModel = makeModel('google/gemini-2.5-flash'); + mockedGetCachedRoutingTable.mockResolvedValue( + routingTable(['google/gemini-2.5-flash', 'anthropic/claude-opus-4.8']) + ); + + const result = await addAutoRoutingModels([efficientModel, visibleModel]); + + expect(result).toEqual([ + { ...efficientModel, autoRouting: { models: ['google/gemini-2.5-flash'] } }, + visibleModel, + ]); + }); + + test('excludes virtual auto ids and dedupes and sorts candidates', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + const balancedModel = makeModel('kilo-auto/balanced'); + const geminiModel = makeModel('google/gemini-2.5-flash'); + const gptModel = makeModel('openai/gpt-5.4-mini'); + mockedGetCachedRoutingTable.mockResolvedValue( + routingTable([ + 'openai/gpt-5.4-mini', + 'kilo-auto/balanced', + 'google/gemini-2.5-flash', + 'openai/gpt-5.4-mini', + ]) + ); + + const result = await addAutoRoutingModels([ + efficientModel, + balancedModel, + geminiModel, + gptModel, + ]); + + expect(result[0]).toEqual({ + ...efficientModel, + autoRouting: { models: ['google/gemini-2.5-flash', 'openai/gpt-5.4-mini'] }, + }); + expect(result.slice(1)).toEqual([balancedModel, geminiModel, gptModel]); + }); + + test('annotates the free auto model from its candidate source', async () => { + const freeModel = makeModel('kilo-auto/free'); + const visibleFreeModel = makeModel('poolside/laguna-m.1:free'); + mockedGetAutoFreeCandidates.mockResolvedValue(['poolside/laguna-m.1:free', 'missing/model']); + + const result = await addAutoRoutingModels([freeModel, visibleFreeModel]); + + expect(result).toEqual([ + { ...freeModel, autoRouting: { models: ['poolside/laguna-m.1:free'] } }, + visibleFreeModel, + ]); + }); + + test('leaves the auto model unannotated when no candidates are visible', async () => { + const efficientModel = makeModel('kilo-auto/efficient'); + mockedGetCachedRoutingTable.mockResolvedValue(routingTable(['denied/model'])); + + const result = await addAutoRoutingModels([efficientModel]); + + expect(result).toEqual([efficientModel]); + }); + + test('skips routing lookups when no auto models are in the catalog', async () => { + const model = makeModel('openai/gpt-5.4-mini'); + + const result = await addAutoRoutingModels([model]); + + expect(result).toEqual([model]); + expect(mockedGetCachedRoutingTable).not.toHaveBeenCalled(); + expect(mockedGetAutoFreeCandidates).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/auto-routing-models.ts b/apps/web/src/lib/ai-gateway/auto-routing-models.ts new file mode 100644 index 0000000000..ffd2f37445 --- /dev/null +++ b/apps/web/src/lib/ai-gateway/auto-routing-models.ts @@ -0,0 +1,45 @@ +import type { OpenRouterModelsResponse } from '@/lib/organizations/organization-types'; +import { KILO_AUTO_EFFICIENT_MODEL, KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; +import { getAutoFreeCandidates } from '@/lib/ai-gateway/auto-model/resolution'; +import { isVirtualAutoModelId } from '@kilocode/auto-routing-contracts'; +import { getCachedRoutingTable } from '@/lib/ai-gateway/auto-routing-table-cache'; + +function visibleConcreteModelIds(models: Iterable, availableModelIds: ReadonlySet) { + return [ + ...new Set([...models].filter(id => availableModelIds.has(id) && !isVirtualAutoModelId(id))), + ].sort((left, right) => left.localeCompare(right)); +} + +export async function addAutoRoutingModels( + models: OpenRouterModelsResponse['data'] +): Promise { + const availableModelIds = new Set(models.map(model => model.id)); + if ( + !availableModelIds.has(KILO_AUTO_EFFICIENT_MODEL.id) && + !availableModelIds.has(KILO_AUTO_FREE_MODEL.id) + ) { + return models; + } + + const [table, autoFreeCandidates] = await Promise.all([ + getCachedRoutingTable(), + getAutoFreeCandidates(null).catch(() => []), + ]); + + const efficientModelIds = visibleConcreteModelIds( + Object.values(table?.routes ?? {}) + .flat() + .map(candidate => candidate.model), + availableModelIds + ); + const freeModelIds = visibleConcreteModelIds(autoFreeCandidates, availableModelIds); + const autoRoutingChoices = new Map([ + [KILO_AUTO_EFFICIENT_MODEL.id, efficientModelIds], + [KILO_AUTO_FREE_MODEL.id, freeModelIds], + ]); + + return models.map(model => { + const modelIds = autoRoutingChoices.get(model.id); + return modelIds?.length ? { ...model, autoRouting: { models: modelIds } } : model; + }); +} diff --git a/apps/web/src/lib/ai-gateway/auto-routing-table-cache.ts b/apps/web/src/lib/ai-gateway/auto-routing-table-cache.ts new file mode 100644 index 0000000000..119f85974c --- /dev/null +++ b/apps/web/src/lib/ai-gateway/auto-routing-table-cache.ts @@ -0,0 +1,24 @@ +import { getBenchmarkRoutingTable } from '@/lib/ai-gateway/auto-routing-benchmark-admin-client'; +import { createCachedFetch } from '@/lib/cached-fetch'; + +const ROUTING_TABLE_TTL_MS = 5 * 60 * 1000; + +/** + * Cache the benchmark routing table in-process. It is fetched on every model + * listing (org endpoint and the tRPC settings query), so an uncached admin-worker + * round-trip per request is wasteful. `createCachedFetch` also serves the + * last-known-good table when a refresh throws, so a transient worker outage does + * not blank the Auto Efficient choices shown in the UI. + */ +export const getCachedRoutingTable = createCachedFetch( + async () => { + const result = await getBenchmarkRoutingTable(); + if (result.status === 200 && 'table' in result.body) { + // `table` may be null, meaning no routing is configured — a valid state to cache. + return result.body.table; + } + throw new Error(`benchmark routing table unavailable (status ${result.status})`); + }, + ROUTING_TABLE_TTL_MS, + null +);