/** * BudgetService Unit Tests * * Tests usage budget tracking, enforcement, and reporting. * Covers issue #329 (ORCH-135) including security hardening. */ 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) => undefined), }; 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 = { "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); }); it("should clamp negative config values to defaults", () => { const badConfig = { get: vi.fn((key: string) => { const config: Record = { "orchestrator.budget.dailyTokenLimit": -100, "orchestrator.budget.perAgentTokenLimit": -50, "orchestrator.budget.maxConcurrentAgents": -1, "orchestrator.budget.maxTaskDurationMinutes": 0, }; return config[key]; }), }; const badService = new BudgetService(badConfig as unknown as ConfigService); const budget = badService.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); }); it("should clamp NaN config values to defaults", () => { const nanConfig = { get: vi.fn((key: string) => { const config: Record = { "orchestrator.budget.dailyTokenLimit": NaN, "orchestrator.budget.maxConcurrentAgents": Infinity, }; return config[key]; }), }; const nanService = new BudgetService(nanConfig as unknown as ConfigService); const budget = nanService.getBudget(); expect(budget.dailyTokenLimit).toBe(10_000_000); expect(budget.maxConcurrentAgents).toBe(10); }); }); 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); }); it("should clamp negative token values to 0", () => { service.recordUsage("agent-1", "task-1", -500, -200); const summary = service.getUsageSummary(); expect(summary.dailyTokensUsed).toBe(0); }); it("should clamp NaN token values to 0", () => { service.recordUsage("agent-1", "task-1", NaN, NaN); const summary = service.getUsageSummary(); expect(summary.dailyTokensUsed).toBe(0); }); it("should clamp Infinity token values to 0", () => { service.recordUsage("agent-1", "task-1", Infinity, -Infinity); const summary = service.getUsageSummary(); expect(summary.dailyTokensUsed).toBe(0); }); it("should skip recording when agentId is empty", () => { service.recordUsage("", "task-1", 1000, 500); const summary = service.getUsageSummary(); expect(summary.dailyTokensUsed).toBe(0); expect(summary.agentUsage).toHaveLength(0); }); it("should skip recording when taskId is empty", () => { service.recordUsage("agent-1", "", 1000, 500); const summary = service.getUsageSummary(); expect(summary.dailyTokensUsed).toBe(0); }); }); 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(`agent-${String(i)}`); } 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(`agent-${String(i)}`); } expect(service.canSpawnAgent().allowed).toBe(false); service.agentStopped("agent-0"); 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 = { "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("trySpawnAgent", () => { it("should atomically check and reserve slot", () => { const result = service.trySpawnAgent("agent-1"); expect(result.allowed).toBe(true); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(1); }); it("should block when at max concurrent and not reserve slot", () => { for (let i = 0; i < 10; i++) { service.trySpawnAgent(`agent-${String(i)}`); } const result = service.trySpawnAgent("agent-11"); expect(result.allowed).toBe(false); expect(result.reason).toContain("Maximum concurrent agents"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(10); }); it("should reject empty agent ID", () => { const result = service.trySpawnAgent(""); expect(result.allowed).toBe(false); expect(result.reason).toContain("Agent ID is required"); }); it("should not double-count same agent ID", () => { service.trySpawnAgent("agent-1"); service.trySpawnAgent("agent-1"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(1); }); }); 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 by ID", () => { service.agentStarted("agent-1"); service.agentStarted("agent-2"); service.agentStarted("agent-3"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(3); }); it("should decrement active count on stop", () => { service.agentStarted("agent-1"); service.agentStarted("agent-2"); service.agentStopped("agent-1"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(1); }); it("should handle stopping non-existent agent gracefully", () => { service.agentStopped("non-existent"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(0); }); it("should not double-count same agent ID on start", () => { service.agentStarted("agent-1"); service.agentStarted("agent-1"); const summary = service.getUsageSummary(); expect(summary.activeAgents).toBe(1); }); }); 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 = { "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 = { "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 = { "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 = { "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); }); }); });