Add agent budget controls

api/nexus/authorization/http_models.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from decimal import Decimal
 from uuid import UUID
 
 from pydantic import BaseModel, Field
@@ -28,3 +29,19 @@ class KillSwitchRequest(BaseModel):
 class KillSwitchResponse(BaseModel):
     agent_key: str
     halted: bool
+
+
+class BudgetRequest(BaseModel):
+    agent_key: str = Field(min_length=1, max_length=128)
+    budget_amount: Decimal = Field(ge=0)
+    currency: str = Field(default="USD", min_length=3, max_length=3)
+    window_seconds: int = Field(gt=0)
+    enabled: bool = True
+
+
+class BudgetResponse(BaseModel):
+    agent_key: str
+    budget_amount: Decimal
+    currency: str
+    window_seconds: int
+    enabled: bool

api/nexus/authorization/models.py

@@ -16,6 +16,7 @@ class DecisionReason(StrEnum):
     FAIL_CLOSED = "fail_closed"
     STATIC_DENY_TEST_TOOL = "denied: static deny test tool"
     KILL_SWITCH_ACTIVE = "denied: kill switch active"
+    BUDGET_EXCEEDED = "denied: budget exceeded"
 
 
 @dataclass(frozen=True)

api/nexus/authorization/routes.py

@@ -9,6 +9,8 @@
     AuthorizeResponse,
     KillSwitchRequest,
     KillSwitchResponse,
+    BudgetRequest,
+    BudgetResponse,
 )
 from nexus.authorization.models import AuthorizationRequest
 from nexus.authorization.service import AuthorizationService
@@ -29,6 +31,7 @@ def _authorization_service(request: Request) -> AuthorizationService:
     return AuthorizationService(
         request.app.state.decision_repo,
         request.app.state.kill_switch_repo,
+        request.app.state.budget_repo,
     )
 
 
@@ -116,3 +119,41 @@ async def set_kill_switch(request: Request):
         agent_key=payload.agent_key,
         halted=payload.halted,
     )
+
+
+@router.post("/budgets", response_model=BudgetResponse)
+async def set_budget(request: Request):
+    try:
+        body = await request.json()
+    except Exception:
+        return _error(400, "malformed_request", "request body is not valid JSON")
+
+    try:
+        payload = BudgetRequest.model_validate(body)
+    except PydanticValidationError as e:
+        details = [
+            {"field": ".".join(str(p) for p in err["loc"]), "issue": err["msg"]}
+            for err in e.errors()
+        ]
+        return _error(
+            422,
+            "validation_failed",
+            "request failed schema validation",
+            details,
+        )
+
+    await request.app.state.budget_repo.set_budget(
+        agent_key=payload.agent_key,
+        budget_amount=payload.budget_amount,
+        currency=payload.currency,
+        window_seconds=payload.window_seconds,
+        enabled=payload.enabled,
+    )
+
+    return BudgetResponse(
+        agent_key=payload.agent_key,
+        budget_amount=str(payload.budget_amount),
+        currency=payload.currency,
+        window_seconds=payload.window_seconds,
+        enabled=payload.enabled,
+    )

api/nexus/authorization/service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from datetime import UTC, datetime
+from decimal import Decimal
 from typing import Protocol
 from uuid import uuid4
 
@@ -26,14 +27,44 @@ async def is_halted(self, agent_key: str) -> bool:
         """Return whether authorization must be denied by kill switch."""
 
 
+class BudgetReader(Protocol):
+    async def get_budget(self, agent_key: str):
+        """Return active budget config for this agent, if any."""
+
+    async def spend_in_window(
+        self,
+        agent_key: str,
+        window_seconds: int,
+        currency: str,
+    ) -> Decimal:
+        """Return accumulated spend for this agent in the budget window."""
+
+
 class AuthorizationService:
     def __init__(
         self,
         writer: DecisionLogWriter,
         kill_switch: KillSwitchReader | None = None,
+        budget: BudgetReader | None = None,
     ):
         self._writer = writer
         self._kill_switch = kill_switch
+        self._budget = budget
+
+    async def _budget_exceeded(self, request: AuthorizationRequest) -> bool:
+        if self._budget is None:
+            return False
+
+        budget = await self._budget.get_budget(request.agent_key)
+        if budget is None:
+            return False
+
+        spent = await self._budget.spend_in_window(
+            request.agent_key,
+            budget.window_seconds,
+            budget.currency,
+        )
+        return spent >= budget.budget_amount
 
     async def authorize(
         self,
@@ -51,6 +82,14 @@ async def authorize(
                 reason=DecisionReason.KILL_SWITCH_ACTIVE,
                 decided_at=datetime.now(UTC),
             )
+        elif await self._budget_exceeded(request):
+            decision = AuthorizationDecision(
+                decision_id=uuid4(),
+                effect=DecisionEffect.DENY,
+                fail_mode_applied=False,
+                reason=DecisionReason.BUDGET_EXCEEDED,
+                decided_at=datetime.now(UTC),
+            )
         elif request.tool_name == "__nexus_deny_test__":
             decision = AuthorizationDecision(
                 decision_id=uuid4(),

api/nexus/main.py

@@ -35,6 +35,7 @@
 from nexus.tenancy.boot_assertion import assert_safe_role
 from nexus.read.routes import router as read_router
 from nexus.repositories.kill_switch import KillSwitchRepository
+from nexus.repositories.budget import BudgetRepository
 
 
 def create_app(settings: Settings | None = None) -> FastAPI:
@@ -54,6 +55,7 @@ async def lifespan(app: FastAPI):
         app.state.agent_run_repo = AgentRunRepository(pool, crypto, limits)
         app.state.decision_repo = DecisionRepository(pool)
         app.state.kill_switch_repo = KillSwitchRepository(pool)
+        app.state.budget_repo = BudgetRepository(pool)
         try:
             yield
         finally:

api/nexus/repositories/budget.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from decimal import Decimal
+
+from nexus.repositories.base import BaseRepository
+
+
+@dataclass(frozen=True)
+class BudgetControl:
+    agent_key: str
+    budget_amount: Decimal
+    currency: str
+    window_seconds: int
+
+
+class BudgetRepository(BaseRepository):
+    async def get_budget(self, agent_key: str) -> BudgetControl | None:
+        async with self._tx() as conn:
+            row = await conn.fetchrow(
+                """
+                SELECT agent_key, budget_amount, currency, window_seconds
+                FROM budget_controls
+                WHERE enabled = true
+                  AND agent_key = $1
+                """,
+                agent_key,
+            )
+
+            if row is None:
+                return None
+
+            return BudgetControl(
+                agent_key=row["agent_key"],
+                budget_amount=row["budget_amount"],
+                currency=row["currency"],
+                window_seconds=row["window_seconds"],
+            )
+
+    async def spend_in_window(
+        self,
+        agent_key: str,
+        window_seconds: int,
+        currency: str,
+    ) -> Decimal:
+        async with self._tx() as conn:
+            total = await conn.fetchval(
+                """
+                SELECT COALESCE(SUM(cost_total), 0)
+                FROM agent_runs
+                WHERE agent_key = $1
+                  AND cost_currency = $2
+                  AND ended_at >= now() - ($3 * interval '1 second')
+                """,
+                agent_key,
+                currency,
+                window_seconds,
+            )
+            return Decimal(total)
+
+    async def set_budget(
+        self,
+        *,
+        agent_key: str,
+        budget_amount: Decimal,
+        currency: str = "USD",
+        window_seconds: int,
+        enabled: bool = True,
+    ) -> None:
+        async with self._tx() as conn:
+            await conn.execute(
+                """
+                INSERT INTO budget_controls
+                  (tenant_id, agent_key, budget_amount, currency,
+                   window_seconds, enabled, updated_at)
+                VALUES (current_setting('app.tenant_id')::uuid,
+                        $1, $2, $3, $4, $5, now())
+                ON CONFLICT (tenant_id, agent_key)
+                DO UPDATE SET
+                    budget_amount = EXCLUDED.budget_amount,
+                    currency = EXCLUDED.currency,
+                    window_seconds = EXCLUDED.window_seconds,
+                    enabled = EXCLUDED.enabled,
+                    updated_at = now()
+                """,
+                agent_key,
+                budget_amount,
+                currency,
+                window_seconds,
+                enabled,
+            )

api/tests/conftest.py

@@ -119,9 +119,12 @@ async def _seeded_key():
                 tenant_id,
             )
             for tbl in (
+                "budget_controls",
+                "decisions",
                 "agent_run_payloads",
                 "run_events",
                 "run_tool_calls",
+                "run_scores",
                 "agent_runs",
                 "agents",
             ):

api/tests/test_authorization_service.py

@@ -1,3 +1,5 @@
+from dataclasses import dataclass
+from decimal import Decimal
 from uuid import uuid4
 
 import pytest
@@ -45,6 +47,25 @@ def make_request():
     )
 
 
+@dataclass(frozen=True)
+class FakeBudgetControl:
+    budget_amount: Decimal
+    currency: str = "USD"
+    window_seconds: int = 3600
+
+
+class FakeBudget:
+    def __init__(self, budget=None, spent=Decimal("0")):
+        self.budget = budget
+        self.spent = Decimal(spent)
+
+    async def get_budget(self, agent_key):
+        return self.budget
+
+    async def spend_in_window(self, agent_key, window_seconds, currency):
+        return self.spent
+
+
 @pytest.mark.asyncio
 async def test_authorization_defaults_to_allow():
     writer = FakeWriter()
@@ -132,3 +153,57 @@ async def test_kill_switch_decision_is_logged():
     assert logged_decision.effect == DecisionEffect.DENY
     assert logged_decision.reason == DecisionReason.KILL_SWITCH_ACTIVE
     assert logged_decision.decision_id == decision.decision_id
+
+
+@pytest.mark.asyncio
+async def test_budget_breach_denies():
+    writer = FakeWriter()
+    service = AuthorizationService(
+        writer,
+        budget=FakeBudget(
+            FakeBudgetControl(budget_amount=Decimal("1.00")),
+            spent=Decimal("1.00"),
+        ),
+    )
+
+    decision = await service.authorize(make_request())
+
+    assert decision.effect == DecisionEffect.DENY
+    assert decision.fail_mode_applied is False
+    assert decision.reason == DecisionReason.BUDGET_EXCEEDED
+
+
+@pytest.mark.asyncio
+async def test_budget_under_threshold_allows():
+    writer = FakeWriter()
+    service = AuthorizationService(
+        writer,
+        budget=FakeBudget(
+            FakeBudgetControl(budget_amount=Decimal("1.00")),
+            spent=Decimal("0.99"),
+        ),
+    )
+
+    decision = await service.authorize(make_request())
+
+    assert decision.effect == DecisionEffect.ALLOW
+    assert decision.fail_mode_applied is False
+    assert decision.reason == DecisionReason.DEFAULT_ALLOW
+
+
+@pytest.mark.asyncio
+async def test_kill_switch_overrides_budget():
+    writer = FakeWriter()
+    service = AuthorizationService(
+        writer,
+        FakeKillSwitch(True),
+        FakeBudget(
+            FakeBudgetControl(budget_amount=Decimal("1.00")),
+            spent=Decimal("1.00"),
+        ),
+    )
+
+    decision = await service.authorize(make_request())
+
+    assert decision.effect == DecisionEffect.DENY
+    assert decision.reason == DecisionReason.KILL_SWITCH_ACTIVE