Skip to content

Commit ad2976a

Browse files
eshurakoveshurakov
committed
feat(cloud-agent-next): reject unavailable models before admission
1 parent 340cf4b commit ad2976a

27 files changed

Lines changed: 1420 additions & 112 deletions

‎apps/web/src/app/api/openrouter/models/validate/route.test.ts‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
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 { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models';
8+
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
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/organizations/organization-models', () => ({
20+
getAvailableModelsForOrganization: jest.fn(),
21+
}));
22+
jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({
23+
listAvailableExperimentModels: jest.fn(),
24+
}));
25+
26+
const mockedGetUserFromAuth = jest.mocked(getUserFromAuth);
27+
const mockedGetEnhancedOpenRouterModels = jest.mocked(getEnhancedOpenRouterModels);
28+
const mockedGetDirectByokModelsForUser = jest.mocked(getDirectByokModelsForUser);
29+
const mockedGetAvailableModelsForOrganization = jest.mocked(getAvailableModelsForOrganization);
30+
const mockedListAvailableExperimentModels = jest.mocked(listAvailableExperimentModels);
31+
32+
function makeModel(id: string): OpenRouterModel {
33+
return {
34+
id,
35+
name: id,
36+
created: 0,
37+
description: '',
38+
architecture: {
39+
input_modalities: ['text'],
40+
output_modalities: ['text'],
41+
tokenizer: 'test',
42+
},
43+
top_provider: { is_moderated: false },
44+
pricing: { prompt: '0', completion: '0' },
45+
context_length: 0,
46+
supported_parameters: ['tools'],
47+
};
48+
}
49+
50+
function request(modelId: string) {
51+
return new NextRequest('http://localhost:3000/api/openrouter/models/validate', {
52+
method: 'POST',
53+
body: JSON.stringify({ modelId }),
54+
});
55+
}
56+
57+
describe('POST /api/openrouter/models/validate', () => {
58+
beforeEach(() => {
59+
jest.resetAllMocks();
60+
mockedGetUserFromAuth.mockResolvedValue({
61+
user: null,
62+
organizationId: null,
63+
authFailedResponse: null,
64+
} as never);
65+
mockedGetEnhancedOpenRouterModels.mockResolvedValue({ data: [makeModel('available/model')] });
66+
mockedGetDirectByokModelsForUser.mockResolvedValue([]);
67+
mockedGetAvailableModelsForOrganization.mockResolvedValue(null);
68+
mockedListAvailableExperimentModels.mockResolvedValue([]);
69+
});
70+
71+
test('confirms a Kilo-eligible catalog model', async () => {
72+
const response = await POST(request('available/model'));
73+
74+
expect(response.status).toBe(200);
75+
await expect(response.json()).resolves.toEqual({ valid: true });
76+
});
77+
78+
test('does not expose details for an unavailable model', async () => {
79+
const response = await POST(request('missing/model'));
80+
81+
expect(response.status).toBe(200);
82+
await expect(response.json()).resolves.toEqual({ valid: false, reason: 'unavailable' });
83+
});
84+
85+
test('uses the public catalog after failed optional authentication', async () => {
86+
mockedGetUserFromAuth.mockResolvedValue({
87+
user: null,
88+
organizationId: null,
89+
authFailedResponse: NextResponse.json({ error: 'Unauthorized' }, { status: 401 }),
90+
} as never);
91+
92+
const response = await POST(request('available/model'));
93+
94+
expect(response.status).toBe(200);
95+
await expect(response.json()).resolves.toEqual({ valid: true });
96+
});
97+
98+
test('returns a service failure when catalog construction fails', async () => {
99+
mockedGetEnhancedOpenRouterModels.mockRejectedValue(new Error('catalog unavailable'));
100+
101+
const response = await POST(request('available/model'));
102+
103+
expect(response.status).toBe(500);
104+
await expect(response.json()).resolves.toEqual({
105+
error: 'Failed to validate model',
106+
message: 'Error from model catalog',
107+
});
108+
});
109+
110+
test('rejects an invalid body without reading a catalog', async () => {
111+
const response = await POST(
112+
new NextRequest('http://localhost:3000/api/openrouter/models/validate', {
113+
method: 'POST',
114+
body: JSON.stringify({ modelId: '' }),
115+
})
116+
);
117+
118+
expect(response.status).toBe(400);
119+
expect(mockedGetEnhancedOpenRouterModels).not.toHaveBeenCalled();
120+
});
121+
});

‎apps/web/src/app/api/openrouter/models/validate/route.ts‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
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 { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models';
9+
import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection';
10+
import { filterByFeature } from '@/lib/ai-gateway/models';
11+
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
12+
import { isKiloAgentModelAvailable } from '@/lib/ai-gateway/validate-kilo-agent-model.server';
13+
14+
const BodySchema = z.object({ modelId: z.string().trim().min(1) });
15+
16+
async function tryGetUserFromAuth() {
17+
try {
18+
return await getUserFromAuth({ adminOnly: false });
19+
} catch (error) {
20+
console.error('[validateOpenRouterModel] failed to get user from auth', error);
21+
return { user: null, organizationId: null };
22+
}
23+
}
24+
25+
export async function POST(request: NextRequest) {
26+
let body: unknown;
27+
try {
28+
body = await request.json();
29+
} catch {
30+
return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 });
31+
}
32+
33+
const bodyResult = BodySchema.safeParse(body);
34+
if (!bodyResult.success) {
35+
return NextResponse.json(
36+
{ error: 'Invalid request body', details: z.treeifyError(bodyResult.error) },
37+
{ status: 400 }
38+
);
39+
}
40+
41+
const feature = validateFeatureHeader(request.headers.get(FEATURE_HEADER));
42+
const auth = await tryGetUserFromAuth();
43+
try {
44+
const organizationModels = auth?.organizationId
45+
? await getAvailableModelsForOrganization(auth.organizationId)
46+
: null;
47+
if (organizationModels) {
48+
const available = isKiloAgentModelAvailable(
49+
bodyResult.data.modelId,
50+
filterByFeature(organizationModels.data, feature)
51+
);
52+
return NextResponse.json(
53+
available ? { valid: true } : { valid: false, reason: 'unavailable' }
54+
);
55+
}
56+
57+
const models = await getEnhancedOpenRouterModels();
58+
if (!Array.isArray(models.data)) {
59+
throw new Error('Model catalog returned invalid data');
60+
}
61+
const byokModels = auth?.user ? await getDirectByokModelsForUser(auth.user.id) : [];
62+
const experimentModels = await listAvailableExperimentModels();
63+
const available = isKiloAgentModelAvailable(
64+
bodyResult.data.modelId,
65+
filterByFeature(models.data.concat(byokModels, experimentModels), feature)
66+
);
67+
return NextResponse.json(available ? { valid: true } : { valid: false, reason: 'unavailable' });
68+
} catch (error) {
69+
captureException(error, {
70+
tags: { endpoint: 'openrouter/models/validate' },
71+
extra: {
72+
action: 'validating_model',
73+
userId: auth?.user?.id,
74+
organizationId: auth?.organizationId,
75+
},
76+
});
77+
return NextResponse.json(
78+
{ error: 'Failed to validate model', message: 'Error from model catalog' },
79+
{ status: 500 }
80+
);
81+
}
82+
}

‎apps/web/src/app/api/organizations/[id]/models/validate/route.test.ts‎

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+
});

‎apps/web/src/app/api/organizations/[id]/models/validate/route.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
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+
import { isKiloAgentModelAvailable } from '@/lib/ai-gateway/validate-kilo-agent-model.server';
8+
9+
const BodySchema = z.object({ modelId: z.string().trim().min(1) });
10+
11+
type ValidationResult = { valid: true } | { valid: false; reason: 'unavailable' };
12+
13+
export async function POST(
14+
request: NextRequest,
15+
{ params }: { params: Promise<{ id: string }> }
16+
): Promise<NextResponse<{ error: string; message?: string } | ValidationResult>> {
17+
let body: unknown;
18+
try {
19+
body = await request.json();
20+
} catch {
21+
return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 });
22+
}
23+
24+
const bodyResult = BodySchema.safeParse(body);
25+
if (!bodyResult.success) {
26+
return NextResponse.json(
27+
{ error: 'Invalid request body', details: z.treeifyError(bodyResult.error) },
28+
{ status: 400 }
29+
);
30+
}
31+
32+
const organizationId = (await params).id;
33+
const feature = validateFeatureHeader(request.headers.get(FEATURE_HEADER));
34+
35+
return handleTRPCRequest<ValidationResult>(request, async caller => {
36+
const result = await caller.organizations.settings.listAvailableModels({ organizationId });
37+
const available = isKiloAgentModelAvailable(
38+
bodyResult.data.modelId,
39+
filterByFeature(result.data, feature)
40+
);
41+
return available ? { valid: true } : { valid: false, reason: 'unavailable' };
42+
});
43+
}

‎apps/web/src/lib/ai-gateway/validate-kilo-agent-model.server.test.ts‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import { describe, expect, test } from '@jest/globals';
2+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
3+
import { isKiloAgentModelAvailable } from './validate-kilo-agent-model.server';
4+
5+
function makeModel(
6+
id: string,
7+
options: { supportedParameters?: string[]; outputModalities?: string[] } = {}
8+
): OpenRouterModel {
9+
return {
10+
id,
11+
name: id,
12+
created: 0,
13+
description: '',
14+
architecture: {
15+
input_modalities: ['text'],
16+
output_modalities: options.outputModalities ?? ['text'],
17+
tokenizer: 'test',
18+
},
19+
top_provider: {
20+
is_moderated: false,
21+
context_length: null,
22+
max_completion_tokens: null,
23+
},
24+
pricing: {
25+
prompt: '0',
26+
completion: '0',
27+
},
28+
context_length: 0,
29+
supported_parameters: options.supportedParameters ?? ['tools'],
30+
};
31+
}
32+
33+
describe('isKiloAgentModelAvailable', () => {
34+
test('accepts an exact listed model', () => {
35+
expect(
36+
isKiloAgentModelAvailable('anthropic/claude-sonnet', [makeModel('anthropic/claude-sonnet')])
37+
).toBe(true);
38+
});
39+
40+
test('rejects a model that is not listed exactly', () => {
41+
expect(
42+
isKiloAgentModelAvailable('anthropic/claude', [makeModel('anthropic/claude-sonnet')])
43+
).toBe(false);
44+
});
45+
46+
test('accepts a listed model regardless of advertised tool capability', () => {
47+
expect(
48+
isKiloAgentModelAvailable('anthropic/claude-sonnet', [
49+
makeModel('anthropic/claude-sonnet', { supportedParameters: ['temperature'] }),
50+
])
51+
).toBe(true);
52+
});
53+
54+
test('accepts a listed model regardless of output modalities', () => {
55+
expect(
56+
isKiloAgentModelAvailable('image/model', [
57+
makeModel('image/model', { outputModalities: ['text', 'image'] }),
58+
])
59+
).toBe(true);
60+
});
61+
});

‎apps/web/src/lib/ai-gateway/validate-kilo-agent-model.server.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
2+
3+
export function isKiloAgentModelAvailable(modelId: string, models: OpenRouterModel[]): boolean {
4+
return models.some(model => model.id === modelId);
5+
}

0 commit comments

Comments
 (0)