stack/apps/orchestrator/src/budget/budget.service.ts

/**
 * Usage Budget Management Service
 *
 * Tracks token usage per agent and enforces budget limits.
 * Provides real-time usage summaries and budget status checks.
 */
import { Injectable, Logger } from "@nestjs/common";
import { ConfigService } from "@nestjs/config";
import type {
  UsageBudget,
  UsageRecord,
  UsageSummary,
  AgentUsageSummary,
  BudgetStatus,
} from "./budget.types";
import { DEFAULT_BUDGET } from "./budget.types";

@Injectable()
export class BudgetService {
  private readonly logger = new Logger(BudgetService.name);
  private readonly budget: UsageBudget;
  private readonly records: UsageRecord[] = [];
  private readonly activeAgentCount = { value: 0 };

  constructor(private readonly configService: ConfigService) {
    this.budget = {
      dailyTokenLimit:
        this.configService.get<number>("orchestrator.budget.dailyTokenLimit") ??
        DEFAULT_BUDGET.dailyTokenLimit,
      perAgentTokenLimit:
        this.configService.get<number>("orchestrator.budget.perAgentTokenLimit") ??
        DEFAULT_BUDGET.perAgentTokenLimit,
      maxConcurrentAgents:
        this.configService.get<number>("orchestrator.budget.maxConcurrentAgents") ??
        DEFAULT_BUDGET.maxConcurrentAgents,
      maxTaskDurationMinutes:
        this.configService.get<number>("orchestrator.budget.maxTaskDurationMinutes") ??
        DEFAULT_BUDGET.maxTaskDurationMinutes,
      enforceHardLimits:
        this.configService.get<boolean>("orchestrator.budget.enforceHardLimits") ??
        DEFAULT_BUDGET.enforceHardLimits,
    };

    this.logger.log(
      `BudgetService initialized: daily=${String(this.budget.dailyTokenLimit)} tokens, ` +
        `perAgent=${String(this.budget.perAgentTokenLimit)} tokens, ` +
        `maxConcurrent=${String(this.budget.maxConcurrentAgents)}`
    );
  }

  /**
   * Record token usage for an agent
   */
  recordUsage(agentId: string, taskId: string, inputTokens: number, outputTokens: number): void {
    const record: UsageRecord = {
      agentId,
      taskId,
      inputTokens,
      outputTokens,
      timestamp: new Date().toISOString(),
    };

    this.records.push(record);

    this.logger.debug(
      `Usage recorded: agent=${agentId} input=${String(inputTokens)} output=${String(outputTokens)}`
    );
  }

  /**
   * Check if an agent can be spawned (concurrency and budget check)
   */
  canSpawnAgent(): { allowed: boolean; reason?: string } {
    if (this.activeAgentCount.value >= this.budget.maxConcurrentAgents) {
      return {
        allowed: false,
        reason: `Maximum concurrent agents reached (${String(this.budget.maxConcurrentAgents)})`,
      };
    }

    const dailyUsed = this.getDailyTokensUsed();
    if (this.budget.enforceHardLimits && dailyUsed >= this.budget.dailyTokenLimit) {
      return {
        allowed: false,
        reason: `Daily token budget exceeded (${String(dailyUsed)}/${String(this.budget.dailyTokenLimit)})`,
      };
    }

    return { allowed: true };
  }

  /**
   * Check if an agent has exceeded its per-task budget
   */
  isAgentOverBudget(agentId: string): { overBudget: boolean; totalTokens: number } {
    const agentRecords = this.records.filter((r) => r.agentId === agentId);
    const totalTokens = agentRecords.reduce((sum, r) => sum + r.inputTokens + r.outputTokens, 0);

    return {
      overBudget: totalTokens >= this.budget.perAgentTokenLimit,
      totalTokens,
    };
  }

  /**
   * Notify that an agent has started (increment active count)
   */
  agentStarted(): void {
    this.activeAgentCount.value++;
  }

  /**
   * Notify that an agent has stopped (decrement active count)
   */
  agentStopped(): void {
    this.activeAgentCount.value = Math.max(0, this.activeAgentCount.value - 1);
  }

  /**
   * Get comprehensive usage summary
   */
  getUsageSummary(): UsageSummary {
    const dailyTokensUsed = this.getDailyTokensUsed();
    const dailyUsagePercent =
      this.budget.dailyTokenLimit > 0 ? (dailyTokensUsed / this.budget.dailyTokenLimit) * 100 : 0;

    return {
      dailyTokensUsed,
      dailyTokenLimit: this.budget.dailyTokenLimit,
      dailyUsagePercent: Math.round(dailyUsagePercent * 100) / 100,
      agentUsage: this.getAgentUsageSummaries(),
      activeAgents: this.activeAgentCount.value,
      maxConcurrentAgents: this.budget.maxConcurrentAgents,
      budgetStatus: this.getBudgetStatus(dailyUsagePercent),
    };
  }

  /**
   * Get the configured budget
   */
  getBudget(): UsageBudget {
    return { ...this.budget };
  }

  /**
   * Get total tokens used today
   */
  private getDailyTokensUsed(): number {
    const todayStart = new Date();
    todayStart.setHours(0, 0, 0, 0);
    const todayIso = todayStart.toISOString();

    return this.records
      .filter((r) => r.timestamp >= todayIso)
      .reduce((sum, r) => sum + r.inputTokens + r.outputTokens, 0);
  }

  /**
   * Get per-agent usage summaries
   */
  private getAgentUsageSummaries(): AgentUsageSummary[] {
    const agentMap = new Map<string, { taskId: string; input: number; output: number }>();

    for (const record of this.records) {
      const existing = agentMap.get(record.agentId);
      if (existing) {
        existing.input += record.inputTokens;
        existing.output += record.outputTokens;
      } else {
        agentMap.set(record.agentId, {
          taskId: record.taskId,
          input: record.inputTokens,
          output: record.outputTokens,
        });
      }
    }

    return Array.from(agentMap.entries()).map(([agentId, data]) => {
      const totalTokens = data.input + data.output;
      const usagePercent =
        this.budget.perAgentTokenLimit > 0
          ? Math.round((totalTokens / this.budget.perAgentTokenLimit) * 10000) / 100
          : 0;

      return {
        agentId,
        taskId: data.taskId,
        inputTokens: data.input,
        outputTokens: data.output,
        totalTokens,
        usagePercent,
      };
    });
  }

  /**
   * Determine overall budget status
   */
  private getBudgetStatus(dailyUsagePercent: number): BudgetStatus {
    if (dailyUsagePercent >= 100) return "exceeded";
    if (dailyUsagePercent >= 95) return "at_limit";
    if (dailyUsagePercent >= 80) return "approaching_limit";
    return "within_budget";
  }
}