diff --git a/server/src/llm/__tests__/usageCounters.test.ts b/server/src/llm/__tests__/usageCounters.test.ts new file mode 100644 index 0000000..694c422 --- /dev/null +++ b/server/src/llm/__tests__/usageCounters.test.ts @@ -0,0 +1,56 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { createUsageCounters } from '../usageCounters.js'; + +describe('usageCounters', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('tracks requests per model', () => { + const counters = createUsageCounters(); + counters.record('free/model'); + counters.record('free/model'); + counters.record('paid/model'); + + const stats = counters.getStats(); + expect(stats['free/model'].total).toBe(2); + expect(stats['paid/model'].total).toBe(1); + }); + + it('resets daily counts at midnight UTC', () => { + // Set time to 23:59:59 UTC + const nearMidnight = new Date('2026-03-09T23:59:59Z'); + vi.setSystemTime(nearMidnight); + + const counters = createUsageCounters(); + counters.record('free/model'); + expect(counters.getStats()['free/model'].total).toBe(1); + + // Advance past midnight UTC + vi.advanceTimersByTime(2000); + counters.record('free/model'); + + // After reset, count should be 1 (just the new one) + expect(counters.getStats()['free/model'].total).toBe(1); + }); + + it('returns empty stats when no requests recorded', () => { + const counters = createUsageCounters(); + expect(counters.getStats()).toEqual({}); + }); + + it('formats a summary string', () => { + const counters = createUsageCounters(); + counters.record('free/model'); + counters.record('free/model'); + counters.record('paid/model'); + + const summary = counters.getSummary(); + expect(summary).toContain('free/model: 2'); + expect(summary).toContain('paid/model: 1'); + }); +}); diff --git a/server/src/llm/usageCounters.ts b/server/src/llm/usageCounters.ts new file mode 100644 index 0000000..16a4144 --- /dev/null +++ b/server/src/llm/usageCounters.ts @@ -0,0 +1,56 @@ +interface ModelStats { + total: number; +} + +export interface UsageCounters { + record(model: string): void; + getStats(): Record; + getSummary(): string; +} + +export function createUsageCounters(): UsageCounters { + let counts = new Map(); + let resetTime = getNextMidnightUTC(); + + function getNextMidnightUTC(): number { + const now = new Date(); + const midnight = new Date(Date.UTC( + now.getUTCFullYear(), now.getUTCMonth(), now.getUTCDate() + 1, + 0, 0, 0, 0, + )); + return midnight.getTime(); + } + + function checkReset(): void { + if (Date.now() >= resetTime) { + counts = new Map(); + resetTime = getNextMidnightUTC(); + } + } + + return { + record(model: string): void { + checkReset(); + counts.set(model, (counts.get(model) ?? 0) + 1); + }, + + getStats(): Record { + checkReset(); + const result: Record = {}; + for (const [model, total] of counts) { + result[model] = { total }; + } + return result; + }, + + getSummary(): string { + checkReset(); + if (counts.size === 0) return 'No LLM requests today'; + const parts: string[] = []; + for (const [model, total] of counts) { + parts.push(`${model}: ${total}`); + } + return `LLM usage today — ${parts.join(', ')}`; + }, + }; +}