Skip to content

feat(models): expose auto-routing choices in models API #4396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
johnnyeric merged 1 commit into from
Jul 6, 2026
Merged

feat(models): expose auto-routing choices in models API #4396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 107 additions & 32 deletions apps/web/src/app/api/openrouter/models/route.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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() }));
Expand All @@ -26,26 +18,46 @@ 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);
const mockedListAvailableExperimentModels = jest.mocked(listAvailableExperimentModels);
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 {
Expand All @@ -64,13 +76,17 @@ function makeModel(id: string): OpenRouterModel {
};
}

function request() {
return new NextRequest('http://localhost:3000/api/openrouter/models');
function request(headers?: Record<string, string>) {
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,
Expand All @@ -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([]);
});

Expand All @@ -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 () => {
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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,
],
});
});
});
53 changes: 7 additions & 46 deletions apps/web/src/app/api/openrouter/models/route.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -25,41 +22,6 @@ async function tryGetUserFromAuth() {
}
}

function visibleConcreteModelIds(models: Iterable<string>, availableModelIds: ReadonlySet<string>) {
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'
Expand All @@ -74,22 +36,21 @@ 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)),
});
}

const data = await getEnhancedOpenRouterModels();
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)),
});
}

Expand All @@ -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) {
Expand Down
Loading