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 { 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 { 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 { 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 { 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 { 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", }; } }