From 9bb4e7239fc2ee3afe843b2e0ccce51255865d44 Mon Sep 17 00:00:00 2001 From: Baraa Abuarra Date: Mon, 8 Jun 2026 13:23:34 +0200 Subject: [PATCH] Add tool governance allowlist --- api/nexus/authorization/http_models.py | 12 +++ api/nexus/authorization/models.py | 1 + api/nexus/authorization/routes.py | 37 +++++++ api/nexus/authorization/service.py | 31 ++++++ api/nexus/main.py | 2 + api/nexus/repositories/tool_permissions.py | 59 ++++++++++++ api/tests/conftest.py | 1 + api/tests/test_authorization_service.py | 54 +++++++++++ api/tests/test_authorize_endpoint.py | 106 ++++++++++++++++++++- migrations/0017_tool_governance.sql | 32 +++++++ 10 files changed, 334 insertions(+), 1 deletion(-) create mode 100644 api/nexus/repositories/tool_permissions.py create mode 100644 migrations/0017_tool_governance.sql diff --git a/api/nexus/authorization/http_models.py b/api/nexus/authorization/http_models.py index e479e8d..17845b5 100644 --- a/api/nexus/authorization/http_models.py +++ b/api/nexus/authorization/http_models.py @@ -45,3 +45,15 @@ class BudgetResponse(BaseModel): currency: str window_seconds: int enabled: bool + + +class ToolPermissionRequest(BaseModel): + agent_key: str = Field(min_length=1, max_length=128) + tool_name: str = Field(min_length=1, max_length=128) + enabled: bool = True + + +class ToolPermissionResponse(BaseModel): + agent_key: str + tool_name: str + enabled: bool diff --git a/api/nexus/authorization/models.py b/api/nexus/authorization/models.py index 7e1b4ae..35ff3ec 100644 --- a/api/nexus/authorization/models.py +++ b/api/nexus/authorization/models.py @@ -17,6 +17,7 @@ class DecisionReason(StrEnum): STATIC_DENY_TEST_TOOL = "denied: static deny test tool" KILL_SWITCH_ACTIVE = "denied: kill switch active" BUDGET_EXCEEDED = "denied: budget exceeded" + TOOL_NOT_ALLOWED = "denied: tool not allowed" @dataclass(frozen=True) diff --git a/api/nexus/authorization/routes.py b/api/nexus/authorization/routes.py index 3e1eb62..040397c 100644 --- a/api/nexus/authorization/routes.py +++ b/api/nexus/authorization/routes.py @@ -11,6 +11,8 @@ KillSwitchResponse, BudgetRequest, BudgetResponse, + ToolPermissionRequest, + ToolPermissionResponse, ) from nexus.authorization.models import AuthorizationRequest from nexus.authorization.service import AuthorizationService @@ -32,6 +34,7 @@ def _authorization_service(request: Request) -> AuthorizationService: request.app.state.decision_repo, request.app.state.kill_switch_repo, request.app.state.budget_repo, + request.app.state.tool_permission_repo, ) @@ -157,3 +160,37 @@ async def set_budget(request: Request): window_seconds=payload.window_seconds, enabled=payload.enabled, ) + + +@router.post("/tool-permissions", response_model=ToolPermissionResponse) +async def set_tool_permission(request: Request): + try: + body = await request.json() + except Exception: + return _error(400, "malformed_request", "request body is not valid JSON") + + try: + payload = ToolPermissionRequest.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.tool_permission_repo.set_permission( + agent_key=payload.agent_key, + tool_name=payload.tool_name, + enabled=payload.enabled, + ) + + return ToolPermissionResponse( + agent_key=payload.agent_key, + tool_name=payload.tool_name, + enabled=payload.enabled, + ) diff --git a/api/nexus/authorization/service.py b/api/nexus/authorization/service.py index 0e7136f..4d5c652 100644 --- a/api/nexus/authorization/service.py +++ b/api/nexus/authorization/service.py @@ -40,16 +40,26 @@ async def spend_in_window( """Return accumulated spend for this agent in the budget window.""" +class ToolPermissionReader(Protocol): + async def has_any_grant(self, agent_key: str) -> bool: + """Return whether this agent has any configured tool grant.""" + + async def is_allowed(self, agent_key: str, tool_name: str) -> bool: + """Return whether this agent may invoke this tool.""" + + class AuthorizationService: def __init__( self, writer: DecisionLogWriter, kill_switch: KillSwitchReader | None = None, budget: BudgetReader | None = None, + tool_permissions: ToolPermissionReader | None = None, ): self._writer = writer self._kill_switch = kill_switch self._budget = budget + self._tool_permissions = tool_permissions async def _budget_exceeded(self, request: AuthorizationRequest) -> bool: if self._budget is None: @@ -66,6 +76,19 @@ async def _budget_exceeded(self, request: AuthorizationRequest) -> bool: ) return spent >= budget.budget_amount + async def _tool_not_allowed(self, request: AuthorizationRequest) -> bool: + if self._tool_permissions is None: + return False + + has_grant = await self._tool_permissions.has_any_grant(request.agent_key) + if not has_grant: + return False + + return not await self._tool_permissions.is_allowed( + request.agent_key, + request.tool_name, + ) + async def authorize( self, request: AuthorizationRequest, @@ -90,6 +113,14 @@ async def authorize( reason=DecisionReason.BUDGET_EXCEEDED, decided_at=datetime.now(UTC), ) + elif await self._tool_not_allowed(request): + decision = AuthorizationDecision( + decision_id=uuid4(), + effect=DecisionEffect.DENY, + fail_mode_applied=False, + reason=DecisionReason.TOOL_NOT_ALLOWED, + 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 ef81623..fbf5e77 100644 --- a/api/nexus/main.py +++ b/api/nexus/main.py @@ -36,6 +36,7 @@ from nexus.read.routes import router as read_router from nexus.repositories.kill_switch import KillSwitchRepository from nexus.repositories.budget import BudgetRepository +from nexus.repositories.tool_permissions import ToolPermissionRepository def create_app(settings: Settings | None = None) -> FastAPI: @@ -56,6 +57,7 @@ async def lifespan(app: FastAPI): app.state.decision_repo = DecisionRepository(pool) app.state.kill_switch_repo = KillSwitchRepository(pool) app.state.budget_repo = BudgetRepository(pool) + app.state.tool_permission_repo = ToolPermissionRepository(pool) try: yield finally: diff --git a/api/nexus/repositories/tool_permissions.py b/api/nexus/repositories/tool_permissions.py new file mode 100644 index 0000000..0122a57 --- /dev/null +++ b/api/nexus/repositories/tool_permissions.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from nexus.repositories.base import BaseRepository + + +class ToolPermissionRepository(BaseRepository): + async def has_any_grant(self, agent_key: str) -> bool: + async with self._tx() as conn: + exists = await conn.fetchval( + """ + SELECT 1 + FROM tool_permissions + WHERE agent_key = $1 + AND enabled = true + LIMIT 1 + """, + agent_key, + ) + return bool(exists) + + async def is_allowed(self, agent_key: str, tool_name: str) -> bool: + async with self._tx() as conn: + exists = await conn.fetchval( + """ + SELECT 1 + FROM tool_permissions + WHERE agent_key = $1 + AND tool_name = $2 + AND enabled = true + LIMIT 1 + """, + agent_key, + tool_name, + ) + return bool(exists) + + async def set_permission( + self, + *, + agent_key: str, + tool_name: str, + enabled: bool = True, + ) -> None: + async with self._tx() as conn: + await conn.execute( + """ + INSERT INTO tool_permissions + (tenant_id, agent_key, tool_name, enabled, updated_at) + VALUES (current_setting('app.tenant_id')::uuid, + $1, $2, $3, now()) + ON CONFLICT (tenant_id, agent_key, tool_name) + DO UPDATE SET + enabled = EXCLUDED.enabled, + updated_at = now() + """, + agent_key, + tool_name, + enabled, + ) diff --git a/api/tests/conftest.py b/api/tests/conftest.py index 233cef2..ce81c3f 100644 --- a/api/tests/conftest.py +++ b/api/tests/conftest.py @@ -119,6 +119,7 @@ async def _seeded_key(): tenant_id, ) for tbl in ( + "tool_permissions", "budget_controls", "decisions", "agent_run_payloads", diff --git a/api/tests/test_authorization_service.py b/api/tests/test_authorization_service.py index ca33b27..cdc6cbb 100644 --- a/api/tests/test_authorization_service.py +++ b/api/tests/test_authorization_service.py @@ -66,6 +66,18 @@ async def spend_in_window(self, agent_key, window_seconds, currency): return self.spent +class FakeToolPermissions: + def __init__(self, has_grant: bool, allowed: bool): + self.has_grant = has_grant + self.allowed = allowed + + async def has_any_grant(self, agent_key): + return self.has_grant + + async def is_allowed(self, agent_key, tool_name): + return self.allowed + + @pytest.mark.asyncio async def test_authorization_defaults_to_allow(): writer = FakeWriter() @@ -207,3 +219,45 @@ async def test_kill_switch_overrides_budget(): assert decision.effect == DecisionEffect.DENY assert decision.reason == DecisionReason.KILL_SWITCH_ACTIVE + + +@pytest.mark.asyncio +async def test_tool_outside_grant_denies(): + writer = FakeWriter() + service = AuthorizationService( + writer, + tool_permissions=FakeToolPermissions(has_grant=True, allowed=False), + ) + + decision = await service.authorize(make_request()) + + assert decision.effect == DecisionEffect.DENY + assert decision.reason == DecisionReason.TOOL_NOT_ALLOWED + + +@pytest.mark.asyncio +async def test_tool_inside_grant_allows(): + writer = FakeWriter() + service = AuthorizationService( + writer, + tool_permissions=FakeToolPermissions(has_grant=True, allowed=True), + ) + + decision = await service.authorize(make_request()) + + assert decision.effect == DecisionEffect.ALLOW + assert decision.reason == DecisionReason.DEFAULT_ALLOW + + +@pytest.mark.asyncio +async def test_absent_tool_grant_defaults_allow(): + writer = FakeWriter() + service = AuthorizationService( + writer, + tool_permissions=FakeToolPermissions(has_grant=False, allowed=False), + ) + + decision = await service.authorize(make_request()) + + assert decision.effect == DecisionEffect.ALLOW + assert decision.reason == DecisionReason.DEFAULT_ALLOW \ No newline at end of file diff --git a/api/tests/test_authorize_endpoint.py b/api/tests/test_authorize_endpoint.py index 8bf0569..f653cd6 100644 --- a/api/tests/test_authorize_endpoint.py +++ b/api/tests/test_authorize_endpoint.py @@ -372,4 +372,108 @@ async def test_budget_deny_is_idempotent(client, auth_headers): 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 + assert r1.json()["reason"] == "denied: budget exceeded" + + +async def test_tool_outside_grant_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_grant = await client.post( + "/v1/tool-permissions", + json={ + "agent_key": agent_key, + "tool_name": "allowed_tool", + "enabled": True, + }, + headers=auth_headers, + ) + assert r_grant.status_code == 200, r_grant.text + + r = await client.post( + "/v1/authorize", + json=_authorize_body(tool_name="blocked_tool"), + headers=auth_headers, + ) + + assert r.status_code == 200, r.text + assert r.json()["effect"] == "deny" + assert r.json()["reason"] == "denied: tool not allowed" + + +async def test_tool_inside_grant_allows(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_grant = await client.post( + "/v1/tool-permissions", + json={ + "agent_key": agent_key, + "tool_name": "customer_lookup", + "enabled": True, + }, + headers=auth_headers, + ) + assert r_grant.status_code == 200, r_grant.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_tool_governance_deny_is_idempotent(client, auth_headers): + agent_key = "support-bot" + tool_call_id = str(uuid.uuid4()) + + 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/tool-permissions", + json={ + "agent_key": agent_key, + "tool_name": "allowed_tool", + "enabled": True, + }, + headers=auth_headers, + ) + + r1 = await client.post( + "/v1/authorize", + json=_authorize_body(tool_call_id, tool_name="blocked_tool"), + headers=auth_headers, + ) + r2 = await client.post( + "/v1/authorize", + json=_authorize_body(tool_call_id, tool_name="blocked_tool"), + 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: tool not allowed" \ No newline at end of file diff --git a/migrations/0017_tool_governance.sql b/migrations/0017_tool_governance.sql new file mode 100644 index 0000000..be3e00b --- /dev/null +++ b/migrations/0017_tool_governance.sql @@ -0,0 +1,32 @@ +-- 0017_tool_governance.sql +-- Run as nexus_owner. +-- PR #34: tenant/agent-scoped tool governance. +-- +-- Static allow-list only. +-- No policy DSL, no tool catalog, no argument inspection. + +SET search_path = nexus, public; + +CREATE TABLE tool_permissions ( + tenant_id uuid NOT NULL, + agent_key text NOT NULL, + tool_name text NOT NULL, + enabled boolean NOT NULL DEFAULT true, + updated_at timestamptz NOT NULL DEFAULT now(), + + PRIMARY KEY (tenant_id, agent_key, tool_name), + + FOREIGN KEY (tenant_id, agent_key) + REFERENCES agents (tenant_id, agent_key) +); + +ALTER TABLE tool_permissions ENABLE ROW LEVEL SECURITY; +ALTER TABLE tool_permissions FORCE ROW LEVEL SECURITY; + +CREATE POLICY tool_permissions_tenant_isolation + ON tool_permissions + USING (tenant_id = current_setting('app.tenant_id')::uuid) + WITH CHECK (tenant_id = current_setting('app.tenant_id')::uuid); + +GRANT SELECT, INSERT, UPDATE ON tool_permissions TO nexus_app; +REVOKE DELETE ON tool_permissions FROM nexus_app;