claude-code/src/utils/__tests__/tokenBudget.test.ts

import { describe, expect, test } from "bun:test";
import {
  parseTokenBudget,
  findTokenBudgetPositions,
  getBudgetContinuationMessage,
} from "../tokenBudget";

describe("parseTokenBudget", () => {
  // --- shorthand at start ---
  test("parses +500k at start", () => {
    expect(parseTokenBudget("+500k")).toBe(500_000);
  });

  test("parses +2.5M at start", () => {
    expect(parseTokenBudget("+2.5M")).toBe(2_500_000);
  });

  test("parses +1b at start", () => {
    expect(parseTokenBudget("+1b")).toBe(1_000_000_000);
  });

  test("parses shorthand with leading whitespace", () => {
    expect(parseTokenBudget("  +500k")).toBe(500_000);
  });

  // --- shorthand at end ---
  test("parses +1.5m at end of sentence", () => {
    expect(parseTokenBudget("do this +1.5m")).toBe(1_500_000);
  });

  test("parses shorthand at end with trailing period", () => {
    expect(parseTokenBudget("please continue +100k.")).toBe(100_000);
  });

  test("parses shorthand at end with trailing whitespace", () => {
    expect(parseTokenBudget("keep going +250k  ")).toBe(250_000);
  });

  // --- verbose ---
  test("parses 'use 2M tokens'", () => {
    expect(parseTokenBudget("use 2M tokens")).toBe(2_000_000);
  });

  test("parses 'spend 500k tokens'", () => {
    expect(parseTokenBudget("spend 500k tokens")).toBe(500_000);
  });

  test("parses verbose with singular 'token'", () => {
    expect(parseTokenBudget("use 1k token")).toBe(1_000);
  });

  test("parses verbose embedded in sentence", () => {
    expect(parseTokenBudget("please use 3.5m tokens for this task")).toBe(
      3_500_000
    );
  });

  // --- no match (returns null) ---
  test("returns null for plain text", () => {
    expect(parseTokenBudget("hello world")).toBeNull();
  });

  test("returns null for bare number without +", () => {
    expect(parseTokenBudget("500k")).toBeNull();
  });

  test("returns null for number without suffix", () => {
    expect(parseTokenBudget("+500")).toBeNull();
  });

  test("returns null for empty string", () => {
    expect(parseTokenBudget("")).toBeNull();
  });

  // --- case insensitivity ---
  test("is case insensitive for suffix", () => {
    expect(parseTokenBudget("+500K")).toBe(500_000);
    expect(parseTokenBudget("+2m")).toBe(2_000_000);
    expect(parseTokenBudget("+1B")).toBe(1_000_000_000);
  });

  // --- priority: start shorthand wins over end/verbose ---
  test("start shorthand takes priority over verbose in same text", () => {
    expect(parseTokenBudget("+100k use 2M tokens")).toBe(100_000);
  });
});

describe("findTokenBudgetPositions", () => {
  test("returns single position for +500k at start", () => {
    const positions = findTokenBudgetPositions("+500k");
    expect(positions).toHaveLength(1);
    expect(positions[0]!.start).toBe(0);
    expect(positions[0]!.end).toBe(5);
  });

  test("returns position for shorthand at end", () => {
    const text = "do this +100k";
    const positions = findTokenBudgetPositions(text);
    expect(positions).toHaveLength(1);
    expect(positions[0]!.start).toBe(8);
    expect(text.slice(positions[0]!.start, positions[0]!.end)).toBe("+100k");
  });

  test("returns position for verbose match", () => {
    const text = "please use 2M tokens here";
    const positions = findTokenBudgetPositions(text);
    expect(positions).toHaveLength(1);
    expect(text.slice(positions[0]!.start, positions[0]!.end)).toBe(
      "use 2M tokens"
    );
  });

  test("returns multiple positions for combined shorthand + verbose", () => {
    const text = "use 2M tokens and then +500k";
    const positions = findTokenBudgetPositions(text);
    expect(positions.length).toBeGreaterThanOrEqual(2);
  });

  test("returns empty array for no match", () => {
    expect(findTokenBudgetPositions("hello world")).toEqual([]);
  });

  test("does not double-count when +500k matches both start and end", () => {
    const positions = findTokenBudgetPositions("+500k");
    expect(positions).toHaveLength(1);
  });
});

describe("getBudgetContinuationMessage", () => {
  test("formats a continuation message with correct values", () => {
    const msg = getBudgetContinuationMessage(50, 250_000, 500_000);
    expect(msg).toContain("50%");
    expect(msg).toContain("250,000");
    expect(msg).toContain("500,000");
    expect(msg).toContain("Keep working");
    expect(msg).toContain("do not summarize");
  });

  test("formats zero values", () => {
    const msg = getBudgetContinuationMessage(0, 0, 100_000);
    expect(msg).toContain("0%");
    expect(msg).toContain("0 / 100,000");
  });

  test("formats large numbers with commas", () => {
    const msg = getBudgetContinuationMessage(75, 7_500_000, 10_000_000);
    expect(msg).toContain("7,500,000");
    expect(msg).toContain("10,000,000");
  });
});