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')
  })
})