Adding unlimited key

This commit is contained in:
2026-07-21 15:07:08 -05:00
parent b9b183df40
commit fe638a14d4
4 changed files with 97 additions and 15 deletions
+39 -1
View File
@@ -1,6 +1,6 @@
import { describe, expect, it, vi } from 'vitest';
import request from 'supertest';
import { createProxyApp } from '../src/proxy.js';
import { createProxyApp, unlimitedKeysFromEnv } from '../src/proxy.js';
const response = (body, options = {}) => new Response(body, {
status: options.status ?? 200,
@@ -69,6 +69,44 @@ describe('proxy', () => {
await request(app).post('/oai/v1/chat/completions').set('X-Real-IP', '198.51.100.1').send({ model: 'm', max_tokens: 1 }).expect(429);
});
it('bypasses rate checks and spend accounting for configured unlimited keys', async () => {
const fetchImpl = vi.fn().mockResolvedValue(response(JSON.stringify({ id: 'c', cost: 1 })));
const rateLimiter = { check: vi.fn(() => ({ allowed: false, retryAfter: 1, reason: 'rate' })), recordCost: vi.fn() };
const app = createProxyApp({ keys: ['key'], fetchImpl, rateLimiter, unlimitedKeys: new Set(['unlimited-key']) });
await request(app)
.post('/oai/v1/chat/completions')
.set('Authorization', 'Bearer unlimited-key')
.send({ model: 'm' })
.expect(200);
expect(rateLimiter.check).not.toHaveBeenCalled();
expect(rateLimiter.recordCost).not.toHaveBeenCalled();
});
it('recognizes x-api-key unlimited keys for Anthropic requests', async () => {
const fetchImpl = vi.fn().mockResolvedValue(response(JSON.stringify({
id: 'c', model: 'm', choices: [{ message: { role: 'assistant', content: 'ok' }, finish_reason: 'stop' }], cost: 1,
})));
const rateLimiter = { check: vi.fn(() => ({ allowed: false, retryAfter: 1, reason: 'rate' })), recordCost: vi.fn() };
const app = createProxyApp({ keys: ['key'], fetchImpl, rateLimiter, unlimitedKeys: new Set(['unlimited-key']) });
await request(app)
.post('/ant/v1/messages')
.set('x-api-key', 'unlimited-key')
.send({ model: 'm', max_tokens: 10, messages: [{ role: 'user', content: 'Hi' }] })
.expect(200);
expect(rateLimiter.check).not.toHaveBeenCalled();
expect(rateLimiter.recordCost).not.toHaveBeenCalled();
});
it('validates the unlimited API key environment variable', () => {
expect(unlimitedKeysFromEnv('["key-a", "key-b"]')).toEqual(new Set(['key-a', 'key-b']));
expect(() => unlimitedKeysFromEnv('key-a')).toThrow('UNLIMITED_API_KEYS must be valid JSON');
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 });