Skip to content

Commit 66ef6fc

Browse files
authored
feat(cloud-agent-next): reject unavailable models before admission (#3471)
* feat(cloud-agent-next): reject unavailable models before admission * fix(cloud-agent-next): address model preflight review findings * fix(ai-gateway): require scoped organization model validation * refactor(ai-gateway): inline model validation predicate * fix(cloud-agent-next): harden model validation admission * fix(cloud-agent-next): skip official validation on route 404
1 parent 39b5af7 commit 66ef6fc

29 files changed

Lines changed: 1632 additions & 116 deletions
Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
import { beforeEach, describe, expect, test } from '@jest/globals';
2+
import { NextRequest, NextResponse } from 'next/server';
3+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
4+
import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter';
5+
import { getUserFromAuth } from '@/lib/user/server';
6+
import { getDirectByokModelsForUser } from '@/lib/ai-gateway/providers/direct-byok';
7+
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
8+
import { ORGANIZATION_ID_HEADER } from '@/lib/constants';
9+
import { POST } from './route';
10+
11+
jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() }));
12+
jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() }));
13+
jest.mock('@/lib/ai-gateway/providers/openrouter', () => ({
14+
getEnhancedOpenRouterModels: jest.fn(),
15+
}));
16+
jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({
17+
getDirectByokModelsForUser: jest.fn(),
18+
}));
19+
jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({
20+
listAvailableExperimentModels: jest.fn(),
21+
}));
22+
23+
const mockedGetUserFromAuth = jest.mocked(getUserFromAuth);
24+
const mockedGetEnhancedOpenRouterModels = jest.mocked(getEnhancedOpenRouterModels);
25+
const mockedGetDirectByokModelsForUser = jest.mocked(getDirectByokModelsForUser);
26+
const mockedListAvailableExperimentModels = jest.mocked(listAvailableExperimentModels);
27+
28+
function makeModel(id: string): OpenRouterModel {
29+
return {
30+
id,
31+
name: id,
32+
created: 0,
33+
description: '',
34+
architecture: {
35+
input_modalities: ['text'],
36+
output_modalities: ['text'],
37+
tokenizer: 'test',
38+
},
39+
top_provider: { is_moderated: false },
40+
pricing: { prompt: '0', completion: '0' },
41+
context_length: 0,
42+
supported_parameters: ['tools'],
43+
};
44+
}
45+
46+
function request(modelId: string, headers?: HeadersInit) {
47+
return new NextRequest('http://localhost:3000/api/openrouter/models/validate', {
48+
method: 'POST',
49+
headers,
50+
body: JSON.stringify({ modelId }),
51+
});
52+
}
53+
54+
describe('POST /api/openrouter/models/validate', () => {
55+
beforeEach(() => {
56+
jest.resetAllMocks();
57+
mockedGetUserFromAuth.mockResolvedValue({
58+
user: null,
59+
organizationId: null,
60+
authFailedResponse: null,
61+
} as never);
62+
mockedGetEnhancedOpenRouterModels.mockResolvedValue({ data: [makeModel('available/model')] });
63+
mockedGetDirectByokModelsForUser.mockResolvedValue([]);
64+
mockedListAvailableExperimentModels.mockResolvedValue([]);
65+
});
66+
67+
test('confirms a Kilo-eligible catalog model', async () => {
68+
const response = await POST(request('available/model'));
69+
70+
expect(response.status).toBe(200);
71+
await expect(response.json()).resolves.toEqual({ valid: true });
72+
});
73+
74+
test('does not expose details for an unavailable model', async () => {
75+
const response = await POST(request('missing/model'));
76+
77+
expect(response.status).toBe(200);
78+
await expect(response.json()).resolves.toEqual({ valid: false, reason: 'unavailable' });
79+
});
80+
81+
test('uses the public catalog after failed optional authentication', async () => {
82+
mockedGetUserFromAuth.mockResolvedValue({
83+
user: null,
84+
organizationId: null,
85+
authFailedResponse: NextResponse.json({ error: 'Unauthorized' }, { status: 401 }),
86+
} as never);
87+
88+
const response = await POST(request('available/model'));
89+
90+
expect(response.status).toBe(200);
91+
await expect(response.json()).resolves.toEqual({ valid: true });
92+
});
93+
94+
test('rejects organization-scoped validation through the personal endpoint', async () => {
95+
const response = await POST(
96+
request('available/model', { [ORGANIZATION_ID_HEADER]: 'organization-id' })
97+
);
98+
99+
expect(response.status).toBe(400);
100+
await expect(response.json()).resolves.toEqual({
101+
error: 'Organization-scoped validation must use /api/organizations/[id]/models/validate',
102+
});
103+
expect(mockedGetUserFromAuth).not.toHaveBeenCalled();
104+
expect(mockedGetEnhancedOpenRouterModels).not.toHaveBeenCalled();
105+
});
106+
107+
test('returns a service failure when catalog construction fails', async () => {
108+
mockedGetEnhancedOpenRouterModels.mockRejectedValue(new Error('catalog unavailable'));
109+
110+
const response = await POST(request('available/model'));
111+
112+
expect(response.status).toBe(500);
113+
await expect(response.json()).resolves.toEqual({
114+
error: 'Failed to validate model',
115+
message: 'Error from model catalog',
116+
});
117+
});
118+
119+
test('loads authenticated auxiliary catalogs concurrently', async () => {
120+
mockedGetUserFromAuth.mockResolvedValue({
121+
user: { id: 'user-id' },
122+
organizationId: null,
123+
authFailedResponse: null,
124+
} as never);
125+
mockedGetEnhancedOpenRouterModels.mockResolvedValue({ data: [] });
126+
let markByokStarted: (() => void) | undefined;
127+
const byokStarted = new Promise<void>(resolve => {
128+
markByokStarted = resolve;
129+
});
130+
type DirectByokModels = Awaited<ReturnType<typeof getDirectByokModelsForUser>>;
131+
let resolveByok: ((models: DirectByokModels) => void) | undefined;
132+
const byokPending = new Promise<DirectByokModels>(resolve => {
133+
resolveByok = resolve;
134+
});
135+
mockedGetDirectByokModelsForUser.mockImplementation(() => {
136+
if (!markByokStarted) throw new Error('BYOK start signal was not initialized');
137+
markByokStarted();
138+
return byokPending;
139+
});
140+
mockedListAvailableExperimentModels.mockResolvedValue([makeModel('experiment/model')]);
141+
142+
const responsePromise = POST(request('experiment/model'));
143+
await byokStarted;
144+
const finishByok = resolveByok;
145+
if (!finishByok) throw new Error('BYOK lookup did not start');
146+
try {
147+
expect(mockedListAvailableExperimentModels).toHaveBeenCalledTimes(1);
148+
} finally {
149+
finishByok([]);
150+
await responsePromise;
151+
}
152+
});
153+
154+
test('rejects an invalid body without reading a catalog', async () => {
155+
const response = await POST(
156+
new NextRequest('http://localhost:3000/api/openrouter/models/validate', {
157+
method: 'POST',
158+
body: JSON.stringify({ modelId: '' }),
159+
})
160+
);
161+
162+
expect(response.status).toBe(400);
163+
expect(mockedGetEnhancedOpenRouterModels).not.toHaveBeenCalled();
164+
});
165+
});
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
import type { NextRequest } from 'next/server';
2+
import { NextResponse } from 'next/server';
3+
import { captureException } from '@sentry/nextjs';
4+
import * as z from 'zod';
5+
import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter';
6+
import { getUserFromAuth } from '@/lib/user/server';
7+
import { getDirectByokModelsForUser } from '@/lib/ai-gateway/providers/direct-byok';
8+
import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection';
9+
import { ORGANIZATION_ID_HEADER } from '@/lib/constants';
10+
import { filterByFeature } from '@/lib/ai-gateway/models';
11+
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
12+
13+
const BodySchema = z.object({ modelId: z.string().trim().min(1) });
14+
15+
async function tryGetUserFromAuth() {
16+
try {
17+
return await getUserFromAuth({ adminOnly: false });
18+
} catch (error) {
19+
console.error('[validateOpenRouterModel] failed to get user from auth', error);
20+
return { user: null, organizationId: null };
21+
}
22+
}
23+
24+
export async function POST(request: NextRequest) {
25+
let body: unknown;
26+
try {
27+
body = await request.json();
28+
} catch {
29+
return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 });
30+
}
31+
32+
const bodyResult = BodySchema.safeParse(body);
33+
if (!bodyResult.success) {
34+
return NextResponse.json(
35+
{ error: 'Invalid request body', details: z.treeifyError(bodyResult.error) },
36+
{ status: 400 }
37+
);
38+
}
39+
40+
if (request.headers.get(ORGANIZATION_ID_HEADER)) {
41+
return NextResponse.json(
42+
{ error: 'Organization-scoped validation must use /api/organizations/[id]/models/validate' },
43+
{ status: 400 }
44+
);
45+
}
46+
47+
const feature = validateFeatureHeader(request.headers.get(FEATURE_HEADER));
48+
const auth = await tryGetUserFromAuth();
49+
try {
50+
const models = await getEnhancedOpenRouterModels();
51+
if (!Array.isArray(models.data)) {
52+
throw new Error('Model catalog returned invalid data');
53+
}
54+
const [byokModels, experimentModels] = await Promise.all([
55+
auth?.user ? getDirectByokModelsForUser(auth.user.id) : [],
56+
listAvailableExperimentModels(),
57+
]);
58+
const available = filterByFeature(
59+
models.data.concat(byokModels, experimentModels),
60+
feature
61+
).some(model => model.id === bodyResult.data.modelId);
62+
return NextResponse.json(available ? { valid: true } : { valid: false, reason: 'unavailable' });
63+
} catch (error) {
64+
captureException(error, {
65+
tags: { endpoint: 'openrouter/models/validate' },
66+
extra: {
67+
action: 'validating_model',
68+
userId: auth?.user?.id,
69+
organizationId: auth?.organizationId,
70+
},
71+
});
72+
return NextResponse.json(
73+
{ error: 'Failed to validate model', message: 'Error from model catalog' },
74+
{ status: 500 }
75+
);
76+
}
77+
}
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import { beforeEach, describe, expect, test } from '@jest/globals';
2+
import { NextRequest, NextResponse } from 'next/server';
3+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
4+
import { handleTRPCRequest } from '@/lib/trpc-route-handler';
5+
import { POST } from './route';
6+
7+
jest.mock('@/lib/trpc-route-handler', () => ({ handleTRPCRequest: jest.fn() }));
8+
9+
const mockedHandleTRPCRequest = jest.mocked(handleTRPCRequest);
10+
const listAvailableModels = jest.fn();
11+
12+
function makeModel(id: string): OpenRouterModel {
13+
return {
14+
id,
15+
name: id,
16+
created: 0,
17+
description: '',
18+
architecture: {
19+
input_modalities: ['text'],
20+
output_modalities: ['text'],
21+
tokenizer: 'test',
22+
},
23+
top_provider: { is_moderated: false },
24+
pricing: { prompt: '0', completion: '0' },
25+
context_length: 0,
26+
supported_parameters: ['tools'],
27+
};
28+
}
29+
30+
function request(modelId: string) {
31+
return new NextRequest('http://localhost:3000/api/organizations/org-1/models/validate', {
32+
method: 'POST',
33+
body: JSON.stringify({ modelId }),
34+
});
35+
}
36+
37+
describe('POST /api/organizations/[id]/models/validate', () => {
38+
beforeEach(() => {
39+
jest.resetAllMocks();
40+
listAvailableModels.mockResolvedValue({ data: [makeModel('available/model')] });
41+
mockedHandleTRPCRequest.mockImplementation(async (request, handler) => {
42+
const result = await handler({
43+
organizations: { settings: { listAvailableModels } },
44+
} as never);
45+
return NextResponse.json(result);
46+
});
47+
});
48+
49+
test('validates against the authorized organization catalog', async () => {
50+
const response = await POST(request('available/model'), {
51+
params: Promise.resolve({ id: 'org-1' }),
52+
});
53+
54+
expect(listAvailableModels).toHaveBeenCalledWith({ organizationId: 'org-1' });
55+
await expect(response.json()).resolves.toEqual({ valid: true });
56+
});
57+
58+
test('reports an organization-unavailable model without policy details', async () => {
59+
const response = await POST(request('missing/model'), {
60+
params: Promise.resolve({ id: 'org-1' }),
61+
});
62+
63+
await expect(response.json()).resolves.toEqual({ valid: false, reason: 'unavailable' });
64+
});
65+
66+
test('rejects an invalid body before invoking organization authorization', async () => {
67+
const response = await POST(
68+
new NextRequest('http://localhost:3000/api/organizations/org-1/models/validate', {
69+
method: 'POST',
70+
body: JSON.stringify({ modelId: '' }),
71+
}),
72+
{ params: Promise.resolve({ id: 'org-1' }) }
73+
);
74+
75+
expect(response.status).toBe(400);
76+
expect(mockedHandleTRPCRequest).not.toHaveBeenCalled();
77+
});
78+
});
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
import type { NextRequest } from 'next/server';
2+
import { NextResponse } from 'next/server';
3+
import * as z from 'zod';
4+
import { handleTRPCRequest } from '@/lib/trpc-route-handler';
5+
import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection';
6+
import { filterByFeature } from '@/lib/ai-gateway/models';
7+
8+
const BodySchema = z.object({ modelId: z.string().trim().min(1) });
9+
10+
type ValidationResult = { valid: true } | { valid: false; reason: 'unavailable' };
11+
12+
export async function POST(
13+
request: NextRequest,
14+
{ params }: { params: Promise<{ id: string }> }
15+
): Promise<NextResponse<{ error: string; message?: string } | ValidationResult>> {
16+
let body: unknown;
17+
try {
18+
body = await request.json();
19+
} catch {
20+
return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 });
21+
}
22+
23+
const bodyResult = BodySchema.safeParse(body);
24+
if (!bodyResult.success) {
25+
return NextResponse.json(
26+
{ error: 'Invalid request body', details: z.treeifyError(bodyResult.error) },
27+
{ status: 400 }
28+
);
29+
}
30+
31+
const organizationId = (await params).id;
32+
const feature = validateFeatureHeader(request.headers.get(FEATURE_HEADER));
33+
34+
return handleTRPCRequest<ValidationResult>(request, async caller => {
35+
const result = await caller.organizations.settings.listAvailableModels({ organizationId });
36+
const available = filterByFeature(result.data, feature).some(
37+
model => model.id === bodyResult.data.modelId
38+
);
39+
return available ? { valid: true } : { valid: false, reason: 'unavailable' };
40+
});
41+
}

services/cloud-agent-next/.dev.vars.example

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,9 @@ INTERNAL_API_SECRET=your-internal-api-secret-here
2424
# @url nextjs
2525
KILOCODE_BACKEND_BASE_URL=http://localhost:3000
2626
# KILO_OPENROUTER_BASE options:
27-
# Real LLM path (dev): http://host.docker.internal:3000/api
27+
# Real LLM path (dev): http://localhost:3000/api
2828
# Fake LLM (E2E harness, local dev only):
29-
# http://host.docker.internal:8811/api (add portOffset if set)
29+
# http://localhost:8811/api (add portOffset if set)
3030
# Start the `fake-llm` dev service before pointing at the fake URL.
3131
# @url nextjs/api
3232
KILO_OPENROUTER_BASE=http://localhost:3000/api

0 commit comments

Comments
 (0)