/ utils / tokenBudget.ts
tokenBudget.ts
 1  // Shorthand (+500k) anchored to start/end to avoid false positives in natural language.
 2  // Verbose (use/spend 2M tokens) matches anywhere.
 3  const SHORTHAND_START_RE = /^\s*\+(\d+(?:\.\d+)?)\s*(k|m|b)\b/i
 4  // Lookbehind (?<=\s) is avoided — it defeats YARR JIT in JSC, and the
 5  // interpreter scans O(n) even with the $ anchor. Capture the whitespace
 6  // instead; callers offset match.index by 1 where position matters.
 7  const SHORTHAND_END_RE = /\s\+(\d+(?:\.\d+)?)\s*(k|m|b)\s*[.!?]?\s*$/i
 8  const VERBOSE_RE = /\b(?:use|spend)\s+(\d+(?:\.\d+)?)\s*(k|m|b)\s*tokens?\b/i
 9  const VERBOSE_RE_G = new RegExp(VERBOSE_RE.source, 'gi')
10  
11  const MULTIPLIERS: Record<string, number> = {
12    k: 1_000,
13    m: 1_000_000,
14    b: 1_000_000_000,
15  }
16  
17  function parseBudgetMatch(value: string, suffix: string): number {
18    return parseFloat(value) * MULTIPLIERS[suffix.toLowerCase()]!
19  }
20  
21  export function parseTokenBudget(text: string): number | null {
22    const startMatch = text.match(SHORTHAND_START_RE)
23    if (startMatch) return parseBudgetMatch(startMatch[1]!, startMatch[2]!)
24    const endMatch = text.match(SHORTHAND_END_RE)
25    if (endMatch) return parseBudgetMatch(endMatch[1]!, endMatch[2]!)
26    const verboseMatch = text.match(VERBOSE_RE)
27    if (verboseMatch) return parseBudgetMatch(verboseMatch[1]!, verboseMatch[2]!)
28    return null
29  }
30  
31  export function findTokenBudgetPositions(
32    text: string,
33  ): Array<{ start: number; end: number }> {
34    const positions: Array<{ start: number; end: number }> = []
35    const startMatch = text.match(SHORTHAND_START_RE)
36    if (startMatch) {
37      const offset =
38        startMatch.index! +
39        startMatch[0].length -
40        startMatch[0].trimStart().length
41      positions.push({
42        start: offset,
43        end: startMatch.index! + startMatch[0].length,
44      })
45    }
46    const endMatch = text.match(SHORTHAND_END_RE)
47    if (endMatch) {
48      // Avoid double-counting when input is just "+500k"
49      const endStart = endMatch.index! + 1 // +1: regex includes leading \s
50      const alreadyCovered = positions.some(
51        p => endStart >= p.start && endStart < p.end,
52      )
53      if (!alreadyCovered) {
54        positions.push({
55          start: endStart,
56          end: endMatch.index! + endMatch[0].length,
57        })
58      }
59    }
60    for (const match of text.matchAll(VERBOSE_RE_G)) {
61      positions.push({ start: match.index, end: match.index + match[0].length })
62    }
63    return positions
64  }
65  
66  export function getBudgetContinuationMessage(
67    pct: number,
68    turnTokens: number,
69    budget: number,
70  ): string {
71    const fmt = (n: number): string => new Intl.NumberFormat('en-US').format(n)
72    return `Stopped at ${pct}% of token target (${fmt(turnTokens)} / ${fmt(budget)}). Keep working \u2014 do not summarize.`
73  }