SRCmail/lib/ai/__tests__/entitlement.test.ts

import { describe, expect, it, beforeEach, afterEach, vi } from 'vitest';
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';

// Real end-to-end seat assignment against the real Ollama was verified live
// (see the commit this test ships with); this covers the rejection branch,
// which is deterministic and cheaper to prove with a unit test than another
// live round trip.

describe('lib/ai/entitlement', () => {
  let stateDir: string;

  beforeEach(async () => {
    vi.resetModules();
    stateDir = await mkdtemp(path.join(tmpdir(), 'ai-entitlement-test-'));
    process.env.ADMIN_STATE_DIR = stateDir;
    delete process.env.AI_SERVER_SEAT_TOTAL;
    // Each test needs a fresh globalThis singleton, not just a fresh module -
    // the module stashes cached state on globalThis specifically to survive
    // HMR, so resetModules() alone doesn't clear it.
    delete (globalThis as Record<symbol, unknown>)[Symbol.for('vncmail.ai.entitlement')];
  });

  afterEach(async () => {
    delete process.env.ADMIN_STATE_DIR;
    await rm(stateDir, { recursive: true, force: true });
  });

  it('assigns a seat on first use and allows the same user again', async () => {
    const { checkAndAssignSeat, setSeatTotal } = await import('../entitlement');
    await setSeatTotal(1);

    const first = await checkAndAssignSeat('alice@example.com');
    expect(first).toEqual({ allowed: true, seatJustAssigned: true });

    const second = await checkAndAssignSeat('alice@example.com');
    expect(second).toEqual({ allowed: true });
  });

  it('rejects a new user once all seats are assigned', async () => {
    const { checkAndAssignSeat, setSeatTotal } = await import('../entitlement');
    await setSeatTotal(1);
    await checkAndAssignSeat('alice@example.com');

    const rejected = await checkAndAssignSeat('bob@example.com');
    expect(rejected.allowed).toBe(false);
    expect(rejected.reason).toMatch(/already assigned/i);
  });

  it('rejects everyone when no seats are configured', async () => {
    const { checkAndAssignSeat } = await import('../entitlement');
    const result = await checkAndAssignSeat('anyone@example.com');
    expect(result.allowed).toBe(false);
    expect(result.reason).toMatch(/no licensed seats/i);
  });

  it('revoking a seat frees it for someone else', async () => {
    const { checkAndAssignSeat, setSeatTotal, revokeSeat } = await import('../entitlement');
    await setSeatTotal(1);
    await checkAndAssignSeat('alice@example.com');
    await revokeSeat('alice@example.com');

    const result = await checkAndAssignSeat('bob@example.com');
    expect(result).toEqual({ allowed: true, seatJustAssigned: true });
  });

  it('persists usage to the metering ledger, append-only', async () => {
    const { recordUsage, readMeteringLedger } = await import('../entitlement');
    await recordUsage({
      timestamp: new Date(0).toISOString(),
      username: 'alice@example.com',
      model: 'qwen2.5:32b',
      promptTokens: 10,
      completionTokens: 5,
      latencyMs: 123,
    });
    await recordUsage({
      timestamp: new Date(0).toISOString(),
      username: 'alice@example.com',
      model: 'qwen2.5:32b',
      promptTokens: 8,
      completionTokens: 3,
      latencyMs: 90,
    });

    const ledger = await readMeteringLedger();
    expect(ledger).toHaveLength(2);
    expect(ledger[0].promptTokens).toBe(10);
    expect(ledger[1].promptTokens).toBe(8);
  });
});