Agent/tests/test_recursive_resource_budget.py

import asyncio
from datetime import datetime, timedelta, timezone
import json
from pathlib import Path
import tempfile
import unittest

from cyber_agent.core.resource_budget import (
    ResourceBudget,
    ResourceBudgetController,
    ResourceBudgetExceeded,
    ResourceBudgetStateError,
    ResourceUsage,
)
from cyber_agent.trace.store import FileSystemTraceStore


class MutableClock:
    def __init__(self) -> None:
        self.current = datetime(2026, 1, 1, tzinfo=timezone.utc)

    def __call__(self) -> datetime:
        return self.current


class ResourceBudgetModelTest(unittest.TestCase):
    def test_environment_defaults_are_demo_limits(self):
        budget = ResourceBudget.from_environment({})
        self.assertTrue(budget.enabled)
        self.assertEqual(50, budget.max_total_agents)
        self.assertEqual(150, budget.max_llm_calls)
        self.assertEqual(1_500_000, budget.max_total_tokens)
        self.assertEqual(15.0, budget.max_total_cost_usd)
        self.assertEqual(3_600, budget.max_duration_seconds)
        self.assertEqual(1, budget.reserved_final_calls)

    def test_environment_is_strict_and_snapshot_round_trips(self):
        environ = {
            "AGENT_RESOURCE_BUDGET_ENABLED": "false",
            "AGENT_MAX_TOTAL_AGENTS": "8",
            "AGENT_MAX_LLM_CALLS": "12",
            "AGENT_MAX_TOTAL_TOKENS": "2000",
            "AGENT_MAX_TOTAL_COST_USD": "1.25",
            "AGENT_MAX_DURATION_SECONDS": "90",
            "AGENT_RESERVED_FINAL_CALLS": "2",
        }
        budget = ResourceBudget.from_environment(environ)
        self.assertEqual(budget, ResourceBudget.from_dict(budget.to_dict()))

        invalid = dict(environ, AGENT_RESOURCE_BUDGET_ENABLED="yes")
        with self.assertRaisesRegex(ValueError, "must be 'true' or 'false'"):
            ResourceBudget.from_environment(invalid)
        for name, value in (
            ("AGENT_MAX_TOTAL_AGENTS", "0"),
            ("AGENT_MAX_LLM_CALLS", "1.2"),
            ("AGENT_MAX_TOTAL_TOKENS", "-1"),
            ("AGENT_MAX_TOTAL_COST_USD", "nan"),
            ("AGENT_MAX_DURATION_SECONDS", "none"),
            ("AGENT_RESERVED_FINAL_CALLS", "0"),
        ):
            with self.subTest(name=name):
                invalid = dict(environ, **{name: value})
                with self.assertRaises(ValueError):
                    ResourceBudget.from_environment(invalid)

    def test_reserved_calls_must_fit(self):
        with self.assertRaisesRegex(ValueError, "less than max_llm_calls"):
            ResourceBudget(max_llm_calls=2, reserved_final_calls=2)

    def test_resource_usage_rejects_inconsistent_or_unknown_state(self):
        usage = ResourceUsage.new(total_agents=1)
        data = usage.to_dict()
        data["total_tokens"] = 1
        with self.assertRaisesRegex(ResourceBudgetStateError, "total_tokens"):
            ResourceUsage.from_dict(data)
        data = usage.to_dict()
        data["future_field"] = True
        with self.assertRaisesRegex(ResourceBudgetStateError, "unknown"):
            ResourceUsage.from_dict(data)


class ResourceBudgetControllerTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp_dir.name)
        self.clock = MutableClock()
        self.controller = ResourceBudgetController(self.store, now=self.clock)
        self.root_trace_id = "root-budget-test"

    async def asyncTearDown(self):
        self.temp_dir.cleanup()

    async def test_initialization_is_idempotent_and_file_is_atomic_json(self):
        budget = ResourceBudget()
        first = await self.controller.initialize(self.root_trace_id, budget)
        second = await self.controller.initialize(
            self.root_trace_id, budget, initial_agents=9
        )
        self.assertEqual(first, second)
        self.assertEqual(1, second.total_agents)

        path = Path(self.temp_dir.name) / self.root_trace_id / "resource_usage.json"
        self.assertEqual(second.to_dict(), json.loads(path.read_text(encoding="utf-8")))
        self.assertEqual([], list(path.parent.glob(".resource_usage.*.tmp")))

        with self.assertRaises(TypeError):
            await self.store.replace_resource_usage(
                self.root_trace_id, {"not_json_serializable": object()}
            )
        self.assertEqual(second.to_dict(), json.loads(path.read_text(encoding="utf-8")))
        self.assertEqual([], list(path.parent.glob(".resource_usage.*.tmp")))

    async def test_missing_or_corrupt_usage_fails_closed(self):
        with self.assertRaisesRegex(ResourceBudgetStateError, "missing"):
            await self.controller.get_usage("missing-root")

        root_dir = Path(self.temp_dir.name) / self.root_trace_id
        root_dir.mkdir()
        (root_dir / "resource_usage.json").write_text("not-json", encoding="utf-8")
        with self.assertRaisesRegex(ResourceBudgetStateError, "cannot be read"):
            await self.controller.get_usage(self.root_trace_id)

    async def test_agent_batch_is_all_or_nothing(self):
        budget = ResourceBudget(max_total_agents=3)
        await self.controller.initialize(self.root_trace_id, budget)
        usage = await self.controller.reserve_agents(self.root_trace_id, budget, 2)
        self.assertEqual(3, usage.total_agents)

        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.reserve_agents(self.root_trace_id, budget, 1)
        self.assertEqual("agents", caught.exception.dimension)
        usage = await self.controller.get_usage(self.root_trace_id)
        self.assertEqual(3, usage.total_agents)
        self.assertEqual("budget_exhausted:agents", usage.exhausted_reason)

    async def test_uncreated_agent_reservation_can_be_released_only_to_root(self):
        budget = ResourceBudget(max_total_agents=5)
        await self.controller.initialize(self.root_trace_id, budget)
        await self.controller.reserve_agents(self.root_trace_id, budget, 3)
        usage = await self.controller.release_agents(self.root_trace_id, budget, 2)
        self.assertEqual(2, usage.total_agents)
        with self.assertRaises(ResourceBudgetStateError):
            await self.controller.release_agents(self.root_trace_id, budget, 2)

    async def test_concurrent_agent_reservations_never_cross_limit(self):
        budget = ResourceBudget(max_total_agents=4)
        await self.controller.initialize(self.root_trace_id, budget)

        results = await asyncio.gather(
            *(self.controller.reserve_agents(self.root_trace_id, budget, 1) for _ in range(8)),
            return_exceptions=True,
        )
        successes = [result for result in results if isinstance(result, ResourceUsage)]
        failures = [result for result in results if isinstance(result, ResourceBudgetExceeded)]
        self.assertEqual(3, len(successes))
        self.assertEqual(5, len(failures))
        self.assertEqual(4, (await self.controller.get_usage(self.root_trace_id)).total_agents)

    async def test_ordinary_calls_preserve_final_validator_slot(self):
        budget = ResourceBudget(max_llm_calls=4, reserved_final_calls=1)
        await self.controller.initialize(self.root_trace_id, budget)
        for _ in range(3):
            await self.controller.reserve_llm_call(self.root_trace_id, budget)

        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.reserve_llm_call(self.root_trace_id, budget)
        self.assertEqual("llm_calls", caught.exception.dimension)
        usage = await self.controller.reserve_llm_call(
            self.root_trace_id, budget, purpose="root_validator"
        )
        self.assertEqual(4, usage.llm_calls)
        with self.assertRaises(ResourceBudgetExceeded):
            await self.controller.reserve_llm_call(
                self.root_trace_id, budget, purpose="root_validator"
            )

    async def test_post_response_usage_is_persisted_before_exceeded(self):
        budget = ResourceBudget(
            max_total_tokens=10,
            max_total_cost_usd=1.0,
        )
        await self.controller.initialize(self.root_trace_id, budget)
        await self.controller.reserve_llm_call(self.root_trace_id, budget)

        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.record_llm_usage(
                self.root_trace_id,
                budget,
                prompt_tokens=8,
                completion_tokens=4,
                cost_usd=0.5,
            )
        self.assertEqual("tokens", caught.exception.dimension)
        usage = await self.controller.get_usage(self.root_trace_id)
        self.assertEqual(12, usage.total_tokens)
        self.assertEqual(0.5, usage.total_cost_usd)
        self.assertEqual("budget_exhausted:tokens", usage.exhausted_reason)
        with self.assertRaises(ResourceBudgetExceeded):
            await self.controller.reserve_agents(self.root_trace_id, budget, 1)
        with self.assertRaises(ResourceBudgetExceeded):
            await self.controller.reserve_llm_call(self.root_trace_id, budget)

    async def test_cost_and_duration_are_enforced(self):
        budget = ResourceBudget(max_total_cost_usd=0.1, max_duration_seconds=5)
        await self.controller.initialize(self.root_trace_id, budget)
        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.record_llm_usage(
                self.root_trace_id,
                budget,
                prompt_tokens=0,
                completion_tokens=0,
                cost_usd=0.2,
            )
        self.assertEqual("cost_usd", caught.exception.dimension)

        other = "duration-root"
        await self.controller.initialize(other, budget)
        self.clock.current += timedelta(seconds=6)
        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.check_time(other, budget)
        self.assertEqual("duration_seconds", caught.exception.dimension)

    async def test_external_llm_usage_is_recorded_before_call_limit_error(self):
        budget = ResourceBudget(max_llm_calls=2, reserved_final_calls=1)
        await self.controller.initialize(self.root_trace_id, budget)
        await self.controller.record_external_llm_usage(
            self.root_trace_id,
            budget,
            prompt_tokens=3,
            completion_tokens=2,
            cost_usd=0.01,
        )
        with self.assertRaises(ResourceBudgetExceeded) as caught:
            await self.controller.record_external_llm_usage(
                self.root_trace_id,
                budget,
                prompt_tokens=7,
                completion_tokens=4,
                cost_usd=0.02,
            )
        self.assertEqual("llm_calls", caught.exception.dimension)
        usage = await self.controller.get_usage(self.root_trace_id)
        self.assertEqual(2, usage.llm_calls)
        self.assertEqual(16, usage.total_tokens)
        self.assertAlmostEqual(0.03, usage.total_cost_usd)

        other = "external-then-validator"
        await self.controller.initialize(other, budget)
        await self.controller.record_external_llm_usage(
            other,
            budget,
            prompt_tokens=1,
            completion_tokens=1,
            cost_usd=0,
        )
        usage = await self.controller.reserve_llm_call(
            other, budget, purpose="root_validator"
        )
        self.assertEqual(2, usage.llm_calls)

    async def test_disabled_budget_records_usage_without_denial(self):
        budget = ResourceBudget(
            enabled=False,
            max_total_agents=1,
            max_llm_calls=2,
            max_total_tokens=1,
            max_total_cost_usd=0.01,
            max_duration_seconds=1,
        )
        await self.controller.initialize(self.root_trace_id, budget)
        await self.controller.reserve_agents(self.root_trace_id, budget, 5)
        for _ in range(4):
            await self.controller.reserve_llm_call(self.root_trace_id, budget)
        await self.controller.record_llm_usage(
            self.root_trace_id,
            budget,
            prompt_tokens=100,
            completion_tokens=100,
            cost_usd=5.0,
        )
        self.clock.current += timedelta(seconds=100)
        usage = await self.controller.check_time(self.root_trace_id, budget)
        self.assertEqual(6, usage.total_agents)
        self.assertEqual(4, usage.llm_calls)
        self.assertEqual(200, usage.total_tokens)
        self.assertIsNone(usage.exhausted_reason)


if __name__ == "__main__":
    unittest.main()