From 7485dec33535b480b44fe8311a996b02a9b34d4f Mon Sep 17 00:00:00 2001 From: root Date: Mon, 9 Mar 2026 15:56:17 +0000 Subject: [PATCH] refactor(llm): remove daily limit gatekeeping from generationQueue Co-Authored-By: Claude Opus 4.6 --- .../src/llm/__tests__/generationQueue.test.ts | 73 +++---------------- server/src/llm/generationQueue.ts | 48 ++---------- 2 files changed, 16 insertions(+), 105 deletions(-) diff --git a/server/src/llm/__tests__/generationQueue.test.ts b/server/src/llm/__tests__/generationQueue.test.ts index ce1f5eb..2bab86f 100644 --- a/server/src/llm/__tests__/generationQueue.test.ts +++ b/server/src/llm/__tests__/generationQueue.test.ts @@ -43,8 +43,8 @@ describe('generationQueue', () => { }; const queue = createGenerationQueue(client, { requestsPerMinute: 60 }); - const p1 = queue.enqueue({ system: 's', user: 'u1' }).then(r => { if (r) results.push(r); }); - const p2 = queue.enqueue({ system: 's', user: 'u2' }).then(r => { if (r) results.push(r); }); + const p1 = queue.enqueue({ system: 's', user: 'u1' }).then(r => { if (r && typeof r === 'string') results.push(r); }); + const p2 = queue.enqueue({ system: 's', user: 'u2' }).then(r => { if (r && typeof r === 'string') results.push(r); }); await vi.advanceTimersByTimeAsync(2000); await Promise.all([p1, p2]); @@ -95,69 +95,18 @@ describe('generationQueue', () => { expect(queue.depth()).toBe(1); }); - it('exposes isDailyLimitReached as false initially', () => { - const client = mockClient('ok'); - const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 5 }); - expect(queue.isDailyLimitReached()).toBe(false); - }); - - it('resolves with null when daily limit is reached', async () => { - const client = mockClient('ok'); - const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 2 }); - - // First two requests should succeed - const p1 = queue.enqueue({ system: 's', user: '1' }); - await vi.advanceTimersByTimeAsync(100); - expect(await p1).toBe('ok'); - - const p2 = queue.enqueue({ system: 's', user: '2' }); - await vi.advanceTimersByTimeAsync(100); - expect(await p2).toBe('ok'); - - // Third request should be rejected (daily limit reached) - expect(queue.isDailyLimitReached()).toBe(true); - const p3 = queue.enqueue({ system: 's', user: '3' }); - const result = await p3; - expect(result).toBeNull(); - expect(client.complete).toHaveBeenCalledTimes(2); - }); - - it('resets daily count at midnight', async () => { - // Set time to 23:59:59 - const nearMidnight = new Date(); - nearMidnight.setHours(23, 59, 59, 0); - vi.setSystemTime(nearMidnight); - - const client = mockClient('ok'); - const queue = createGenerationQueue(client, { requestsPerMinute: 60, requestsPerDay: 1 }); - - // Use up the daily limit - const p1 = queue.enqueue({ system: 's', user: '1' }); - await vi.advanceTimersByTimeAsync(100); - expect(await p1).toBe('ok'); - expect(queue.isDailyLimitReached()).toBe(true); - - // Advance past midnight - await vi.advanceTimersByTimeAsync(2000); - - // Should be able to make requests again - expect(queue.isDailyLimitReached()).toBe(false); - const p2 = queue.enqueue({ system: 's', user: '2' }); - await vi.advanceTimersByTimeAsync(100); - expect(await p2).toBe('ok'); - }); - - it('works without requestsPerDay (no daily limit)', async () => { - const client = mockClient('ok'); + it('passes through rate limit info from client', async () => { + const rateLimitResult = { rateLimited: true as const, resetAt: 1741305600000 }; + const client: OpenRouterClient = { + complete: vi.fn().mockResolvedValue(rateLimitResult), + }; const queue = createGenerationQueue(client, { requestsPerMinute: 60 }); - expect(queue.isDailyLimitReached()).toBe(false); - - // Should process requests without daily limit - const p1 = queue.enqueue({ system: 's', user: '1' }); + const promise = queue.enqueue({ system: 's', user: 'u' }); await vi.advanceTimersByTimeAsync(100); - expect(await p1).toBe('ok'); - expect(queue.isDailyLimitReached()).toBe(false); + const result = await promise; + + expect(result).toEqual(rateLimitResult); }); it('can be cleared', async () => { diff --git a/server/src/llm/generationQueue.ts b/server/src/llm/generationQueue.ts index 4113eea..332be58 100644 --- a/server/src/llm/generationQueue.ts +++ b/server/src/llm/generationQueue.ts @@ -1,64 +1,31 @@ -import type { OpenRouterClient, CompletionRequest } from './openRouterClient.js'; +import type { OpenRouterClient, CompletionRequest, CompletionResult } from './openRouterClient.js'; interface QueueItem { request: CompletionRequest; - resolve: (value: string | null) => void; + resolve: (value: CompletionResult) => void; } export interface GenerationQueue { - enqueue(request: CompletionRequest): Promise; + enqueue(request: CompletionRequest): Promise; depth(): number; clear(): void; - isDailyLimitReached(): boolean; } export function createGenerationQueue( client: OpenRouterClient, - options: { requestsPerMinute: number; requestsPerDay?: number }, + options: { requestsPerMinute: number }, ): GenerationQueue { const queue: QueueItem[] = []; const intervalMs = Math.ceil(60000 / options.requestsPerMinute); let processing = false; - let dailyCount = 0; - let dailyResetTime = getNextMidnight(); - - function getNextMidnight(): number { - const now = new Date(); - const midnight = new Date(now); - midnight.setHours(24, 0, 0, 0); - return midnight.getTime(); - } - - function checkDailyReset(): void { - if (Date.now() >= dailyResetTime) { - dailyCount = 0; - dailyResetTime = getNextMidnight(); - } - } - - function isDailyLimitReached(): boolean { - checkDailyReset(); - return options.requestsPerDay != null && dailyCount >= options.requestsPerDay; - } - function scheduleNext(): void { if (processing || queue.length === 0) return; - if (isDailyLimitReached()) { - // Drain the queue with null responses - const remaining = queue.splice(0); - for (const item of remaining) { - item.resolve(null); - } - return; - } - processing = true; const item = queue.shift()!; client.complete(item.request).then(result => { - dailyCount++; item.resolve(result); }).catch(() => { item.resolve(null); @@ -71,10 +38,7 @@ export function createGenerationQueue( } return { - enqueue(request: CompletionRequest): Promise { - if (isDailyLimitReached()) { - return Promise.resolve(null); - } + enqueue(request: CompletionRequest): Promise { return new Promise(resolve => { queue.push({ request, resolve }); scheduleNext(); @@ -91,7 +55,5 @@ export function createGenerationQueue( item.resolve(null); } }, - - isDailyLimitReached, }; }