import { describe, it, expect, beforeEach } from 'vitest'; import { eq } from 'drizzle-orm'; import { db } from '@/lib/db'; import { aiUsageLedger } from '@/lib/db/schema/ai-usage'; import { systemSettings } from '@/lib/db/schema/system'; import { checkBudget, currentPeriodTokens, getAiBudget, periodBreakdown, periodStart, recordAiUsage, setAiBudget, } from '@/lib/services/ai-budget.service'; import { makePort } from '../helpers/factories'; beforeEach(async () => { await db.delete(systemSettings).where(eq(systemSettings.key, 'ai.budget')); }); describe('ai-budget service', () => { it('defaults to disabled with sane caps', async () => { const port = await makePort(); const b = await getAiBudget(port.id); expect(b.enabled).toBe(false); expect(b.softCapTokens).toBeGreaterThan(0); expect(b.hardCapTokens).toBeGreaterThanOrEqual(b.softCapTokens); expect(b.period).toBe('month'); }); it('round-trips a custom budget and rejects soft > hard', async () => { const port = await makePort(); const saved = await setAiBudget( port.id, { enabled: true, softCapTokens: 10_000, hardCapTokens: 50_000, period: 'week' }, 'u1', ); expect(saved.enabled).toBe(true); expect(saved.softCapTokens).toBe(10_000); expect(saved.hardCapTokens).toBe(50_000); expect(saved.period).toBe('week'); await expect(setAiBudget(port.id, { softCapTokens: 60_000 }, 'u1')).rejects.toThrow( /soft.*cannot exceed hard/i, ); }); it('records usage and rolls into currentPeriodTokens', async () => { const port = await makePort(); await recordAiUsage({ portId: port.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 1000, outputTokens: 200, }); await recordAiUsage({ portId: port.id, feature: 'summary', provider: 'claude', model: 'claude-haiku-4-5', inputTokens: 500, outputTokens: 100, }); const total = await currentPeriodTokens(port.id); expect(total).toBe(1800); const ocrOnly = await currentPeriodTokens(port.id, 'ocr'); expect(ocrOnly).toBe(1200); const breakdown = await periodBreakdown(port.id); expect(breakdown.find((r) => r.feature === 'ocr')?.tokens).toBe(1200); expect(breakdown.find((r) => r.feature === 'summary')?.tokens).toBe(600); }); it('checkBudget passes when disabled regardless of usage', async () => { const port = await makePort(); // Stuff the ledger with a huge usage row. await recordAiUsage({ portId: port.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 10_000_000, outputTokens: 0, }); const r = await checkBudget({ portId: port.id, estimatedTokens: 5000 }); expect(r.ok).toBe(true); }); it('checkBudget refuses when hard cap is already exceeded', async () => { const port = await makePort(); await setAiBudget( port.id, { enabled: true, softCapTokens: 1000, hardCapTokens: 5000, period: 'month' }, 'u1', ); await recordAiUsage({ portId: port.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 6000, outputTokens: 0, }); const r = await checkBudget({ portId: port.id, estimatedTokens: 100 }); expect(r.ok).toBe(false); if (!r.ok) { expect(r.reason).toBe('hard-cap-exceeded'); } }); it('checkBudget refuses when estimated would push past the cap', async () => { const port = await makePort(); await setAiBudget( port.id, { enabled: true, softCapTokens: 1000, hardCapTokens: 5000, period: 'month' }, 'u1', ); await recordAiUsage({ portId: port.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 4500, outputTokens: 0, }); const r = await checkBudget({ portId: port.id, estimatedTokens: 1000 }); expect(r.ok).toBe(false); if (!r.ok) { expect(r.reason).toBe('estimated-exceeds-cap'); } }); it('checkBudget signals soft-cap warning when usage > soft but < hard', async () => { const port = await makePort(); await setAiBudget( port.id, { enabled: true, softCapTokens: 1000, hardCapTokens: 5000, period: 'month' }, 'u1', ); await recordAiUsage({ portId: port.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 1500, outputTokens: 0, }); const r = await checkBudget({ portId: port.id, estimatedTokens: 100 }); expect(r.ok).toBe(true); if (r.ok) { expect(r.softCap).toBe(true); } }); it('periodStart honors day/week/month boundaries (UTC)', () => { const wed = new Date(Date.UTC(2026, 3, 29, 14, 30)); // Wed 2026-04-29 14:30 UTC expect(periodStart('day', wed).toISOString()).toBe('2026-04-29T00:00:00.000Z'); // 2026-04-27 was a Monday — week starts there. expect(periodStart('week', wed).toISOString()).toBe('2026-04-27T00:00:00.000Z'); expect(periodStart('month', wed).toISOString()).toBe('2026-04-01T00:00:00.000Z'); }); it('isolates usage by port (cross-port rows do not leak into rollups)', async () => { const portA = await makePort(); const portB = await makePort(); await recordAiUsage({ portId: portA.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 100, outputTokens: 0, }); await recordAiUsage({ portId: portB.id, feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 999_999, outputTokens: 0, }); expect(await currentPeriodTokens(portA.id)).toBe(100); expect(await currentPeriodTokens(portB.id)).toBe(999_999); }); }); describe('ai-budget ledger writes never throw', () => { it('returns even when given a non-existent port (logs and continues)', async () => { // recordAiUsage swallows DB errors so the user-facing call doesn't fail // because the audit write hiccupped. await expect( recordAiUsage({ portId: 'nonexistent-port-id', feature: 'ocr', provider: 'openai', model: 'gpt-4o-mini', inputTokens: 1, outputTokens: 0, }), ).resolves.toBeUndefined(); // No row should have been inserted with a nonexistent portId due to FK. const rows = await db .select() .from(aiUsageLedger) .where(eq(aiUsageLedger.portId, 'nonexistent-port-id')); expect(rows).toHaveLength(0); }); });