1  import { getBudgetContinuationMessage } from '../utils/tokenBudget.js'
 2  
 3  const COMPLETION_THRESHOLD = 0.9
 4  const DIMINISHING_THRESHOLD = 500
 5  
 6  export type BudgetTracker = {
 7    continuationCount: number
 8    lastDeltaTokens: number
 9    lastGlobalTurnTokens: number
10    startedAt: number
11  }
12  
13  export function createBudgetTracker(): BudgetTracker {
14    return {
15      continuationCount: 0,
16      lastDeltaTokens: 0,
17      lastGlobalTurnTokens: 0,
18      startedAt: Date.now(),
19    }
20  }
21  
22  type ContinueDecision = {
23    action: 'continue'
24    nudgeMessage: string
25    continuationCount: number
26    pct: number
27    turnTokens: number
28    budget: number
29  }
30  
31  type StopDecision = {
32    action: 'stop'
33    completionEvent: {
34      continuationCount: number
35      pct: number
36      turnTokens: number
37      budget: number
38      diminishingReturns: boolean
39      durationMs: number
40    } | null
41  }
42  
43  export type TokenBudgetDecision = ContinueDecision | StopDecision
44  
45  export function checkTokenBudget(
46    tracker: BudgetTracker,
47    agentId: string | undefined,
48    budget: number | null,
49    globalTurnTokens: number,
50  ): TokenBudgetDecision {
51    if (agentId || budget === null || budget <= 0) {
52      return { action: 'stop', completionEvent: null }
53    }
54  
55    const turnTokens = globalTurnTokens
56    const pct = Math.round((turnTokens / budget) * 100)
57    const deltaSinceLastCheck = globalTurnTokens - tracker.lastGlobalTurnTokens
58  
59    const isDiminishing =
60      tracker.continuationCount >= 3 &&
61      deltaSinceLastCheck < DIMINISHING_THRESHOLD &&
62      tracker.lastDeltaTokens < DIMINISHING_THRESHOLD
63  
64    if (!isDiminishing && turnTokens < budget * COMPLETION_THRESHOLD) {
65      tracker.continuationCount++
66      tracker.lastDeltaTokens = deltaSinceLastCheck
67      tracker.lastGlobalTurnTokens = globalTurnTokens
68      return {
69        action: 'continue',
70        nudgeMessage: getBudgetContinuationMessage(pct, turnTokens, budget),
71        continuationCount: tracker.continuationCount,
72        pct,
73        turnTokens,
74        budget,
75      }
76    }
77  
78    if (isDiminishing || tracker.continuationCount > 0) {
79      return {
80        action: 'stop',
81        completionEvent: {
82          continuationCount: tracker.continuationCount,
83          pct,
84          turnTokens,
85          budget,
86          diminishingReturns: isDiminishing,
87          durationMs: Date.now() - tracker.startedAt,
88        },
89      }
90    }
91  
92    return { action: 'stop', completionEvent: null }
93  }