From 1541b9df6dac56ac12d51da278186439d0524c41 Mon Sep 17 00:00:00 2001 From: Baraa Abuarra Date: Mon, 8 Jun 2026 12:15:27 +0200 Subject: [PATCH] Add agent budget controls --- api/nexus/authorization/http_models.py | 17 +++ api/nexus/authorization/models.py | 1 + api/nexus/authorization/routes.py | 41 +++++++ api/nexus/authorization/service.py | 39 +++++++ api/nexus/main.py | 2 + api/nexus/repositories/budget.py | 91 +++++++++++++++ api/tests/conftest.py | 3 + api/tests/test_authorization_service.py | 75 +++++++++++++ api/tests/test_authorize_endpoint.py | 140 +++++++++++++++++++++++- migrations/0016_budget_controls.sql | 38 +++++++ 10 files changed, 446 insertions(+), 1 deletion(-) create mode 100644 api/nexus/repositories/budget.py create mode 100644 migrations/0016_budget_controls.sql diff --git a/api/nexus/authorization/http_models.py b/api/nexus/authorization/http_models.py index 94cb686..e479e8d 100644 --- a/api/nexus/authorization/http_models.py +++ b/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 diff --git a/api/nexus/authorization/models.py b/api/nexus/authorization/models.py index 0be65ca..7e1b4ae 100644 --- a/api/nexus/authorization/models.py +++ b/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) diff --git a/api/nexus/authorization/routes.py b/api/nexus/authorization/routes.py index f0a30a1..3e1eb62 100644 --- a/api/nexus/authorization/routes.py +++ b/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, + ) diff --git a/api/nexus/authorization/service.py b/api/nexus/authorization/service.py index 241c7f7..0e7136f 100644 --- a/api/nexus/authorization/service.py +++ b/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(), diff --git a/api/nexus/main.py b/api/nexus/main.py index 670cb02..ef81623 100644 --- a/api/nexus/main.py +++ b/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: diff --git a/api/nexus/repositories/budget.py b/api/nexus/repositories/budget.py new file mode 100644 index 0000000..0e7e9c5 --- /dev/null +++ b/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, + ) diff --git a/api/tests/conftest.py b/api/tests/conftest.py index 76cdc9e..233cef2 100644 --- a/api/tests/conftest.py +++ b/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", ): diff --git a/api/tests/test_authorization_service.py b/api/tests/test_authorization_service.py index 8c665b1..ca33b27 100644 --- a/api/tests/test_authorization_service.py +++ b/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 diff --git a/api/tests/test_authorize_endpoint.py b/api/tests/test_authorize_endpoint.py index 7e8d26e..8bf0569 100644 --- a/api/tests/test_authorize_endpoint.py +++ b/api/tests/test_authorize_endpoint.py @@ -17,6 +17,19 @@ def _authorize_body( "tool_call_id": tool_call_id or str(uuid.uuid4()), } +def _run_body(agent_key="support-bot", cost="0.00"): + return { + "run_id": str(uuid.uuid4()), + "agent": {"key": agent_key}, + "status": "completed", + "started_at": "2026-01-01T00:00:00Z", + "ended_at": "2026-01-01T00:00:01Z", + "cost": {"total": cost, "currency": "USD"}, + "events": [], + "tool_calls": [], + "payloads": [], + } + async def test_authorize_requires_authentication(client): r = await client.post("/v1/authorize", json=_authorize_body()) @@ -234,4 +247,129 @@ async def test_kill_switch_deny_is_idempotent(client, auth_headers): assert r1.json()["effect"] == "deny" assert r2.json()["effect"] == "deny" assert r1.json()["reason"] == r2.json()["reason"] - assert r1.json()["reason"] == "denied: kill switch active" \ No newline at end of file + assert r1.json()["reason"] == "denied: kill switch active" + + +async def test_budget_breach_forces_deny(client, auth_headers): + agent_key = "support-bot" + r_seed = await client.post( + "/v1/runs", + json=_run_body(agent_key=agent_key), + headers=auth_headers, + ) + + assert r_seed.status_code == 200, r_seed.text + + r_budget = await client.post( + "/v1/budgets", + json={ + "agent_key": agent_key, + "budget_amount": "0.00", + "currency": "USD", + "window_seconds": 3600, + "enabled": True, + }, + headers=auth_headers, + ) + + assert r_budget.status_code == 200, r_budget.text + + r = await client.post( + "/v1/authorize", + json=_authorize_body(tool_name="customer_lookup"), + headers=auth_headers, + ) + + assert r.status_code == 200, r.text + assert r.json()["effect"] == "deny" + assert r.json()["reason"] == "denied: budget exceeded" + + +async def test_budget_disabled_restores_default_allow(client, auth_headers): + agent_key = "support-bot" + + r_seed = await client.post( + "/v1/runs", + json=_run_body(agent_key=agent_key), + headers=auth_headers, + ) + + assert r_seed.status_code == 200, r_seed.text + + await client.post( + "/v1/budgets", + json={ + "agent_key": agent_key, + "budget_amount": "0.00", + "currency": "USD", + "window_seconds": 3600, + "enabled": True, + }, + headers=auth_headers, + ) + + r_disable = await client.post( + "/v1/budgets", + json={ + "agent_key": agent_key, + "budget_amount": "0.00", + "currency": "USD", + "window_seconds": 3600, + "enabled": False, + }, + headers=auth_headers, + ) + + assert r_disable.status_code == 200, r_disable.text + + r = await client.post( + "/v1/authorize", + json=_authorize_body(tool_name="customer_lookup"), + headers=auth_headers, + ) + + assert r.status_code == 200, r.text + assert r.json()["effect"] == "allow" + assert r.json()["reason"] == "default_allow" + + +async def test_budget_deny_is_idempotent(client, auth_headers): + tool_call_id = str(uuid.uuid4()) + + r_seed = await client.post( + "/v1/runs", + json=_run_body(agent_key="support-bot"), + headers=auth_headers, + ) + + assert r_seed.status_code == 200, r_seed.text + + await client.post( + "/v1/budgets", + json={ + "agent_key": "support-bot", + "budget_amount": "0.00", + "currency": "USD", + "window_seconds": 3600, + "enabled": True, + }, + headers=auth_headers, + ) + + r1 = await client.post( + "/v1/authorize", + json=_authorize_body(tool_call_id), + headers=auth_headers, + ) + r2 = await client.post( + "/v1/authorize", + json=_authorize_body(tool_call_id), + headers=auth_headers, + ) + + assert r1.status_code == 200, r1.text + assert r2.status_code == 200, r2.text + assert r1.json()["decision_id"] == r2.json()["decision_id"] + assert r1.json()["effect"] == "deny" + assert r2.json()["effect"] == "deny" + assert r1.json()["reason"] == "denied: budget exceeded" \ No newline at end of file diff --git a/migrations/0016_budget_controls.sql b/migrations/0016_budget_controls.sql new file mode 100644 index 0000000..848c37c --- /dev/null +++ b/migrations/0016_budget_controls.sql @@ -0,0 +1,38 @@ +-- 0016_budget_controls.sql +-- Run as nexus_owner. +-- PR #33: tenant-scoped budget controls. +-- +-- Budget configuration is tenant-isolated and agent-specific in PR #33 v1. +-- Spend state is derived from existing agent_runs.cost_total records. +-- This migration intentionally does not create a second spend ledger. + +SET search_path = nexus, public; + +CREATE TABLE budget_controls ( + tenant_id uuid NOT NULL, + agent_key text NOT NULL, + budget_amount numeric(18,8) NOT NULL, + currency text NOT NULL DEFAULT 'USD', + window_seconds integer NOT NULL, + enabled boolean NOT NULL DEFAULT true, + updated_at timestamptz NOT NULL DEFAULT now(), + + PRIMARY KEY (tenant_id, agent_key), + + FOREIGN KEY (tenant_id, agent_key) + REFERENCES agents (tenant_id, agent_key), + + CHECK (budget_amount >= 0), + CHECK (window_seconds > 0) +); + +ALTER TABLE budget_controls ENABLE ROW LEVEL SECURITY; +ALTER TABLE budget_controls FORCE ROW LEVEL SECURITY; + +CREATE POLICY budget_controls_tenant_isolation + ON budget_controls + USING (tenant_id = current_setting('app.tenant_id')::uuid) + WITH CHECK (tenant_id = current_setting('app.tenant_id')::uuid); + +GRANT SELECT, INSERT, UPDATE ON budget_controls TO nexus_app; +REVOKE DELETE ON budget_controls FROM nexus_app;