ai-control-center 1.15.2

package/lib/utils/context-compressor.js

@@ -0,0 +1,172 @@
+/**
+ * Context Compressor — reduces prompt sizes for stateless AI calls (Gemini).
+ *
+ * Gemini is stateless — every call re-sends the full context. This module
+ * creates compact context digests per stage to avoid sending the full
+ * codebase tree, README, and git history every time.
+ *
+ * Cache is in-memory (Map) and resets on pipeline start.
+ */
+import { execSync } from 'child_process';
+import { existsSync, readFileSync, readdirSync, statSync } from 'fs';
+import { resolve, relative, extname } from 'path';
+import { getRootDir } from './pipeline.js';
+
+// In-memory cache — keyed by featureId, reset on new pipeline
+const _cache = new Map();
+
+/**
+ * Reset the context cache (call when starting a new pipeline run).
+ */
+export function resetContextCache() {
+  _cache.clear();
+}
+
+/**
+ * Get a cached value or compute and store it.
+ */
+function cached(featureId, key, computeFn) {
+  const cacheKey = `${featureId}:${key}`;
+  if (_cache.has(cacheKey)) return _cache.get(cacheKey);
+  const value = computeFn();
+  _cache.set(cacheKey, value);
+  return value;
+}
+
+/**
+ * Generate a compact file tree listing (names + sizes only, no content).
+ */
+function compactFileTree(root) {
+  const SKIP = new Set(['node_modules', '.git', '.ai-workflow', 'dist', 'build', 'coverage', '.next']);
+  const SOURCE_EXTS = new Set(['.js', '.ts', '.jsx', '.tsx', '.py', '.go', '.java', '.cls', '.trigger', '.vue', '.svelte']);
+  const lines = [];
+
+  function walk(dir, depth = 0) {
+    if (depth > 4) return; // max 4 levels deep
+    try {
+      const entries = readdirSync(dir).sort();
+      for (const entry of entries) {
+        if (entry.startsWith('.') && depth === 0 && !entry.startsWith('.claude')) continue;
+        if (SKIP.has(entry)) continue;
+        const fullPath = resolve(dir, entry);
+        const stat = statSync(fullPath);
+        if (stat.isDirectory()) {
+          lines.push(`${'  '.repeat(depth)}${entry}/`);
+          walk(fullPath, depth + 1);
+        } else if (SOURCE_EXTS.has(extname(entry))) {
+          const sizeKb = (stat.size / 1024).toFixed(1);
+          lines.push(`${'  '.repeat(depth)}${entry} (${sizeKb}KB)`);
+        }
+      }
+    } catch { /* permission or read errors */ }
+  }
+
+  walk(root);
+  return lines.join('\n');
+}
+
+/**
+ * Get git diff for changed files only (for review stage).
+ */
+function getGitDiff(root) {
+  try {
+    const diff = execSync('git diff HEAD~1 --stat --no-color 2>/dev/null || git diff --cached --stat --no-color 2>/dev/null || echo "No diff available"', {
+      cwd: root, encoding: 'utf8', maxBuffer: 50 * 1024,
+    });
+    return diff.trim();
+  } catch {
+    return 'No git diff available';
+  }
+}
+
+/**
+ * Get the full diff content for changed files.
+ */
+function getGitDiffContent(root) {
+  try {
+    const diff = execSync('git diff HEAD~1 --no-color 2>/dev/null || git diff --cached --no-color 2>/dev/null || echo "No diff available"', {
+      cwd: root, encoding: 'utf8', maxBuffer: 200 * 1024,
+    });
+    return diff.trim().slice(0, 50000); // cap at 50KB
+  } catch {
+    return 'No git diff available';
+  }
+}
+
+/**
+ * Get list of relevant source files for a feature (no content, just paths).
+ */
+function getRelevantFileList(root) {
+  try {
+    const files = execSync('git ls-files --others --cached --exclude-standard 2>/dev/null | head -100', {
+      cwd: root, encoding: 'utf8',
+    });
+    return files.trim();
+  } catch {
+    return compactFileTree(root);
+  }
+}
+
+/**
+ * Compress context for a specific pipeline stage.
+ *
+ * @param {string} stage       - 'pm', 'review', 'arch', 'impl'
+ * @param {string} featureId   - Current feature ID
+ * @param {string} fullContext  - The original full context string
+ * @returns {string} Compressed context
+ */
+export function compressForStage(stage, featureId, fullContext) {
+  const root = getRootDir();
+  const originalLength = fullContext.length;
+
+  let compressed;
+
+  switch (stage) {
+    case 'pm':
+    case 'spec': {
+      // PM spec: only include feature description + relevant file list (not full tree)
+      const fileList = cached(featureId, 'fileList', () => getRelevantFileList(root));
+      compressed = `## Project Structure (file list only)\n${fileList}\n\n` +
+        fullContext.replace(/## (Full |Complete )?(Codebase |Source )?(Tree|Directory|Structure)[\s\S]*?(?=## |$)/gi, '')
+          .replace(/## Git (History|Log)[\s\S]*?(?=## |$)/gi, '')
+          .replace(/## README[\s\S]*?(?=## |$)/gi, '');
+      break;
+    }
+
+    case 'review': {
+      // Review: only include changed files diff (not full codebase tree)
+      const diffStat = cached(featureId, 'diffStat', () => getGitDiff(root));
+      const diffContent = cached(featureId, 'diffContent', () => getGitDiffContent(root));
+      compressed = `## Changed Files\n${diffStat}\n\n## Diff\n${diffContent}\n\n` +
+        fullContext.replace(/## (Full |Complete )?(Codebase |Source )?(Tree|Directory|Structure)[\s\S]*?(?=## |$)/gi, '')
+          .replace(/## Git (History|Log)[\s\S]*?(?=## |$)/gi, '');
+      break;
+    }
+
+    case 'arch':
+    case 'architect': {
+      // Architecture: include compact tree but skip git history
+      const tree = cached(featureId, 'compactTree', () => compactFileTree(root));
+      compressed = `## Project Structure (compact)\n${tree}\n\n` +
+        fullContext.replace(/## Git (History|Log)[\s\S]*?(?=## |$)/gi, '');
+      break;
+    }
+
+    default:
+      compressed = fullContext;
+  }
+
+  // Remove excessive whitespace
+  compressed = compressed.replace(/\n{3,}/g, '\n\n').trim();
+
+  const compressedTokens = Math.ceil(compressed.length / 4);
+  const originalTokens = Math.ceil(originalLength / 4);
+  const reduction = originalLength > 0 ? Math.round((1 - compressed.length / originalLength) * 100) : 0;
+
+  if (reduction > 5) {
+    // Only log if meaningful reduction achieved
+    console.log(`[SYSTEM] Context compressed: ${originalTokens.toLocaleString()} → ${compressedTokens.toLocaleString()} tokens (${reduction}% reduction)`);
+  }
+
+  return compressed;
+}

package/lib/utils/correlation.js

@@ -0,0 +1,63 @@
+/**
+ * Correlation ID system for distributed tracing.
+ */
+
+import { randomUUID } from 'crypto';
+
+let currentCorrelationId = null;
+
+export function generateCorrelationId() {
+  return randomUUID();
+}
+
+export function withCorrelation(id, fn) {
+  const previous = currentCorrelationId;
+  currentCorrelationId = id;
+  try {
+    return fn();
+  } finally {
+    currentCorrelationId = previous;
+  }
+}
+
+export function getCorrelationId() {
+  return currentCorrelationId;
+}
+
+export function setCorrelationId(id) {
+  currentCorrelationId = id;
+}
+
+export function clearCorrelationId() {
+  currentCorrelationId = null;
+}
+
+export function correlatedLog(level, message, meta) {
+  const prefix = currentCorrelationId ? `[PIPELINE:${currentCorrelationId}]` : '[PIPELINE]';
+  const logFn = console[level] || console.log;
+  if (meta !== undefined) {
+    logFn(`${prefix} ${message}`, meta);
+  } else {
+    logFn(`${prefix} ${message}`);
+  }
+}
+
+export function buildTraceTimeline(correlationId, events) {
+  if (!events || events.length === 0) {
+    return { correlationId, startedAt: null, completedAt: null, durationMs: 0, stages: [] };
+  }
+
+  const sorted = [...events].sort((a, b) => new Date(a.timestamp) - new Date(b.timestamp));
+  const startedAt = sorted[0].timestamp;
+  const completedAt = sorted[sorted.length - 1].timestamp;
+  const durationMs = new Date(completedAt) - new Date(startedAt);
+
+  const stages = sorted.map((event) => ({
+    name: event.name || event.stage,
+    timestamp: event.timestamp,
+    duration: event.duration || null,
+    status: event.status || 'completed',
+  }));
+
+  return { correlationId, startedAt, completedAt, durationMs, stages };
+}

package/lib/utils/cost-tracker.js

@@ -0,0 +1,423 @@
+/**
+ * Cost Tracker — tracks AI usage and estimated costs per feature per provider.
+ *
+ * Stores usage data in .ai-workflow/costs.jsonl (append-only JSONL format).
+ * Each line is a JSON object with: timestamp, provider, model, stage, featureId,
+ * inputTokens, outputTokens, estimatedCost.
+ *
+ * Token estimation: since CLI tools don't always return token counts,
+ * we estimate from prompt length (÷4) and response length (÷4).
+ *
+ * Pricing is based on published per-1M-token rates (June 2025).
+ */
+import { appendFileSync, existsSync, mkdirSync, readFileSync } from 'fs';
+import { resolve } from 'path';
+import { bus } from '../shared/event-bus.js';
+import { ErrorCodes, TokenBudgetError } from './errors.js';
+import { getWorkflowDir } from './pipeline.js';
+
+// ─── Pricing Table (per 1M tokens, USD) ──────────────────────────────────────
+
+const PRICING = {
+  // Gemini
+  'gemini-2.5-pro':       { input: 1.25,  output: 10.00 },
+  'gemini-2.5-flash':     { input: 0.15,  output: 0.60  },
+  'gemini-2.0-flash':     { input: 0.10,  output: 0.40  },
+  // Claude
+  'claude-sonnet-4-6':    { input: 3.00,  output: 15.00 },
+  'claude-haiku-4-5-20251001': { input: 0.80, output: 4.00 },
+  // Copilot (GitHub-hosted — technically free/included in subscription,
+  // but we track "equivalent" cost for comparison)
+  'claude-sonnet-4.6':    { input: 3.00,  output: 15.00 },
+  'claude-haiku-4.5':     { input: 0.80,  output: 4.00  },
+  'claude-opus-4.6':      { input: 15.00, output: 75.00 },
+  // Ollama (local — zero cost but track usage)
+  'ollama':               { input: 0.00,  output: 0.00  },
+  // OpenClaw (uses Copilot under the hood)
+  'openclaw':             { input: 0.00,  output: 0.00  },
+};
+
+// Fallback pricing for unknown models
+const DEFAULT_PRICING = { input: 1.00, output: 5.00 };
+
+// ─── Core Functions ──────────────────────────────────────────────────────────
+
+/**
+ * Get the cost file path.
+ */
+function getCostFile() {
+  const dir = getWorkflowDir();
+  if (!existsSync(dir)) mkdirSync(dir, { recursive: true });
+  return resolve(dir, 'costs.jsonl');
+}
+
+/**
+ * Estimate token count from text length.
+ * Rough approximation: 1 token ≈ 4 characters for English text.
+ */
+function estimateTokens(text) {
+  if (!text) return 0;
+  return Math.ceil(String(text).length / 4);
+}
+
+/**
+ * Get pricing for a model.
+ */
+function getModelPricing(model) {
+  if (!model) return DEFAULT_PRICING;
+  // Try exact match first
+  if (PRICING[model]) return PRICING[model];
+  // Try partial match (e.g. "gemini-2.5-pro-latest" → "gemini-2.5-pro")
+  const key = Object.keys(PRICING).find(k => model.includes(k) || k.includes(model));
+  return key ? PRICING[key] : DEFAULT_PRICING;
+}
+
+/**
+ * Calculate estimated cost for a usage entry.
+ */
+function calculateCost(model, inputTokens, outputTokens) {
+  const pricing = getModelPricing(model);
+  const inputCost  = (inputTokens / 1_000_000) * pricing.input;
+  const outputCost = (outputTokens / 1_000_000) * pricing.output;
+  return Math.round((inputCost + outputCost) * 1_000_000) / 1_000_000; // 6 decimal places
+}
+
+/**
+ * Track an AI usage event.
+ *
+ * @param {object} params
+ * @param {string} params.provider   - 'gemini' | 'claude' | 'copilot' | 'ollama' | 'openclaw'
+ * @param {string} params.model      - Specific model name
+ * @param {string} params.stage      - Pipeline stage: 'pm' | 'architect' | 'implement' | 'review' | 'ask' | 'chat'
+ * @param {string} params.featureId  - Feature identifier (optional)
+ * @param {number} params.inputTokens  - Input token count (or estimated)
+ * @param {number} params.outputTokens - Output token count (or estimated)
+ * @param {string} params.promptText   - Raw prompt text (for estimation if tokens not provided)
+ * @param {string} params.responseText - Raw response text (for estimation if tokens not provided)
+ * @param {number} params.durationMs   - Call duration in milliseconds (optional)
+ */
+export function trackUsage(params) {
+  try {
+    const {
+      provider = 'unknown',
+      model = 'unknown',
+      stage = 'unknown',
+      featureId = null,
+      promptText = '',
+      responseText = '',
+      durationMs = null,
+    } = params;
+
+    const inputTokens  = params.inputTokens  || estimateTokens(promptText);
+    const outputTokens = params.outputTokens || estimateTokens(responseText);
+    const estimatedCost = calculateCost(model, inputTokens, outputTokens);
+
+    const entry = {
+      ts: new Date().toISOString(),
+      provider,
+      model,
+      stage,
+      featureId,
+      inputTokens,
+      outputTokens,
+      estimatedCost,
+      durationMs,
+    };
+
+    const costFile = getCostFile();
+    appendFileSync(costFile, JSON.stringify(entry) + '\n', 'utf8');
+
+    return entry;
+  } catch {
+    // Non-fatal — never crash the pipeline for cost tracking
+    return null;
+  }
+}
+
+/**
+ * Read all cost entries from the JSONL file.
+ *
+ * @param {object} filters
+ * @param {string} filters.featureId - Filter by feature
+ * @param {string} filters.provider  - Filter by provider
+ * @param {string} filters.stage     - Filter by stage
+ * @param {number} filters.last      - Only last N entries
+ * @returns {Array} Parsed cost entries
+ */
+export function getCostEntries(filters = {}) {
+  try {
+    const costFile = getCostFile();
+    if (!existsSync(costFile)) return [];
+
+    let entries = readFileSync(costFile, 'utf8')
+      .split('\n')
+      .filter(Boolean)
+      .map(line => {
+        try { return JSON.parse(line); }
+        catch { return null; }
+      })
+      .filter(Boolean);
+
+    if (filters.featureId) entries = entries.filter(e => e.featureId === filters.featureId);
+    if (filters.provider)  entries = entries.filter(e => e.provider === filters.provider);
+    if (filters.stage)     entries = entries.filter(e => e.stage === filters.stage);
+    if (filters.last)      entries = entries.slice(-filters.last);
+
+    return entries;
+  } catch {
+    return [];
+  }
+}
+
+/**
+ * Get cost summary — total and per-provider breakdown.
+ *
+ * @param {string} featureId - Optional: get costs for a specific feature
+ * @returns {object} Summary with totals and per-provider/per-model breakdown
+ */
+export function getCostSummary(featureId = null) {
+  const entries = getCostEntries(featureId ? { featureId } : {});
+
+  const summary = {
+    totalCost: 0,
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    totalCalls: entries.length,
+    byProvider: {},
+    byModel: {},
+    byStage: {},
+    byFeature: {},
+  };
+
+  for (const entry of entries) {
+    summary.totalCost += entry.estimatedCost || 0;
+    summary.totalInputTokens += entry.inputTokens || 0;
+    summary.totalOutputTokens += entry.outputTokens || 0;
+
+    // By provider
+    if (!summary.byProvider[entry.provider]) {
+      summary.byProvider[entry.provider] = { calls: 0, cost: 0, inputTokens: 0, outputTokens: 0 };
+    }
+    summary.byProvider[entry.provider].calls++;
+    summary.byProvider[entry.provider].cost += entry.estimatedCost || 0;
+    summary.byProvider[entry.provider].inputTokens += entry.inputTokens || 0;
+    summary.byProvider[entry.provider].outputTokens += entry.outputTokens || 0;
+
+    // By model
+    if (!summary.byModel[entry.model]) {
+      summary.byModel[entry.model] = { calls: 0, cost: 0, inputTokens: 0, outputTokens: 0 };
+    }
+    summary.byModel[entry.model].calls++;
+    summary.byModel[entry.model].cost += entry.estimatedCost || 0;
+    summary.byModel[entry.model].inputTokens += entry.inputTokens || 0;
+    summary.byModel[entry.model].outputTokens += entry.outputTokens || 0;
+
+    // By stage
+    if (!summary.byStage[entry.stage]) {
+      summary.byStage[entry.stage] = { calls: 0, cost: 0 };
+    }
+    summary.byStage[entry.stage].calls++;
+    summary.byStage[entry.stage].cost += entry.estimatedCost || 0;
+
+    // By feature (only if featureId present)
+    if (entry.featureId) {
+      if (!summary.byFeature[entry.featureId]) {
+        summary.byFeature[entry.featureId] = { calls: 0, cost: 0 };
+      }
+      summary.byFeature[entry.featureId].calls++;
+      summary.byFeature[entry.featureId].cost += entry.estimatedCost || 0;
+    }
+  }
+
+  // Round totals
+  summary.totalCost = Math.round(summary.totalCost * 1_000_000) / 1_000_000;
+  for (const p of Object.values(summary.byProvider)) p.cost = Math.round(p.cost * 1_000_000) / 1_000_000;
+  for (const m of Object.values(summary.byModel))    m.cost = Math.round(m.cost * 1_000_000) / 1_000_000;
+  for (const s of Object.values(summary.byStage))    s.cost = Math.round(s.cost * 1_000_000) / 1_000_000;
+  for (const f of Object.values(summary.byFeature))   f.cost = Math.round(f.cost * 1_000_000) / 1_000_000;
+
+  return summary;
+}
+
+/**
+ * Format cost summary for display (terminal or Telegram).
+ *
+ * @param {object} summary - From getCostSummary()
+ * @param {string} format  - 'text' | 'html'
+ * @returns {string}
+ */
+export function formatCostSummary(summary, format = 'text') {
+  if (!summary || summary.totalCalls === 0) {
+    return format === 'html'
+      ? '<i>No AI usage tracked yet.</i>'
+      : 'No AI usage tracked yet.';
+  }
+
+  if (format === 'html') {
+    let html = `💰 <b>AI Cost Summary</b>\n\n`;
+    html += `Total: <b>$${summary.totalCost.toFixed(4)}</b> across ${summary.totalCalls} calls\n`;
+    html += `Tokens: ${summary.totalInputTokens.toLocaleString()} in / ${summary.totalOutputTokens.toLocaleString()} out\n\n`;
+
+    html += `<b>By Provider:</b>\n`;
+    for (const [provider, data] of Object.entries(summary.byProvider)) {
+      html += `  ${provider}: $${data.cost.toFixed(4)} (${data.calls} calls)\n`;
+    }
+
+    if (Object.keys(summary.byModel).length > 1) {
+      html += `\n<b>By Model:</b>\n`;
+      for (const [model, data] of Object.entries(summary.byModel)) {
+        html += `  ${model}: $${data.cost.toFixed(4)} (${data.calls} calls)\n`;
+      }
+    }
+
+    if (Object.keys(summary.byStage).length > 0) {
+      html += `\n<b>By Stage:</b>\n`;
+      for (const [stage, data] of Object.entries(summary.byStage)) {
+        html += `  ${stage}: $${data.cost.toFixed(4)} (${data.calls} calls)\n`;
+      }
+    }
+
+    return html;
+  }
+
+  // Plain text format
+  let text = `\n💰 AI Cost Summary\n${'─'.repeat(40)}\n`;
+  text += `Total: $${summary.totalCost.toFixed(4)} across ${summary.totalCalls} calls\n`;
+  text += `Tokens: ${summary.totalInputTokens.toLocaleString()} in / ${summary.totalOutputTokens.toLocaleString()} out\n\n`;
+
+  text += `By Provider:\n`;
+  for (const [provider, data] of Object.entries(summary.byProvider)) {
+    text += `  ${provider.padEnd(10)} $${data.cost.toFixed(4).padStart(8)} (${data.calls} calls)\n`;
+  }
+
+  text += `\nBy Model:\n`;
+  for (const [model, data] of Object.entries(summary.byModel)) {
+    text += `  ${model.padEnd(30)} $${data.cost.toFixed(4).padStart(8)} (${data.calls} calls)\n`;
+  }
+
+  text += `\nBy Stage:\n`;
+  for (const [stage, data] of Object.entries(summary.byStage)) {
+    text += `  ${stage.padEnd(15)} $${data.cost.toFixed(4).padStart(8)} (${data.calls} calls)\n`;
+  }
+
+  return text;
+}
+
+// ─── Token Budget Per Pipeline Run ────────────────────────────────────────────
+
+const DEFAULT_TOKEN_BUDGET = 100_000;
+const DEFAULT_COST_BUDGET = 2.00;
+
+/**
+ * Token budget tracker for a single pipeline run.
+ * Tracks cumulative token usage and emits warnings/events at thresholds.
+ */
+export class TokenBudget {
+  constructor(featureId, options = {}) {
+    this.featureId = featureId;
+    this.maxTokens = options.maxTokens || DEFAULT_TOKEN_BUDGET;
+    this.maxCost = options.maxCost || DEFAULT_COST_BUDGET;
+    this.usedTokens = 0;
+    this.estimatedCost = 0;
+    this._warned80 = false;
+    this._exceeded = false;
+    this.savedByCheckpoints = 0;
+  }
+
+  /**
+   * Track token usage from a completed AI call.
+   */
+  track(inputTokens, outputTokens, cost = 0) {
+    const total = (inputTokens || 0) + (outputTokens || 0);
+    this.usedTokens += total;
+    this.estimatedCost += cost;
+
+    if (!this._warned80 && this.percentUsed() >= 80) {
+      this._warned80 = true;
+      const msg = `[COST] Token budget 80% used (${this.usedTokens.toLocaleString()}/${this.maxTokens.toLocaleString()}) for ${this.featureId}`;
+      console.log(msg);
+      try {
+        bus.emitEvent('token_budget_warning', {
+          featureId: this.featureId,
+          used: this.usedTokens,
+          limit: this.maxTokens,
+          percent: this.percentUsed(),
+          message: `⚠️ <b>Token budget 80% used</b>\n\n${this.usedTokens.toLocaleString()} / ${this.maxTokens.toLocaleString()} tokens for <code>${this.featureId}</code>\n\nEstimated cost: $${this.estimatedCost.toFixed(4)}`,
+        });
+      } catch { /* non-fatal */ }
+    }
+
+    if (!this._exceeded && this.isExceeded()) {
+      this._exceeded = true;
+      try {
+        bus.emitEvent('token_budget_exceeded', {
+          featureId: this.featureId,
+          used: this.usedTokens,
+          limit: this.maxTokens,
+          message: `🛑 <b>Token budget exceeded</b>\n\n${this.usedTokens.toLocaleString()} / ${this.maxTokens.toLocaleString()} tokens for <code>${this.featureId}</code>\n\nPipeline paused. Use /approve to continue or /reject to stop.`,
+        });
+      } catch { /* non-fatal */ }
+    }
+  }
+
+  /**
+   * Record tokens saved by loading from checkpoint instead of calling AI.
+   */
+  trackCheckpointSaving(estimatedTokens) {
+    this.savedByCheckpoints += estimatedTokens;
+  }
+
+  remaining() {
+    return Math.max(0, this.maxTokens - this.usedTokens);
+  }
+
+  percentUsed() {
+    return Math.round((this.usedTokens / this.maxTokens) * 100);
+  }
+
+  isExceeded() {
+    return this.usedTokens >= this.maxTokens || this.estimatedCost >= this.maxCost;
+  }
+
+  toJSON() {
+    return {
+      used: this.usedTokens,
+      limit: this.maxTokens,
+      percent: this.percentUsed(),
+      estimatedCost: Math.round(this.estimatedCost * 1_000_000) / 1_000_000,
+      maxCost: this.maxCost,
+      remaining: this.remaining(),
+      exceeded: this.isExceeded(),
+      savedByCheckpoints: this.savedByCheckpoints,
+    };
+  }
+
+  /**
+   * Format budget as a visual gauge for display.
+   */
+  formatGauge() {
+    const pct = this.percentUsed();
+    const filled = Math.round(pct / 10);
+    const empty = 10 - filled;
+    const bar = '='.repeat(filled) + (filled < 10 ? '>' : '') + ' '.repeat(Math.max(0, empty - 1));
+    return `[${bar}] ${pct}% ($${this.estimatedCost.toFixed(2)} / $${this.maxCost.toFixed(2)})`;
+  }
+}
+
+// Active budget instance per feature
+const _activeBudgets = new Map();
+
+export function getOrCreateBudget(featureId, config = {}) {
+  if (!_activeBudgets.has(featureId)) {
+    _activeBudgets.set(featureId, new TokenBudget(featureId, config));
+  }
+  return _activeBudgets.get(featureId);
+}
+
+export function resetBudget(featureId) {
+  _activeBudgets.delete(featureId);
+}
+
+export function getActiveBudget(featureId) {
+  return _activeBudgets.get(featureId) || null;
+}