stack/apps/api/src/token-budget/token-budget.service.ts

import { Injectable, Logger, NotFoundException } from "@nestjs/common";
import { PrismaService } from "../prisma/prisma.service";
import type { TokenBudget } from "@prisma/client";
import type { TaskComplexity, BudgetAnalysis } from "./interfaces";
import { COMPLEXITY_BUDGETS, BUDGET_THRESHOLDS } from "./interfaces";
import type { AllocateBudgetDto } from "./dto";
import { BudgetAnalysisDto } from "./dto";

/**
 * Token Budget Service
 * Tracks token usage and prevents premature done claims with significant budget remaining
 */
@Injectable()
export class TokenBudgetService {
  private readonly logger = new Logger(TokenBudgetService.name);

  constructor(private readonly prisma: PrismaService) {}

  /**
   * Allocate budget for a new task
   */
  async allocateBudget(dto: AllocateBudgetDto): Promise<TokenBudget> {
    this.logger.log(`Allocating ${String(dto.allocatedTokens)} tokens for task ${dto.taskId}`);

    const budget = await this.prisma.tokenBudget.create({
      data: {
        taskId: dto.taskId,
        workspaceId: dto.workspaceId,
        agentId: dto.agentId,
        allocatedTokens: dto.allocatedTokens,
        estimatedComplexity: dto.complexity,
      },
    });

    return budget;
  }

  /**
   * Update usage after agent response
   * Uses atomic increment operations to prevent race conditions
   */
  async updateUsage(
    taskId: string,
    inputTokens: number,
    outputTokens: number
  ): Promise<TokenBudget> {
    this.logger.debug(
      `Updating usage for task ${taskId}: +${String(inputTokens)} input, +${String(outputTokens)} output`
    );

    // First verify budget exists
    const budget = await this.prisma.tokenBudget.findUnique({
      where: { taskId },
    });

    if (!budget) {
      throw new NotFoundException(`Token budget not found for task ${taskId}`);
    }

    // Use atomic increment operations to prevent race conditions
    const totalIncrement = inputTokens + outputTokens;
    const newTotalTokens = budget.totalTokensUsed + totalIncrement;
    const utilization = newTotalTokens / budget.allocatedTokens;

    // Update budget with atomic increments
    const updatedBudget = await this.prisma.tokenBudget.update({
      where: { taskId },
      data: {
        inputTokensUsed: { increment: inputTokens },
        outputTokensUsed: { increment: outputTokens },
        totalTokensUsed: { increment: totalIncrement },
        budgetUtilization: utilization,
      },
    });

    return updatedBudget;
  }

  /**
   * Analyze budget for suspicious patterns
   */
  async analyzeBudget(taskId: string): Promise<BudgetAnalysis> {
    this.logger.debug(`Analyzing budget for task ${taskId}`);

    const budget = await this.prisma.tokenBudget.findUnique({
      where: { taskId },
    });

    if (!budget) {
      throw new NotFoundException(`Token budget not found for task ${taskId}`);
    }

    const usedTokens = budget.totalTokensUsed;
    const allocatedTokens = budget.allocatedTokens;
    const remainingTokens = allocatedTokens - usedTokens;
    const utilizationPercentage = (usedTokens / allocatedTokens) * 100;

    // Detect suspicious patterns
    const suspiciousPattern = this.detectSuspiciousPattern(budget);

    // Determine recommendation
    let recommendation: "accept" | "continue" | "review";
    if (suspiciousPattern.triggered) {
      if (suspiciousPattern.severity === "high") {
        recommendation = "continue";
      } else {
        recommendation = "review";
      }
    } else {
      recommendation = "accept";
    }

    return new BudgetAnalysisDto({
      taskId,
      allocatedTokens,
      usedTokens,
      remainingTokens,
      utilizationPercentage,
      suspiciousPattern: suspiciousPattern.triggered,
      suspiciousReason: suspiciousPattern.reason ?? null,
      recommendation,
    });
  }

  /**
   * Check if done claim is suspicious (>20% budget remaining)
   */
  async checkSuspiciousDoneClaim(
    taskId: string
  ): Promise<{ suspicious: boolean; reason?: string }> {
    this.logger.debug(`Checking done claim for task ${taskId}`);

    const budget = await this.prisma.tokenBudget.findUnique({
      where: { taskId },
    });

    if (!budget) {
      throw new NotFoundException(`Token budget not found for task ${taskId}`);
    }

    const suspiciousPattern = this.detectSuspiciousPattern(budget);

    if (suspiciousPattern.triggered && suspiciousPattern.reason) {
      return {
        suspicious: true,
        reason: suspiciousPattern.reason,
      };
    }

    if (suspiciousPattern.triggered) {
      return {
        suspicious: true,
      };
    }

    return { suspicious: false };
  }

  /**
   * Get budget utilization percentage
   */
  async getBudgetUtilization(taskId: string): Promise<number> {
    const budget = await this.prisma.tokenBudget.findUnique({
      where: { taskId },
    });

    if (!budget) {
      throw new NotFoundException(`Token budget not found for task ${taskId}`);
    }

    const utilizationPercentage = (budget.totalTokensUsed / budget.allocatedTokens) * 100;

    return utilizationPercentage;
  }

  /**
   * Mark task as completed
   */
  async markCompleted(taskId: string): Promise<void> {
    this.logger.log(`Marking budget as completed for task ${taskId}`);

    const budget = await this.prisma.tokenBudget.findUnique({
      where: { taskId },
    });

    if (!budget) {
      throw new NotFoundException(`Token budget not found for task ${taskId}`);
    }

    await this.prisma.tokenBudget.update({
      where: { taskId },
      data: {
        completedAt: new Date(),
      },
    });
  }

  /**
   * Get complexity-based budget allocation
   */
  getDefaultBudgetForComplexity(complexity: TaskComplexity): number {
    return COMPLEXITY_BUDGETS[complexity];
  }

  /**
   * Detect suspicious patterns in budget usage
   * @private
   */
  private detectSuspiciousPattern(budget: TokenBudget): {
    triggered: boolean;
    reason?: string;
    severity: "low" | "medium" | "high";
    recommendation: "accept" | "continue" | "review";
  } {
    const utilization = budget.totalTokensUsed / budget.allocatedTokens;
    const remainingPercentage = (1 - utilization) * 100;

    // Pattern 1: Very low utilization (<10%)
    if (utilization < BUDGET_THRESHOLDS.VERY_LOW_UTILIZATION) {
      return {
        triggered: true,
        reason: `Very low budget utilization (${(utilization * 100).toFixed(1)}%). This suggests minimal work was performed.`,
        severity: "high",
        recommendation: "continue",
      };
    }

    // Pattern 2: Done claimed with >20% budget remaining
    if (utilization < 1 - BUDGET_THRESHOLDS.SUSPICIOUS_REMAINING) {
      return {
        triggered: true,
        reason: `Task claimed done with ${remainingPercentage.toFixed(1)}% budget remaining (${String(budget.allocatedTokens - budget.totalTokensUsed)} tokens). This may indicate premature completion.`,
        severity: "medium",
        recommendation: "review",
      };
    }

    // Pattern 3: Extremely high utilization (>95%) - might indicate inefficiency
    if (utilization > BUDGET_THRESHOLDS.VERY_HIGH_UTILIZATION) {
      return {
        triggered: true,
        reason: `Very high budget utilization (${(utilization * 100).toFixed(1)}%). Task may need more budget or review for efficiency.`,
        severity: "low",
        recommendation: "review",
      };
    }

    return {
      triggered: false,
      severity: "low",
      recommendation: "accept",
    };
  }
}