pangolin/server/lib/aiBudgetEnforcement.test.ts

import { execFileSync } from "child_process";
import fs from "fs";
import os from "os";
import path from "path";
import { fileURLToPath } from "url";
import { assertEquals } from "@test/assert";

const here = path.dirname(fileURLToPath(import.meta.url));
const repoRoot = path.resolve(here, "..", "..");

function tsconfigPath(): string {
    const tsconfig = ["tsconfig.json", "tsconfig.oss.json"]
        .map((file) => path.join(repoRoot, file))
        .find((file) => fs.existsSync(file));
    if (!tsconfig) {
        throw new Error("No tsconfig found for @server path aliases");
    }
    return tsconfig;
}

function runTsx(workdir: string, script: string): string {
    const tsxCli = path.join(
        repoRoot,
        "node_modules",
        "tsx",
        "dist",
        "cli.mjs"
    );
    if (!fs.existsSync(tsxCli)) {
        throw new Error("tsx is not installed; run npm ci first");
    }
    return execFileSync(
        process.execPath,
        [tsxCli, "--tsconfig", tsconfigPath(), script],
        {
            cwd: workdir,
            timeout: 120000,
            encoding: "utf8",
            env: { ...process.env, NODE_ENV: "test" }
        }
    );
}

function createTestEnvironment(): string {
    for (const generated of ["server/build.ts", "server/db/index.ts"]) {
        if (!fs.existsSync(path.join(repoRoot, generated))) {
            throw new Error(
                `Missing ${generated}; run npm run set:oss && npm run set:sqlite first`
            );
        }
    }
    const workdir = fs.mkdtempSync(
        path.join(os.tmpdir(), "pangolin-ai-budget-")
    );
    fs.mkdirSync(path.join(workdir, "config", "db"), { recursive: true });
    fs.copyFileSync(
        path.join(repoRoot, "config", "config.example.yml"),
        path.join(workdir, "config", "config.yml")
    );
    for (const linked of ["server", "test", "node_modules"]) {
        fs.symlinkSync(
            path.join(repoRoot, linked),
            path.join(workdir, linked),
            process.platform === "win32" ? "junction" : "dir"
        );
    }
    return workdir;
}

function testBudgetsUseRecordedUsageWindows() {
    console.log("Running AI budget window tests...");
    const workdir = createTestEnvironment();
    runTsx(workdir, path.join(repoRoot, "server", "db", "migrate.ts"));

    const output = runTsx(
        workdir,
        path.join(repoRoot, "server", "lib", "aiBudgetEnforcement.childTest.ts")
    );
    const match = output.match(/OBSERVATIONS (.*)/);
    if (!match) {
        throw new Error("child did not print observations: " + output);
    }
    const o = JSON.parse(match[1]);

    // Rolling windows over usage stored by recordUsage. These all summed to 0
    // while the lower bound was in ms and createdAt was in seconds.
    assertEquals(o.providerHourlyUsd, 3, "hourly budget must count last hour");
    assertEquals(o.providerDailyUsd, 7, "daily budget must count last 24h");
    assertEquals(o.providerWeeklyUsd, 17, "weekly budget must count last 7d");
    assertEquals(o.providerLifetimeUsd, 17, "lifetime budget must count all");
    assertEquals(o.providerDailyTokens, 750, "token budgets must use tokens");

    // Model budgets match the requested model in JS, and a request with no
    // known price must add 0 rather than poisoning the sum.
    assertEquals(o.wildcardHourlyUsd, 2, "wildcard model hourly usd");
    assertEquals(o.wildcardDailyUsd, 6, "wildcard model daily usd");
    assertEquals(o.wildcardDailyTokens, 650, "wildcard model daily tokens");
    assertEquals(o.exactModelDailyUsd, 1, "exact model daily usd");

    assertEquals(
        o.boundaryDailyUsd,
        1,
        "the window must include its cutoff and exclude the second before it"
    );

    // checkBudgets end to end.
    assertEquals(o.overHardBlocked, true, "over a hard budget must block");
    assertEquals(
        o.overHardBlockingBudgetMatches,
        true,
        "the blocking budget must be the breached one"
    );
    assertEquals(o.overHardBlockedAgain, true, "must keep blocking while over");
    assertEquals(
        o.overHardBreachRows,
        1,
        "a sustained breach must write one event (breach events are in ms)"
    );
    assertEquals(o.underHardBlocked, false, "under budget must not block");
    assertEquals(o.overSoftBlocked, false, "soft budgets must never block");
    assertEquals(o.overSoftBreachRows, 1, "soft breaches must be recorded");
    assertEquals(
        o.expiredBlocked,
        false,
        "usage older than the window must not count"
    );

    console.log("AI budget window tests passed");
}

function main() {
    try {
        testBudgetsUseRecordedUsageWindows();
    } catch (error) {
        console.error(`FAIL: ${(error as Error).message}`);
        process.exit(1);
    }
    console.log("aiBudgetEnforcement tests complete");
}

main();