diff --git a/app/ai/agent_loop.py b/app/ai/agent_loop.py index 9feff45..c9cf5bb 100644 --- a/app/ai/agent_loop.py +++ b/app/ai/agent_loop.py @@ -122,6 +122,10 @@ async def _execute_tool( """Execute a single tool call via the registry. Returns the tool result as a string, or an error message. + + Note: access control (allowlist + required_permission) is enforced by + ``_check_tool_access`` in ``run_react_loop`` BEFORE any execution path + (dry-run, approval, execute) is reached. """ tool = tool_registry.get(tool_name) if tool is None: @@ -135,6 +139,89 @@ async def _execute_tool( return f"Error: {exc}" +def _extract_allowed_tool_names(tools: list[dict[str, Any]] | None) -> set[str]: + """Extract the tool names actually offered to the LLM. + + Always returns a set (possibly empty) — empty means no tools were + offered, so the caller fails closed on ANY tool call. + """ + if not tools: + return set() + names: set[str] = set() + for t in tools: + fn = (t or {}).get("function") or {} + name = fn.get("name") + if name: + names.add(str(name)) + return names + + +def _normalize_permissions(ctx: dict[str, Any] | None) -> dict[str, Any]: + """Normalize a user/agent context into the resolved-permissions shape + expected by ``check_permission`` (permissions / denied / is_system_admin). + + Session user contexts carry ``denied_permissions`` while resolved + permission dicts use ``denied`` — both are accepted here. + """ + if not isinstance(ctx, dict): + return {"permissions": [], "denied": [], "is_system_admin": False} + return { + "permissions": list(ctx.get("permissions", []) or []), + "denied": list(ctx.get("denied", ctx.get("denied_permissions", [])) or []), + "is_system_admin": bool(ctx.get("is_system_admin", False)), + } + + +def _check_tool_access( + tool_registry: ToolRegistry, + tool_name: str, + allowed_tools: set[str] | None, + user_permissions: dict[str, Any] | None, +) -> str | None: + """F01 guard: enforce allowlist + required_permission before execution. + + Must run before EVERY execution path (dry-run, approval, execute). + Returns None when access is granted, otherwise an error observation. + + Checks (fail-closed): + 1. Allowlist — the tool must be among the schemas actually offered to + the LLM. A hallucinated/injected tool name never reaches a handler. + 2. required_permission — when the tool declares one, the acting user's + CURRENT permissions must grant it. Without a permission context the + call is rejected (deny list first, system admin bypass). + """ + tool = tool_registry.get(tool_name) + if tool is None: + return None # "not found" is handled by _execute_tool + + # 1. Allowlist: only tools offered to the LLM may run. + if allowed_tools is not None and tool_name not in allowed_tools: + logger.warning( + "F01 guard: tool '%s' is registered but NOT offered to this agent — rejected", + tool_name, + ) + return f"Error: Tool '{tool_name}' is not available to this agent" + + # 2. required_permission: enforce against the user's CURRENT permissions. + required = getattr(tool, "required_permission", None) + if isinstance(required, str) and required: + resolved = _normalize_permissions(user_permissions) + from app.core.permissions import check_permission + + if not check_permission(resolved, required): + logger.warning( + "F01 guard: tool '%s' requires '%s' which the acting user lacks — rejected", + tool_name, + required, + ) + return ( + f"Error: Permission '{required}' required for tool '{tool_name}' " + "and not granted to the acting user" + ) + + return None + + async def run_react_loop( agent_definition: Any, # AgentDefinition from automation models messages: list[dict[str, Any]], @@ -151,6 +238,7 @@ async def run_react_loop( dry_run: bool = False, require_approval: bool = False, approval_tools: list[str] | None = None, + user_permissions: dict[str, Any] | None = None, ) -> ReActResult: """Execute a ReAct loop: LLM reasoning → tool execution → repeat. @@ -200,6 +288,10 @@ async def run_react_loop( "agent_name": getattr(agent_definition, "name", "Agent"), } + # F01 (Astra P0): allowlist — only tools actually offered to the LLM + # may ever execute. Empty set = no tools offered = every call rejected. + allowed_tool_names: set[str] = _extract_allowed_tool_names(tools) + # Audit helper — records every tool call in the audit log. async def _audit_tool_call( step_number: int, @@ -366,7 +458,15 @@ async def run_react_loop( args = {} logger.warning("Invalid JSON arguments for tool '%s': %s", tool_name, tc["arguments"]) - if dry_run: + # F01 (Astra P0): enforce allowlist + required_permission before + # every execution path (dry-run, approval, execute). Fail-closed: + # a hallucinated or injected tool name never reaches a handler. + guard_error = _check_tool_access( + tool_registry, tool_name, allowed_tool_names, user_permissions + ) + if guard_error is not None: + observation = guard_error + elif dry_run: observation = json.dumps( { "dry_run": True, diff --git a/app/ai/agent_stream.py b/app/ai/agent_stream.py index affc6c2..89d94b2 100644 --- a/app/ai/agent_stream.py +++ b/app/ai/agent_stream.py @@ -20,7 +20,8 @@ import asyncio import json import logging import uuid -from typing import TYPE_CHECKING, Any, AsyncGenerator +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Any from app.ai.agent_loop import ReActStep, run_react_loop @@ -71,6 +72,7 @@ async def stream_react_loop( max_steps: int = 20, timeout_seconds: int = 300, trace_id: str | None = None, + user_permissions: dict[str, Any] | None = None, ) -> AsyncGenerator[str, None]: """Run the ReAct loop and yield SSE-formatted events. @@ -116,6 +118,7 @@ async def stream_react_loop( timeout_seconds=timeout_seconds, trace_id=trace_id, on_step=on_step, + user_permissions=user_permissions, # F01: enforce at execution time ) await queue.put( _sse( diff --git a/app/core/permissions.py b/app/core/permissions.py index ef72db1..7e6f2f3 100644 --- a/app/core/permissions.py +++ b/app/core/permissions.py @@ -431,7 +431,12 @@ def check_permission(resolved: dict[str, Any], required: str) -> bool: return True permissions = set(resolved.get("permissions", [])) - denied = set(resolved.get("denied", [])) + # F01/Astra: session user contexts carry ``denied_permissions`` while + # resolved permission dicts use ``denied`` — accept both so the deny + # list is never silently ignored. + denied = set( + resolved.get("denied", resolved.get("denied_permissions", [])) or [] + ) # Check deny list first for d in denied: diff --git a/app/plugins/builtins/ai_assistant/services.py b/app/plugins/builtins/ai_assistant/services.py index 5f97abf..00ecbf8 100644 --- a/app/plugins/builtins/ai_assistant/services.py +++ b/app/plugins/builtins/ai_assistant/services.py @@ -234,6 +234,10 @@ async def stream_chat_comm( tools.append(crm_api_tool) tool_schemas = [t.to_openai_schema() for t in tools] if tools else None + # F01 (Astra P0): allowlist — only the tools offered above may execute. + # A hallucinated/injected tool name must never reach a handler. + allowed_tool_names = {t.name for t in tools} + # Build LLM params params, model_id = await build_litellm_params(db, agent, messages, tenant_id) @@ -290,8 +294,17 @@ async def stream_chat_comm( except json.JSONDecodeError: tool_args = {} + # F01 (Astra P0): allowlist enforcement — reject any tool + # name that was not offered to the LLM before it reaches a + # handler. registered ≠ permitted. tool = registry.get(tool_name) - if tool is None: + if tool_name not in allowed_tool_names: + logger.warning( + "stream_chat_comm F01 guard: tool '%s' is registered but NOT offered to this agent — rejected", + tool_name, + ) + result = f"Error: Tool '{tool_name}' is not available to this agent" + elif tool is None: result = f"Tool '{tool_name}' not found" else: result = await execute_tool_call(tool, tool_args, user_context) diff --git a/app/plugins/builtins/automation/agent_routes.py b/app/plugins/builtins/automation/agent_routes.py index 6f0291c..a1fac82 100644 --- a/app/plugins/builtins/automation/agent_routes.py +++ b/app/plugins/builtins/automation/agent_routes.py @@ -664,6 +664,7 @@ async def stream_agent_run( tenant_id=tenant_id, user_id=user_id, agent_run_id=aid, + user_permissions=current_user, # F01: enforce allowlist+permission at execution time ), media_type="text/event-stream", ) diff --git a/app/plugins/builtins/automation/agent_runner.py b/app/plugins/builtins/automation/agent_runner.py index 6ccd4d6..41af1d6 100644 --- a/app/plugins/builtins/automation/agent_runner.py +++ b/app/plugins/builtins/automation/agent_runner.py @@ -260,6 +260,7 @@ async def run_agent( timeout_seconds=max_duration, require_approval=bool(getattr(agent, "require_approval", False)), approval_tools=getattr(agent, "approval_tools", None), + user_permissions=perm_ctx.user_permissions, # F01: enforce at execution time ), timeout=max_duration + 10, # Extra buffer beyond loop's own timeout ) diff --git a/tests/test_agent_loop.py b/tests/test_agent_loop.py index b54e374..045d02b 100644 --- a/tests/test_agent_loop.py +++ b/tests/test_agent_loop.py @@ -16,7 +16,6 @@ import pytest from app.ai.agent_loop import ReActResult, ReActStep, run_react_loop - # ────────────────────────────────────────────────────────────────────────── # Test helpers # ────────────────────────────────────────────────────────────────────────── @@ -416,6 +415,306 @@ async def test_tool_not_found(agent_def, tenant_id, user_id, mock_db): assert "not found" in (result.steps[0].observation or "").lower() +def _make_guard_registry(tools: dict[str, Any]) -> MagicMock: + """Build a mock ToolRegistry with REAL attribute values. + + Unlike ``_make_tool_registry`` (whose MagicMock tools auto-create any + attribute), this sets ``required_permission`` explicitly so the F01 + guard behaves like production AITool objects. + """ + registry = MagicMock() + + def _get(name: str) -> Any: + return tools.get(name) + + registry.get = _get + return registry + + +def _guard_tool(name: str, handler: AsyncMock, required_permission: str | None = None) -> MagicMock: + t = MagicMock() + t.name = name + t.handler = handler + t.required_permission = required_permission + return t + + +# ────────────────────────────────────────────────────────────────────────── +# F01 (Astra P0): allowlist + required_permission enforced at execution time +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_f01_registered_but_not_offered_tool_rejected(agent_def, tenant_id, user_id, mock_db): + """F01 acceptance: a tool registered in the registry but NOT offered + to the LLM is rejected — the handler must never run (stays null).""" + handler = AsyncMock(return_value="secret audit data") + registry = _make_guard_registry({ + "audit_restricted": _guard_tool("audit_restricted", handler, required_permission="system:admin"), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="calling restricted tool", + tool_calls=[_make_tool_call(name="audit_restricted")], + cost_usd=0.001, + ), + _make_llm_response(content="understood", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "run audit_restricted"}], + tools=[{"type": "function", "function": {"name": "audit_allowed"}}], # only audit_allowed offered + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + user_permissions={"permissions": ["system:admin"], "denied": [], "is_system_admin": True}, + ) + + assert result.status == "completed" + handler.assert_not_awaited() # core F01 acceptance: handler stays null + assert "not available" in (result.steps[0].observation or "") + + +@pytest.mark.asyncio +async def test_f01_offered_tool_user_lacks_permission_rejected(agent_def, tenant_id, user_id, mock_db): + """Offered tool with required_permission: user WITHOUT the permission + is rejected even though the tool was offered to the LLM.""" + handler = AsyncMock(return_value="mail sent") + registry = _make_guard_registry({ + "send_mail": _guard_tool("send_mail", handler, required_permission="mail:write"), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="sending", + tool_calls=[_make_tool_call(name="send_mail")], + cost_usd=0.001, + ), + _make_llm_response(content="ok", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "send a mail"}], + tools=[{"type": "function", "function": {"name": "send_mail"}}], + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + user_permissions={"permissions": ["mail:read"], "denied": [], "is_system_admin": False}, + ) + + assert result.status == "completed" + handler.assert_not_awaited() + assert "Permission 'mail:write' required" in (result.steps[0].observation or "") + + +@pytest.mark.asyncio +async def test_f01_offered_tool_with_permission_executes(agent_def, tenant_id, user_id, mock_db): + """Offered tool + user HAS the required permission → executes normally.""" + handler = AsyncMock(return_value="mail sent") + registry = _make_guard_registry({ + "send_mail": _guard_tool("send_mail", handler, required_permission="mail:write"), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="sending", + tool_calls=[_make_tool_call(name="send_mail")], + cost_usd=0.001, + ), + _make_llm_response(content="done", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "send a mail"}], + tools=[{"type": "function", "function": {"name": "send_mail"}}], + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + user_permissions={"permissions": ["mail:write", "mail:read"], "denied": [], "is_system_admin": False}, + ) + + assert result.status == "completed" + handler.assert_awaited_once() + assert result.steps[0].observation == "mail sent" + + +@pytest.mark.asyncio +async def test_f01_fail_closed_without_permission_context(agent_def, tenant_id, user_id, mock_db): + """Tool with required_permission but NO permission context → rejected + (fail-closed: missing context is not implicit access).""" + handler = AsyncMock(return_value="boom") + registry = _make_guard_registry({ + "dangerous": _guard_tool("dangerous", handler, required_permission="system:admin"), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="calling", + tool_calls=[_make_tool_call(name="dangerous")], + cost_usd=0.001, + ), + _make_llm_response(content="ok", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "call it"}], + tools=[{"type": "function", "function": {"name": "dangerous"}}], + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + user_permissions=None, # no permission context at all + ) + + assert result.status == "completed" + handler.assert_not_awaited() + assert "Permission 'system:admin' required" in (result.steps[0].observation or "") + + +@pytest.mark.asyncio +async def test_f01_deny_list_session_shape_respected(agent_def, tenant_id, user_id, mock_db): + """Session user contexts carry ``denied_permissions`` (not ``denied``). + An explicit deny must reject even when the permission is also granted.""" + handler = AsyncMock(return_value="mail sent") + registry = _make_guard_registry({ + "send_mail": _guard_tool("send_mail", handler, required_permission="mail:write"), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="sending", + tool_calls=[_make_tool_call(name="send_mail")], + cost_usd=0.001, + ), + _make_llm_response(content="ok", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "send"}], + tools=[{"type": "function", "function": {"name": "send_mail"}}], + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + # session shape: granted AND denied via denied_permissions key + user_permissions={ + "permissions": ["mail:write"], + "denied_permissions": ["mail:write"], + "is_system_admin": False, + }, + ) + + assert result.status == "completed" + handler.assert_not_awaited() + assert "Permission 'mail:write' required" in (result.steps[0].observation or "") + + +@pytest.mark.asyncio +async def test_f01_permission_revoked_during_run_takes_effect_next_action(agent_def, tenant_id, user_id, mock_db): + """F01 acceptance: revoking permissions DURING a run takes effect on + the next action — the guard reads the CURRENT context every call.""" + handler = AsyncMock(return_value="mail sent") + registry = _make_guard_registry({ + "send_mail": _guard_tool("send_mail", handler, required_permission="mail:write"), + }) + + perms: dict[str, Any] = { + "permissions": ["mail:write"], + "denied": [], + "is_system_admin": False, + } + + call_count = 0 + + async def llm_with_revocation(**kwargs: Any) -> dict[str, Any]: + nonlocal call_count + call_count += 1 + if call_count == 1: + return _make_llm_response( + content="first send", + tool_calls=[_make_tool_call(name="send_mail")], + cost_usd=0.001, + ) + # Revoke the permission before the second LLM response + if call_count == 2: + perms["permissions"] = [] # permission revoked mid-run + return _make_llm_response( + content="second send", + tool_calls=[_make_tool_call(name="send_mail")], + cost_usd=0.001, + ) + return _make_llm_response(content="finished", cost_usd=0.001) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = llm_with_revocation + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "send twice"}], + tools=[{"type": "function", "function": {"name": "send_mail"}}], + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + user_permissions=perms, + ) + + assert result.status == "completed" + handler.assert_awaited_once() # first call ran, second was rejected + assert result.steps[0].observation == "mail sent" + assert "Permission 'mail:write' required" in (result.steps[1].observation or "") + + +@pytest.mark.asyncio +async def test_f01_dry_run_still_guards_non_offered_tools(agent_def, tenant_id, user_id, mock_db): + """Dry-run must not even SIMULATE a non-offered tool — the guard runs + before the dry-run path.""" + handler = AsyncMock(return_value="x") + registry = _make_guard_registry({ + "secret_tool": _guard_tool("secret_tool", handler), + }) + + with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm: + mock_llm.side_effect = [ + _make_llm_response( + content="trying", + tool_calls=[_make_tool_call(name="secret_tool")], + cost_usd=0.001, + ), + _make_llm_response(content="done", cost_usd=0.001), + ] + + result = await run_react_loop( + agent_definition=agent_def, + messages=[{"role": "user", "content": "try it"}], + tools=[], # nothing offered + tool_registry=registry, + db=mock_db, + tenant_id=tenant_id, + user_id=user_id, + dry_run=True, + ) + + assert result.status == "completed" + obs = result.steps[0].observation or "" + assert "not available" in obs + assert "would_execute" not in obs # guard fired, not the dry-run simulation + + @pytest.mark.asyncio async def test_multiple_tool_calls_per_step(agent_def, tenant_id, user_id, mock_db): """Multiple tool calls in a single LLM response are all executed."""