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"); }); });