adding model limits

test/proxy.test.js

@@ -1,6 +1,6 @@
-import { describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import request from 'supertest';
-import { createProxyApp, unlimitedKeysFromEnv } from '../src/proxy.js';
+import { allowedModelsFromEnv, createProxyApp, DEFAULT_ALLOWED_MODELS, unlimitedKeysFromEnv } from '../src/proxy.js';
 
 const response = (body, options = {}) => new Response(body, {
   status: options.status ?? 200,
@@ -8,6 +8,17 @@ const response = (body, options = {}) => new Response(body, {
 });
 
 describe('proxy', () => {
+  const originalAllowedModels = process.env.ALLOWED_MODELS;
+
+  beforeEach(() => {
+    process.env.ALLOWED_MODELS = '';
+  });
+
+  afterEach(() => {
+    if (originalAllowedModels === undefined) delete process.env.ALLOWED_MODELS;
+    else process.env.ALLOWED_MODELS = originalAllowedModels;
+  });
+
   it('forwards models and rotates credentials', async () => {
     const fetchImpl = vi.fn()
       .mockResolvedValueOnce(response('{"object":"list","data":[]}'))
@@ -107,13 +118,65 @@ describe('proxy', () => {
     expect(() => unlimitedKeysFromEnv('[""]')).toThrow('UNLIMITED_API_KEYS must be a JSON array of non-empty strings');
   });
 
-  it('adds the global context limit to model listings', async () => {
-    const fetchImpl = vi.fn().mockResolvedValue(response('{"data":[{"id":"m"}]}'));
-    const app = createProxyApp({ keys: ['key'], fetchImpl, maxContextTokens: 123 });
+  it('validates the allowed models environment variable', () => {
+    expect(allowedModelsFromEnv(' model-a,model-b ')).toEqual(new Set(['model-a', 'model-b']));
+    expect(allowedModelsFromEnv('')).toBeNull();
+    expect(() => allowedModelsFromEnv('model-a,')).toThrow('ALLOWED_MODELS must be a comma-separated list of non-empty model IDs');
+  });
+
+  it('uses the built-in allowlist when ALLOWED_MODELS is unset', () => {
+    delete process.env.ALLOWED_MODELS;
+    expect(allowedModelsFromEnv()).toEqual(new Set(DEFAULT_ALLOWED_MODELS));
+  });
+
+  it('adds the global context limit and filters model listings for restricted clients', async () => {
+    const fetchImpl = vi.fn().mockResolvedValue(response('{"data":[{"id":"m"},{"id":"hidden"}]}'));
+    const app = createProxyApp({ keys: ['key'], fetchImpl, maxContextTokens: 123, allowedModels: new Set(['m']) });
     const result = await request(app).get('/oai/v1/models').expect(200);
+    expect(result.body.data).toHaveLength(1);
+    expect(result.body.data[0].id).toBe('m');
     expect(result.body.data[0]).toMatchObject({ context_length: 123, context_window: 123, max_context_tokens: 123 });
   });
 
+  it('enforces allowed models on every inference API and lets unlimited keys bypass it', async () => {
+    const fetchImpl = vi.fn().mockResolvedValue(response(JSON.stringify({
+      id: 'c', model: 'blocked', choices: [{ message: { role: 'assistant', content: 'ok' }, finish_reason: 'stop' }], usage: {},
+    })));
+    const app = createProxyApp({
+      keys: ['key'], fetchImpl, allowedModels: new Set(['allowed']), unlimitedKeys: new Set(['unlimited-key']),
+    });
+
+    await request(app).post('/oai/v1/chat/completions').send({ model: 'blocked' }).expect(400, {
+      error: { message: 'Model "blocked" is not allowed', type: 'invalid_request_error' },
+    });
+    await request(app).post('/oai/v1/responses').send({ model: 'blocked', input: 'Hi' }).expect(400, {
+      error: { message: 'Model "blocked" is not allowed', type: 'invalid_request_error' },
+    });
+    await request(app).post('/ant/v1/messages').send({ model: 'blocked', max_tokens: 1, messages: [] }).expect(400, {
+      type: 'error', error: { type: 'invalid_request_error', message: 'Model "blocked" is not allowed' },
+    });
+    expect(fetchImpl).not.toHaveBeenCalled();
+
+    await request(app)
+      .post('/oai/v1/chat/completions')
+      .set('Authorization', 'Bearer unlimited-key')
+      .send({ model: 'blocked' })
+      .expect(200);
+  });
+
+  it('lets unlimited keys see the complete model list', async () => {
+    const fetchImpl = vi.fn().mockImplementation(() => response('{"data":[{"id":"allowed"},{"id":"blocked"}]}'));
+    const app = createProxyApp({
+      keys: ['key'], fetchImpl, allowedModels: new Set(['allowed']), unlimitedKeys: new Set(['unlimited-key']),
+    });
+
+    const restricted = await request(app).get('/ant/v1/models').expect(200);
+    expect(restricted.body.data.map((model) => model.id)).toEqual(['allowed']);
+
+    const unlimited = await request(app).get('/oai/v1/models').set('x-api-key', 'unlimited-key').expect(200);
+    expect(unlimited.body.data.map((model) => model.id)).toEqual(['allowed', 'blocked']);
+  });
+
   it('passes through streaming responses', async () => {
     const fetchImpl = vi.fn().mockResolvedValue(response('data: {"delta":"Hi"}\n\ndata: [DONE]\n\n', {
       headers: { 'content-type': 'text/event-stream' },