maestro/src/db/repository.gateway-usage.test.ts

/**
 * AAO Gateway Phase 2b — gateway_key_usage repository tests.
 *
 * Coverage:
 *  - getGatewayKeyUsage returns null on miss, row on hit
 *  - incrementGatewayKeyUsage UPSERTs first call, accumulates on second
 *  - incrementGatewayKeyUsage clamps negative deltas to zero
 *  - listGatewayKeyUsagesByKey orders by period_start DESC
 *  - Cascade delete removes usage rows when key is deleted
 *  - tokens_budget / rate_limit_rpm fields persist & round-trip
 *  - updateGatewayVirtualKey patches fields independently and supports null reset
 */
import { describe, expect, it, beforeEach } from 'vitest';
import { Repository } from './repository.js';
import { hashKey } from '../gateway/key-format.js';

function makeRepo(): Repository {
  // ':memory:' triggers Repository.initSchema (creates all tables fresh).
  return new Repository(':memory:');
}

function seedKey(repo: Repository, raw: string, team = 'alpha'): string {
  return repo.createGatewayVirtualKey({
    keyHash: hashKey(raw),
    keyPrefix: raw.slice(0, 14),
    team,
  }).id;
}

describe('gateway_key_usage repository', () => {
  let repo: Repository;
  beforeEach(() => {
    repo = makeRepo();
  });

  it('getGatewayKeyUsage returns null when no row exists', () => {
    const id = seedKey(repo, 'sk-aao-test-1');
    expect(repo.getGatewayKeyUsage(id, '2026-05')).toBeNull();
  });

  it('incrementGatewayKeyUsage creates row on first call, accumulates on second', () => {
    const id = seedKey(repo, 'sk-aao-test-2');
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 100, tokensOut: 50, requests: 1 });
    let usage = repo.getGatewayKeyUsage(id, '2026-05');
    expect(usage).not.toBeNull();
    expect(usage!.tokensIn).toBe(100);
    expect(usage!.tokensOut).toBe(50);
    expect(usage!.requests).toBe(1);

    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 200, tokensOut: 80, requests: 1 });
    usage = repo.getGatewayKeyUsage(id, '2026-05');
    expect(usage!.tokensIn).toBe(300);
    expect(usage!.tokensOut).toBe(130);
    expect(usage!.requests).toBe(2);
  });

  it('different periods get separate rows', () => {
    const id = seedKey(repo, 'sk-aao-test-3');
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-04', tokensIn: 10, requests: 1 });
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 20, requests: 1 });
    expect(repo.getGatewayKeyUsage(id, '2026-04')!.tokensIn).toBe(10);
    expect(repo.getGatewayKeyUsage(id, '2026-05')!.tokensIn).toBe(20);
  });

  it('clamps negative deltas to zero', () => {
    const id = seedKey(repo, 'sk-aao-test-4');
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 100 });
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: -500, tokensOut: -10, requests: -1 });
    const usage = repo.getGatewayKeyUsage(id, '2026-05')!;
    expect(usage.tokensIn).toBe(100); // unchanged
    expect(usage.tokensOut).toBe(0);
    expect(usage.requests).toBe(0);
  });

  it('listGatewayKeyUsagesByKey returns rows newest period first', () => {
    const id = seedKey(repo, 'sk-aao-test-5');
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-03', requests: 1 });
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', requests: 1 });
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-04', requests: 1 });
    const list = repo.listGatewayKeyUsagesByKey(id);
    expect(list.map(u => u.periodStart)).toEqual(['2026-05', '2026-04', '2026-03']);
  });

  it('listGatewayKeyUsagesByKey honors limit', () => {
    const id = seedKey(repo, 'sk-aao-test-6');
    for (const m of ['2026-01', '2026-02', '2026-03', '2026-04']) {
      repo.incrementGatewayKeyUsage({ keyId: id, period: m, requests: 1 });
    }
    const list = repo.listGatewayKeyUsagesByKey(id, { limit: 2 });
    expect(list).toHaveLength(2);
    expect(list[0]!.periodStart).toBe('2026-04');
  });

  it('cascade-deletes usage rows when key is hard-deleted', () => {
    const id = seedKey(repo, 'sk-aao-test-7');
    repo.incrementGatewayKeyUsage({ keyId: id, period: '2026-05', tokensIn: 1 });
    expect(repo.getGatewayKeyUsage(id, '2026-05')).not.toBeNull();

    // The Repository's defense-in-depth guard refuses to delete
    // config-import rows. Seeded above as 'admin' so this is legal.
    repo.deleteGatewayVirtualKey(id);
    expect(repo.getGatewayKeyUsage(id, '2026-05')).toBeNull();
  });
});

describe('gateway_virtual_keys budget/rate fields', () => {
  let repo: Repository;
  beforeEach(() => {
    repo = makeRepo();
  });

  it('round-trips tokensBudget and rateLimitRpm on create', () => {
    const created = repo.createGatewayVirtualKey({
      keyHash: hashKey('sk-aao-create-budget'),
      keyPrefix: 'sk-aao-create',
      team: 'alpha',
      tokensBudget: 1_000_000,
      rateLimitRpm: 60,
    });
    expect(created.tokensBudget).toBe(1_000_000);
    expect(created.rateLimitRpm).toBe(60);

    const refreshed = repo.findGatewayVirtualKeyById(created.id)!;
    expect(refreshed.tokensBudget).toBe(1_000_000);
    expect(refreshed.rateLimitRpm).toBe(60);
  });

  it('defaults to null when budget/rate omitted', () => {
    const created = repo.createGatewayVirtualKey({
      keyHash: hashKey('sk-aao-no-limits'),
      keyPrefix: 'sk-aao-no-lim',
      team: 'beta',
    });
    expect(created.tokensBudget).toBeNull();
    expect(created.rateLimitRpm).toBeNull();
  });

  it('coerces zero or negative limits to null (defensive)', () => {
    const created = repo.createGatewayVirtualKey({
      keyHash: hashKey('sk-aao-bad-limits'),
      keyPrefix: 'sk-aao-bad',
      team: 'gamma',
      tokensBudget: 0,
      rateLimitRpm: -5,
    });
    expect(created.tokensBudget).toBeNull();
    expect(created.rateLimitRpm).toBeNull();
  });

  it('updateGatewayVirtualKey patches only specified fields', () => {
    const created = repo.createGatewayVirtualKey({
      keyHash: hashKey('sk-aao-patch'),
      keyPrefix: 'sk-aao-patch',
      team: 'alpha',
      tokensBudget: 1000,
      rateLimitRpm: 60,
      allowedModels: ['qwen3:8b'],
    });
    // Patch only tokensBudget — other fields untouched
    const after1 = repo.updateGatewayVirtualKey(created.id, { tokensBudget: 5000 });
    expect(after1.tokensBudget).toBe(5000);
    expect(after1.rateLimitRpm).toBe(60);
    expect(after1.allowedModels).toEqual(['qwen3:8b']);

    // Patch allowedModels alone
    const after2 = repo.updateGatewayVirtualKey(created.id, { allowedModels: ['qwen3:14b'] });
    expect(after2.allowedModels).toEqual(['qwen3:14b']);
    expect(after2.tokensBudget).toBe(5000);

    // Reset rate limit to null (unlimited)
    const after3 = repo.updateGatewayVirtualKey(created.id, { rateLimitRpm: null });
    expect(after3.rateLimitRpm).toBeNull();

    // Reset allowedModels to null (no allowlist)
    const after4 = repo.updateGatewayVirtualKey(created.id, { allowedModels: null });
    expect(after4.allowedModels).toBeNull();
  });

  it('updateGatewayVirtualKey throws for unknown id', () => {
    expect(() => repo.updateGatewayVirtualKey('does-not-exist', { tokensBudget: 1 })).toThrow();
  });
});