diff --git a/server/src/config/__tests__/llmConfig.test.ts b/server/src/config/__tests__/llmConfig.test.ts index eb8749c..d9b09a3 100644 --- a/server/src/config/__tests__/llmConfig.test.ts +++ b/server/src/config/__tests__/llmConfig.test.ts @@ -17,7 +17,8 @@ describe('llmConfig', () => { expect(config.model).toBe('arcee-ai/trinity-large-preview:free'); expect(config.maxTokens).toBe(200); expect(config.temperature).toBe(0.8); - expect(config.requestsPerMinute).toBe(10); + expect(config.requestsPerMinute).toBe(20); + expect(config.requestsPerDay).toBe(1000); expect(config.timeoutMs).toBe(15000); expect(config.enabled).toBe(false); }); diff --git a/server/src/config/llmConfig.ts b/server/src/config/llmConfig.ts index 5dbd49e..d1a2d87 100644 --- a/server/src/config/llmConfig.ts +++ b/server/src/config/llmConfig.ts @@ -4,6 +4,7 @@ export interface LlmConfig { maxTokens: number; temperature: number; requestsPerMinute: number; + requestsPerDay: number; timeoutMs: number; enabled: boolean; } @@ -15,7 +16,8 @@ export function getLlmConfig(): LlmConfig { model: process.env.LLM_MODEL ?? 'arcee-ai/trinity-large-preview:free', maxTokens: 200, temperature: 0.8, - requestsPerMinute: 10, + requestsPerMinute: 20, + requestsPerDay: 1000, timeoutMs: 15000, enabled: apiKey.length > 0, }; diff --git a/server/src/llm/__tests__/generationQueue.test.ts b/server/src/llm/__tests__/generationQueue.test.ts index 5d0eb91..ce1f5eb 100644 --- a/server/src/llm/__tests__/generationQueue.test.ts +++ b/server/src/llm/__tests__/generationQueue.test.ts @@ -95,6 +95,71 @@ 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'); + 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' }); + await vi.advanceTimersByTimeAsync(100); + expect(await p1).toBe('ok'); + expect(queue.isDailyLimitReached()).toBe(false); + }); + it('can be cleared', async () => { const client = mockClient('ok'); const queue = createGenerationQueue(client, { requestsPerMinute: 6 }); diff --git a/server/src/llm/__tests__/llmService.test.ts b/server/src/llm/__tests__/llmService.test.ts index e5507d6..0e7048c 100644 --- a/server/src/llm/__tests__/llmService.test.ts +++ b/server/src/llm/__tests__/llmService.test.ts @@ -38,6 +38,25 @@ describe('llmService', () => { expect(result).toBe('A brave warrior from the north.'); }); + it('returns false for isDailyLimitReached when disabled', () => { + delete process.env.OPENROUTER_API_KEY; + const service = createLlmService(); + expect(service.isDailyLimitReached()).toBe(false); + }); + + it('exposes isDailyLimitReached when enabled', () => { + process.env.OPENROUTER_API_KEY = 'test-key'; + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ + choices: [{ message: { content: 'ok' } }], + }), + }); + + const service = createLlmService(); + expect(service.isDailyLimitReached()).toBe(false); + }); + it('exposes queue depth and clear', async () => { process.env.OPENROUTER_API_KEY = 'test-key'; globalThis.fetch = vi.fn().mockResolvedValue({ diff --git a/server/src/llm/generationQueue.ts b/server/src/llm/generationQueue.ts index 8b0c800..16388bb 100644 --- a/server/src/llm/generationQueue.ts +++ b/server/src/llm/generationQueue.ts @@ -9,24 +9,59 @@ export interface GenerationQueue { enqueue(request: CompletionRequest): Promise; depth(): number; clear(): void; + isDailyLimitReached(): boolean; } export function createGenerationQueue( client: OpenRouterClient, - options: { requestsPerMinute: number }, + options: { requestsPerMinute: number; requestsPerDay?: 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(() => { + dailyCount++; item.resolve(null); }).finally(() => { processing = false; @@ -38,6 +73,9 @@ export function createGenerationQueue( return { enqueue(request: CompletionRequest): Promise { + if (isDailyLimitReached()) { + return Promise.resolve(null); + } return new Promise(resolve => { queue.push({ request, resolve }); scheduleNext(); @@ -54,5 +92,7 @@ export function createGenerationQueue( item.resolve(null); } }, + + isDailyLimitReached, }; } diff --git a/server/src/llm/llmService.ts b/server/src/llm/llmService.ts index 074bc24..bd5eb5d 100644 --- a/server/src/llm/llmService.ts +++ b/server/src/llm/llmService.ts @@ -8,6 +8,7 @@ export interface LlmService { generate(templateName: string, variables: Record): Promise; queueDepth(): number; clear(): void; + isDailyLimitReached(): boolean; } export function createLlmService(): LlmService { @@ -18,12 +19,14 @@ export function createLlmService(): LlmService { generate: async () => null, queueDepth: () => 0, clear: () => {}, + isDailyLimitReached: () => false, }; } const client = createOpenRouterClient(config); const queue = createGenerationQueue(client, { requestsPerMinute: config.requestsPerMinute, + requestsPerDay: config.requestsPerDay, }); return { @@ -47,5 +50,9 @@ export function createLlmService(): LlmService { clear(): void { queue.clear(); }, + + isDailyLimitReached(): boolean { + return queue.isDailyLimitReached(); + }, }; }