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:
296
apps/orchestrator/src/budget/budget.service.spec.ts
Normal file
296
apps/orchestrator/src/budget/budget.service.spec.ts
Normal file
@@ -0,0 +1,296 @@
|
||||
/**
|
||||
* BudgetService Unit Tests
|
||||
*
|
||||
* Tests usage budget tracking, enforcement, and reporting.
|
||||
* Covers issue #329 (ORCH-135)
|
||||
*/
|
||||
import { describe, it, expect, beforeEach, vi } from "vitest";
|
||||
import { BudgetService } from "./budget.service";
|
||||
import { ConfigService } from "@nestjs/config";
|
||||
|
||||
describe("BudgetService", () => {
|
||||
let service: BudgetService;
|
||||
|
||||
const mockConfigService = {
|
||||
get: vi.fn((_key: string, defaultValue?: unknown) => defaultValue),
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
service = new BudgetService(mockConfigService as unknown as ConfigService);
|
||||
});
|
||||
|
||||
describe("initialization", () => {
|
||||
it("should initialize with default budget values", () => {
|
||||
const budget = service.getBudget();
|
||||
expect(budget.dailyTokenLimit).toBe(10_000_000);
|
||||
expect(budget.perAgentTokenLimit).toBe(2_000_000);
|
||||
expect(budget.maxConcurrentAgents).toBe(10);
|
||||
expect(budget.maxTaskDurationMinutes).toBe(120);
|
||||
expect(budget.enforceHardLimits).toBe(false);
|
||||
});
|
||||
|
||||
it("should use config values when provided", () => {
|
||||
const customConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 5_000_000,
|
||||
"orchestrator.budget.perAgentTokenLimit": 1_000_000,
|
||||
"orchestrator.budget.maxConcurrentAgents": 5,
|
||||
"orchestrator.budget.maxTaskDurationMinutes": 60,
|
||||
"orchestrator.budget.enforceHardLimits": true,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const customService = new BudgetService(customConfig as unknown as ConfigService);
|
||||
const budget = customService.getBudget();
|
||||
|
||||
expect(budget.dailyTokenLimit).toBe(5_000_000);
|
||||
expect(budget.perAgentTokenLimit).toBe(1_000_000);
|
||||
expect(budget.maxConcurrentAgents).toBe(5);
|
||||
expect(budget.maxTaskDurationMinutes).toBe(60);
|
||||
expect(budget.enforceHardLimits).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("recordUsage", () => {
|
||||
it("should record token usage", () => {
|
||||
service.recordUsage("agent-1", "task-1", 1000, 500);
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.dailyTokensUsed).toBe(1500);
|
||||
});
|
||||
|
||||
it("should accumulate usage across multiple records", () => {
|
||||
service.recordUsage("agent-1", "task-1", 1000, 500);
|
||||
service.recordUsage("agent-1", "task-1", 2000, 1000);
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.dailyTokensUsed).toBe(4500);
|
||||
});
|
||||
|
||||
it("should track usage per agent", () => {
|
||||
service.recordUsage("agent-1", "task-1", 1000, 500);
|
||||
service.recordUsage("agent-2", "task-2", 3000, 1500);
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.agentUsage).toHaveLength(2);
|
||||
|
||||
const agent1 = summary.agentUsage.find((a) => a.agentId === "agent-1");
|
||||
const agent2 = summary.agentUsage.find((a) => a.agentId === "agent-2");
|
||||
|
||||
expect(agent1?.totalTokens).toBe(1500);
|
||||
expect(agent2?.totalTokens).toBe(4500);
|
||||
});
|
||||
});
|
||||
|
||||
describe("canSpawnAgent", () => {
|
||||
it("should allow spawning when under limits", () => {
|
||||
const result = service.canSpawnAgent();
|
||||
expect(result.allowed).toBe(true);
|
||||
});
|
||||
|
||||
it("should block spawning when at max concurrent agents", () => {
|
||||
for (let i = 0; i < 10; i++) {
|
||||
service.agentStarted();
|
||||
}
|
||||
|
||||
const result = service.canSpawnAgent();
|
||||
expect(result.allowed).toBe(false);
|
||||
expect(result.reason).toContain("Maximum concurrent agents");
|
||||
});
|
||||
|
||||
it("should allow spawning after agent stops", () => {
|
||||
for (let i = 0; i < 10; i++) {
|
||||
service.agentStarted();
|
||||
}
|
||||
|
||||
expect(service.canSpawnAgent().allowed).toBe(false);
|
||||
|
||||
service.agentStopped();
|
||||
|
||||
expect(service.canSpawnAgent().allowed).toBe(true);
|
||||
});
|
||||
|
||||
it("should block spawning when daily budget exceeded with hard limits", () => {
|
||||
const strictConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 1000,
|
||||
"orchestrator.budget.enforceHardLimits": true,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const strictService = new BudgetService(strictConfig as unknown as ConfigService);
|
||||
strictService.recordUsage("agent-1", "task-1", 800, 300);
|
||||
|
||||
const result = strictService.canSpawnAgent();
|
||||
expect(result.allowed).toBe(false);
|
||||
expect(result.reason).toContain("Daily token budget exceeded");
|
||||
});
|
||||
|
||||
it("should allow spawning when over budget without hard limits", () => {
|
||||
service.recordUsage("agent-1", "task-1", 5_000_000, 5_000_000);
|
||||
|
||||
const result = service.canSpawnAgent();
|
||||
expect(result.allowed).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("isAgentOverBudget", () => {
|
||||
it("should return false when agent is within budget", () => {
|
||||
service.recordUsage("agent-1", "task-1", 100_000, 50_000);
|
||||
|
||||
const result = service.isAgentOverBudget("agent-1");
|
||||
expect(result.overBudget).toBe(false);
|
||||
expect(result.totalTokens).toBe(150_000);
|
||||
});
|
||||
|
||||
it("should return true when agent exceeds per-agent limit", () => {
|
||||
service.recordUsage("agent-1", "task-1", 1_000_000, 1_000_000);
|
||||
|
||||
const result = service.isAgentOverBudget("agent-1");
|
||||
expect(result.overBudget).toBe(true);
|
||||
expect(result.totalTokens).toBe(2_000_000);
|
||||
});
|
||||
|
||||
it("should return false for unknown agent", () => {
|
||||
const result = service.isAgentOverBudget("non-existent");
|
||||
expect(result.overBudget).toBe(false);
|
||||
expect(result.totalTokens).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("agentStarted / agentStopped", () => {
|
||||
it("should track active agent count", () => {
|
||||
service.agentStarted();
|
||||
service.agentStarted();
|
||||
service.agentStarted();
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.activeAgents).toBe(3);
|
||||
});
|
||||
|
||||
it("should decrement active count on stop", () => {
|
||||
service.agentStarted();
|
||||
service.agentStarted();
|
||||
service.agentStopped();
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.activeAgents).toBe(1);
|
||||
});
|
||||
|
||||
it("should not go below zero", () => {
|
||||
service.agentStopped();
|
||||
service.agentStopped();
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
expect(summary.activeAgents).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getUsageSummary", () => {
|
||||
it("should return complete summary with no usage", () => {
|
||||
const summary = service.getUsageSummary();
|
||||
|
||||
expect(summary.dailyTokensUsed).toBe(0);
|
||||
expect(summary.dailyTokenLimit).toBe(10_000_000);
|
||||
expect(summary.dailyUsagePercent).toBe(0);
|
||||
expect(summary.agentUsage).toHaveLength(0);
|
||||
expect(summary.activeAgents).toBe(0);
|
||||
expect(summary.maxConcurrentAgents).toBe(10);
|
||||
expect(summary.budgetStatus).toBe("within_budget");
|
||||
});
|
||||
|
||||
it("should calculate usage percentage correctly", () => {
|
||||
const customConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 10_000,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const customService = new BudgetService(customConfig as unknown as ConfigService);
|
||||
customService.recordUsage("agent-1", "task-1", 5000, 0);
|
||||
|
||||
const summary = customService.getUsageSummary();
|
||||
expect(summary.dailyUsagePercent).toBe(50);
|
||||
});
|
||||
|
||||
it("should report 'approaching_limit' at 80%", () => {
|
||||
const customConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 10_000,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const customService = new BudgetService(customConfig as unknown as ConfigService);
|
||||
customService.recordUsage("agent-1", "task-1", 8500, 0);
|
||||
|
||||
const summary = customService.getUsageSummary();
|
||||
expect(summary.budgetStatus).toBe("approaching_limit");
|
||||
});
|
||||
|
||||
it("should report 'at_limit' at 95%", () => {
|
||||
const customConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 10_000,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const customService = new BudgetService(customConfig as unknown as ConfigService);
|
||||
customService.recordUsage("agent-1", "task-1", 9600, 0);
|
||||
|
||||
const summary = customService.getUsageSummary();
|
||||
expect(summary.budgetStatus).toBe("at_limit");
|
||||
});
|
||||
|
||||
it("should report 'exceeded' over 100%", () => {
|
||||
const customConfig = {
|
||||
get: vi.fn((key: string) => {
|
||||
const config: Record<string, unknown> = {
|
||||
"orchestrator.budget.dailyTokenLimit": 10_000,
|
||||
};
|
||||
return config[key];
|
||||
}),
|
||||
};
|
||||
|
||||
const customService = new BudgetService(customConfig as unknown as ConfigService);
|
||||
customService.recordUsage("agent-1", "task-1", 11_000, 0);
|
||||
|
||||
const summary = customService.getUsageSummary();
|
||||
expect(summary.budgetStatus).toBe("exceeded");
|
||||
});
|
||||
|
||||
it("should calculate per-agent usage percentage", () => {
|
||||
service.recordUsage("agent-1", "task-1", 500_000, 500_000);
|
||||
|
||||
const summary = service.getUsageSummary();
|
||||
const agent = summary.agentUsage.find((a) => a.agentId === "agent-1");
|
||||
|
||||
expect(agent?.usagePercent).toBe(50);
|
||||
});
|
||||
});
|
||||
|
||||
describe("getBudget", () => {
|
||||
it("should return a copy of the budget", () => {
|
||||
const budget1 = service.getBudget();
|
||||
const budget2 = service.getBudget();
|
||||
|
||||
expect(budget1).toEqual(budget2);
|
||||
expect(budget1).not.toBe(budget2); // Different reference
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user