shipdeck/lib/ultimate/anthropic/budget-manager.js

Ship MVPs in 48 hours. Fix bugs in 30 seconds. The command deck for developers who ship.

/**
 * Budget Manager for Soft Caps and Cost Control
 * Implements progressive warnings without hard blocking
 */

const EventEmitter = require('events');

class BudgetManager extends EventEmitter {
  constructor(options = {}) {
    super();
    
    // Soft cap configuration
    this.dailyLimit = options.dailyLimit || 10; // $10 default
    this.warningThresholds = options.warningThresholds || [0.5, 0.8, 0.9]; // 50%, 80%, 90%
    this.hardCapEnabled = options.hardCapEnabled || false;
    this.alertEmail = options.alertEmail || null;
    
    // Current usage tracking
    this.todayUsage = 0;
    this.lastResetDate = new Date().toDateString();
    this.warningsIssued = new Set();
    
    // Cost preview cache
    this.costPreviews = new Map();
    
    // User preferences
    this.userPreferences = {
      qualityMode: options.qualityMode || 'balanced', // 'economy' | 'balanced' | 'premium'
      allowOverage: options.allowOverage !== false,
      notificationPreference: options.notifications || 'all' // 'all' | 'critical' | 'none'
    };
  }

  /**
   * Check if we should issue a budget warning
   */
  checkBudget(estimatedCost) {
    this._resetIfNewDay();
    
    const projectedTotal = this.todayUsage + estimatedCost;
    const percentUsed = (projectedTotal / this.dailyLimit) * 100;
    
    const result = {
      allowed: true,
      currentUsage: this.todayUsage,
      estimatedCost,
      projectedTotal,
      dailyLimit: this.dailyLimit,
      percentUsed,
      warnings: [],
      suggestions: []
    };

    // Check warning thresholds
    for (const threshold of this.warningThresholds) {
      const thresholdPercent = threshold * 100;
      const thresholdKey = `${thresholdPercent}%`;
      
      if (percentUsed >= thresholdPercent && !this.warningsIssued.has(thresholdKey)) {
        this.warningsIssued.add(thresholdKey);
        
        const warning = {
          level: this._getWarningLevel(threshold),
          message: `Budget warning: ${thresholdPercent}% of daily limit reached`,
          threshold: thresholdPercent,
          action: this._getSuggestedAction(threshold)
        };
        
        result.warnings.push(warning);
        this.emit('budget-warning', warning);
        
        // Send notification if configured
        this._sendNotification(warning);
      }
    }

    // Check hard cap if enabled
    if (this.hardCapEnabled && projectedTotal > this.dailyLimit) {
      if (!this.userPreferences.allowOverage) {
        result.allowed = false;
        result.warnings.push({
          level: 'critical',
          message: 'Budget limit exceeded - request blocked',
          threshold: 100,
          action: 'upgrade'
        });
      }
    }

    // Add cost optimization suggestions
    result.suggestions = this._getCostSuggestions(percentUsed, estimatedCost);

    return result;
  }

  /**
   * Update usage after successful API call
   */
  recordUsage(actualCost) {
    this._resetIfNewDay();
    this.todayUsage += actualCost;
    
    // Emit usage event for monitoring
    this.emit('usage-recorded', {
      cost: actualCost,
      total: this.todayUsage,
      remaining: Math.max(0, this.dailyLimit - this.todayUsage)
    });
  }

  /**
   * Get cost preview for a request
   */
  getCostPreview(task, model, estimatedTokens) {
    const cacheKey = `${model}-${estimatedTokens.input}-${estimatedTokens.output}`;
    
    if (this.costPreviews.has(cacheKey)) {
      return this.costPreviews.get(cacheKey);
    }
    
    const preview = {
      task: task.substring(0, 50) + '...',
      model,
      estimatedTokens,
      estimatedCost: this._calculateCost(model, estimatedTokens),
      alternativeModels: this._getAlternativeModels(model, estimatedTokens),
      savingsOpportunity: null
    };
    
    // Calculate potential savings
    if (preview.alternativeModels.length > 0) {
      const cheapest = preview.alternativeModels[0];
      const savings = preview.estimatedCost - cheapest.cost;
      if (savings > 0) {
        preview.savingsOpportunity = {
          model: cheapest.model,
          savings,
          percentSaved: (savings / preview.estimatedCost * 100).toFixed(1)
        };
      }
    }
    
    // Cache for 5 minutes
    this.costPreviews.set(cacheKey, preview);
    setTimeout(() => this.costPreviews.delete(cacheKey), 5 * 60 * 1000);
    
    return preview;
  }

  /**
   * Set user quality preference
   */
  setQualityMode(mode) {
    if (!['economy', 'balanced', 'premium'].includes(mode)) {
      throw new Error('Invalid quality mode');
    }
    
    this.userPreferences.qualityMode = mode;
    
    // Adjust thresholds based on mode
    if (mode === 'economy') {
      this.warningThresholds = [0.3, 0.5, 0.7]; // Warn earlier
    } else if (mode === 'premium') {
      this.warningThresholds = [0.7, 0.85, 0.95]; // Warn later
    } else {
      this.warningThresholds = [0.5, 0.8, 0.9]; // Default
    }
  }

  /**
   * Get budget status summary
   */
  getStatus() {
    this._resetIfNewDay();
    
    return {
      daily: {
        limit: this.dailyLimit,
        used: this.todayUsage,
        remaining: Math.max(0, this.dailyLimit - this.todayUsage),
        percentUsed: (this.todayUsage / this.dailyLimit * 100).toFixed(1)
      },
      settings: {
        hardCapEnabled: this.hardCapEnabled,
        qualityMode: this.userPreferences.qualityMode,
        allowOverage: this.userPreferences.allowOverage
      },
      warnings: Array.from(this.warningsIssued),
      recommendations: this._getRecommendations()
    };
  }

  /**
   * Reset daily usage if it's a new day
   */
  _resetIfNewDay() {
    const today = new Date().toDateString();
    if (today !== this.lastResetDate) {
      this.todayUsage = 0;
      this.lastResetDate = today;
      this.warningsIssued.clear();
      this.emit('daily-reset', { date: today });
    }
  }

  /**
   * Get warning level based on threshold
   */
  _getWarningLevel(threshold) {
    if (threshold >= 0.9) return 'critical';
    if (threshold >= 0.8) return 'high';
    if (threshold >= 0.5) return 'medium';
    return 'low';
  }

  /**
   * Get suggested action based on threshold
   */
  _getSuggestedAction(threshold) {
    if (threshold >= 0.9) {
      return 'Consider delaying non-critical requests or upgrading limit';
    }
    if (threshold >= 0.8) {
      return 'Switch to economy mode for remaining requests';
    }
    return 'Monitor usage - approaching daily limit';
  }

  /**
   * Get cost optimization suggestions
   */
  _getCostSuggestions(percentUsed, estimatedCost) {
    const suggestions = [];
    
    if (percentUsed > 70 && this.userPreferences.qualityMode !== 'economy') {
      suggestions.push({
        type: 'mode_change',
        action: 'Switch to economy mode',
        impact: 'Reduce costs by 40-60%'
      });
    }
    
    if (estimatedCost > 0.50) {
      suggestions.push({
        type: 'task_split',  
        action: 'Consider breaking into smaller subtasks',
        impact: 'Enable caching and reduce token usage'
      });
    }
    
    if (percentUsed > 50) {
      suggestions.push({
        type: 'caching',
        action: 'Check if similar requests were made recently',
        impact: 'Cached responses are free'
      });
    }
    
    return suggestions;
  }

  /**
   * Calculate cost for model and tokens
   */
  _calculateCost(model, tokens) {
    const modelCosts = {
      'claude-3-5-haiku-20241022': { input: 0.001, output: 0.005 },
      'claude-3-5-sonnet-20241022': { input: 0.003, output: 0.015 },
      'claude-opus-4-1-20250805': { input: 0.015, output: 0.075 }
    };
    
    const costs = modelCosts[model] || modelCosts['claude-3-5-sonnet-20241022'];
    
    return (tokens.input / 1000) * costs.input + 
           (tokens.output / 1000) * costs.output;
  }

  /**
   * Get alternative model options
   */
  _getAlternativeModels(currentModel, tokens) {
    const models = [
      { name: 'haiku', id: 'claude-3-5-haiku-20241022' },
      { name: 'sonnet', id: 'claude-3-5-sonnet-20241022' },
      { name: 'opus', id: 'claude-opus-4-1-20250805' }
    ];
    
    return models
      .filter(m => m.id !== currentModel)
      .map(m => ({
        model: m.name,
        cost: this._calculateCost(m.id, tokens)
      }))
      .sort((a, b) => a.cost - b.cost);
  }

  /**
   * Get recommendations based on usage
   */
  _getRecommendations() {
    const recs = [];
    const percentUsed = (this.todayUsage / this.dailyLimit) * 100;
    
    if (percentUsed > 80) {
      recs.push('Consider implementing response caching to reduce API calls');
    }
    
    if (this.userPreferences.qualityMode === 'premium' && percentUsed > 60) {
      recs.push('Switch to balanced mode to extend daily budget');
    }
    
    if (!this.hardCapEnabled && percentUsed > 90) {
      recs.push('Enable hard cap to prevent unexpected overages');
    }
    
    return recs;
  }

  /**
   * Send notification for warnings
   */
  _sendNotification(warning) {
    // In production, this would send email/slack/webhook
    if (this.userPreferences.notificationPreference === 'none') return;
    
    if (warning.level === 'critical' || 
        this.userPreferences.notificationPreference === 'all') {
      console.log(`📧 Budget Alert: ${warning.message}`);
      // TODO: Implement actual notification service
    }
  }

  /**
   * Override daily limit
   */
  setDailyLimit(newLimit) {
    this.dailyLimit = newLimit;
    this.emit('limit-changed', { newLimit });
  }

  /**
   * Enable/disable hard cap
   */
  setHardCap(enabled) {
    this.hardCapEnabled = enabled;
    this.emit('hardcap-changed', { enabled });
  }
}

module.exports = { BudgetManager };