PubScript/packages/ai/src/budget.test.ts

import { describe, expect, it } from "vitest";

import {
  allocateBudget,
  countTokens,
  DEFAULT_BUDGET,
  enforceBudget,
  measureBudget,
} from "./index.js";
import type { Message } from "./index.js";

describe("countTokens", () => {
  it("counts tokens for a known string", () => {
    // "hello world" is 2 words ≈ 3-4 tokens; assert a stable lower bound.
    const n = countTokens("hello world");
    expect(n).toBeGreaterThan(0);
    expect(Number.isInteger(n)).toBe(true);
  });

  it("returns 0 for an empty string", () => {
    expect(countTokens("")).toBe(0);
  });

  it("scales with input length", () => {
    const short = countTokens("a");
    const long = countTokens("a".repeat(1000));
    expect(long).toBeGreaterThan(short);
  });
});

describe("allocateBudget", () => {
  it("splits a 128k context window into per-slot caps", () => {
    const budget = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 128_000,
    });
    expect(budget.system).toBe(12_800);
    expect(budget.context).toBe(51_200);
    expect(budget.history).toBe(38_400);
    expect(budget.output).toBe(25_600);
    expect(budget.total).toBe(128_000);
  });

  it("scales the output budget with context window size", () => {
    const small = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 8_000,
    });
    const large = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 128_000,
    });
    expect(large.output).toBeGreaterThan(small.output);
  });

  it("never allocates more than half the context window to output", () => {
    const budget = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 16_000,
    });
    expect(budget.output).toBeLessThanOrEqual(16_000 / 2);
  });
});

describe("measureBudget", () => {
  it("breaks down tokens by dimension", () => {
    const messages: Message[] = [
      { role: "system", content: "You are a helpful editor." },
      { role: "user", content: "Summarize this chapter." },
    ];
    const measured = measureBudget(
      messages.find((m) => m.role === "system")!.content,
      "retrieved context",
      messages.filter((m) => m.role !== "system"),
      DEFAULT_BUDGET.output,
    );
    expect(measured.system).toBeGreaterThan(0);
    expect(measured.context).toBeGreaterThan(0);
    expect(measured.history).toBeGreaterThan(0);
    expect(measured.output).toBe(DEFAULT_BUDGET.output);
  });
});

describe("enforceBudget", () => {
  it("passes when all dimensions are within budget", () => {
    const budget = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 128_000,
    });
    const measured = {
      system: 100,
      context: 200,
      history: 300,
      output: 500,
    };
    expect(() => enforceBudget(measured, budget)).not.toThrow();
  });

  it("throws when the context dimension overflows", () => {
    const budget = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 8_000,
    });
    const measured = {
      system: 100,
      context: budget.context + 1,
      history: 100,
      output: 100,
    };
    expect(() => enforceBudget(measured, budget)).toThrow(/context/);
  });

  it("throws when the total overflows even if each dimension fits", () => {
    // Caps sum to more than the total, so each dimension can fit individually
    // while the aggregate still exceeds the context window.
    const budget = allocateBudget({
      system: 30,
      context: 30,
      history: 30,
      output: 30,
      total: 1_000,
    });
    const measured = { system: 300, context: 300, history: 300, output: 300 };
    expect(() => enforceBudget(measured, budget)).toThrow(/total/);
  });

  it("reports all overflowing dimensions", () => {
    const budget = allocateBudget({
      system: 10,
      context: 40,
      history: 30,
      output: 20,
      total: 8_000,
    });
    const measured = {
      system: budget.system + 1,
      context: budget.context + 1,
      history: 100,
      output: 100,
    };
    expect(() => enforceBudget(measured, budget)).toThrow(/system/);
    expect(() => enforceBudget(measured, budget)).toThrow(/context/);
  });
});