init claude-code

query/tokenBudget.ts

@@ -0,0 +1,93 @@
+import { getBudgetContinuationMessage } from '../utils/tokenBudget.js'
+
+const COMPLETION_THRESHOLD = 0.9
+const DIMINISHING_THRESHOLD = 500
+
+export type BudgetTracker = {
+  continuationCount: number
+  lastDeltaTokens: number
+  lastGlobalTurnTokens: number
+  startedAt: number
+}
+
+export function createBudgetTracker(): BudgetTracker {
+  return {
+    continuationCount: 0,
+    lastDeltaTokens: 0,
+    lastGlobalTurnTokens: 0,
+    startedAt: Date.now(),
+  }
+}
+
+type ContinueDecision = {
+  action: 'continue'
+  nudgeMessage: string
+  continuationCount: number
+  pct: number
+  turnTokens: number
+  budget: number
+}
+
+type StopDecision = {
+  action: 'stop'
+  completionEvent: {
+    continuationCount: number
+    pct: number
+    turnTokens: number
+    budget: number
+    diminishingReturns: boolean
+    durationMs: number
+  } | null
+}
+
+export type TokenBudgetDecision = ContinueDecision | StopDecision
+
+export function checkTokenBudget(
+  tracker: BudgetTracker,
+  agentId: string | undefined,
+  budget: number | null,
+  globalTurnTokens: number,
+): TokenBudgetDecision {
+  if (agentId || budget === null || budget <= 0) {
+    return { action: 'stop', completionEvent: null }
+  }
+
+  const turnTokens = globalTurnTokens
+  const pct = Math.round((turnTokens / budget) * 100)
+  const deltaSinceLastCheck = globalTurnTokens - tracker.lastGlobalTurnTokens
+
+  const isDiminishing =
+    tracker.continuationCount >= 3 &&
+    deltaSinceLastCheck < DIMINISHING_THRESHOLD &&
+    tracker.lastDeltaTokens < DIMINISHING_THRESHOLD
+
+  if (!isDiminishing && turnTokens < budget * COMPLETION_THRESHOLD) {
+    tracker.continuationCount++
+    tracker.lastDeltaTokens = deltaSinceLastCheck
+    tracker.lastGlobalTurnTokens = globalTurnTokens
+    return {
+      action: 'continue',
+      nudgeMessage: getBudgetContinuationMessage(pct, turnTokens, budget),
+      continuationCount: tracker.continuationCount,
+      pct,
+      turnTokens,
+      budget,
+    }
+  }
+
+  if (isDiminishing || tracker.continuationCount > 0) {
+    return {
+      action: 'stop',
+      completionEvent: {
+        continuationCount: tracker.continuationCount,
+        pct,
+        turnTokens,
+        budget,
+        diminishingReturns: isDiminishing,
+        durationMs: Date.now() - tracker.startedAt,
+      },
+    }
+  }
+
+  return { action: 'stop', completionEvent: null }
+}