pn-new-crm/tests/integration/ai-budget.test.ts

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);
  });
});