stack/apps/orchestrator/src/budget/budget.service.spec.ts

/**
 * 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<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);
    });

    it("should clamp negative config values to defaults", () => {
      const badConfig = {
        get: vi.fn((key: string) => {
          const config: Record<string, unknown> = {
            "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<string, unknown> = {
            "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<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("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<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);
    });
  });
});