diff --git a/engine/src/agent_control_engine/selectors.py b/engine/src/agent_control_engine/selectors.py index 92ee0e15..1fda2022 100644 --- a/engine/src/agent_control_engine/selectors.py +++ b/engine/src/agent_control_engine/selectors.py @@ -17,6 +17,8 @@ def select_data(step: Step, path: str) -> Any: """ if not path or path == "*": return step.model_dump(mode="json") + if path == "canonical_name": + return step.canonical_name or step.name parts = path.split(".") current: Any = step diff --git a/engine/tests/test_selectors.py b/engine/tests/test_selectors.py index 577922ff..59dedde6 100644 --- a/engine/tests/test_selectors.py +++ b/engine/tests/test_selectors.py @@ -31,6 +31,7 @@ def llm_step_payload() -> Step: "path,expected", [ ("name", "search_database"), + ("canonical_name", "search_database"), ("input.query", "SELECT * FROM users"), ("input.limit", 10), ("input.nested.key", "value"), @@ -85,6 +86,17 @@ def test_select_data_none_handling(): assert result is None +def test_select_data_prefers_explicit_canonical_name() -> None: + payload = Step( + type="tool", + name="writer.web_search", + canonical_name="web_search", + input={}, + ) + + assert select_data(payload, "canonical_name") == "web_search" + + def test_list_selection(): """Test that selecting a path pointing to a list returns the whole list.""" # Given: a payload with a list in the output diff --git a/models/src/agent_control_models/agent.py b/models/src/agent_control_models/agent.py index 6a0eedba..540b050e 100644 --- a/models/src/agent_control_models/agent.py +++ b/models/src/agent_control_models/agent.py @@ -150,6 +150,15 @@ class Step(BaseModel): name: str = Field( ..., min_length=1, description="Step name (tool name or model/chain id)" ) + canonical_name: str | None = Field( + default=None, + min_length=1, + exclude_if=lambda value: value is None, + description=( + "Optional integration-independent identity for a qualified step name " + "(for example, 'web_search' for 'writer.web_search')." + ), + ) input: JSONValue = Field( ..., description="Input content for this step" ) diff --git a/models/src/agent_control_models/controls.py b/models/src/agent_control_models/controls.py index 1e2bb9e9..3a4729d9 100644 --- a/models/src/agent_control_models/controls.py +++ b/models/src/agent_control_models/controls.py @@ -27,7 +27,8 @@ class ControlSelector(BaseModel): default="*", description=( "Path to data using dot notation. " - "Examples: 'input', 'output', 'context.user_id', 'name', 'type', '*'" + "Examples: 'input', 'output', 'context.user_id', 'name', " + "'canonical_name', 'type', '*'" ), ) @@ -43,7 +44,15 @@ def validate_path(cls, v: str | None) -> str: ) # Valid root fields - valid_roots = {"input", "output", "name", "type", "context", "*"} + valid_roots = { + "input", + "output", + "name", + "canonical_name", + "type", + "context", + "*", + } root = v.split(".")[0] if root not in valid_roots: @@ -61,6 +70,7 @@ def validate_path(cls, v: str | None) -> str: {"path": "input"}, {"path": "*"}, {"path": "name"}, + {"path": "canonical_name"}, {"path": "output"}, ] } diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 767a3e02..bfd2348f 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -517,6 +517,7 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult: async def evaluate_controls( step_name: str, *, + canonical_step_name: str | None = None, input: Any | None = None, output: Any | None = None, context: dict[str, Any] | None = None, @@ -547,6 +548,7 @@ async def evaluate_controls( step_dict: dict[str, Any] = { "type": step_type, "name": step_name, + "canonical_name": canonical_step_name, "input": input if input is not None else default_value, "output": output if output is not None else default_value, } diff --git a/sdks/python/src/agent_control/integrations/_core.py b/sdks/python/src/agent_control/integrations/_core.py index 27693dbf..318ca7fd 100644 --- a/sdks/python/src/agent_control/integrations/_core.py +++ b/sdks/python/src/agent_control/integrations/_core.py @@ -48,6 +48,7 @@ async def _evaluate_and_enforce( agent_name: str, step_name: str, *, + canonical_step_name: str | None = None, input: Any | None = None, output: Any | None = None, context: dict[str, Any] | None = None, @@ -58,6 +59,7 @@ async def _evaluate_and_enforce( result = await agent_control.evaluate_controls( step_name=step_name, + canonical_step_name=canonical_step_name, input=input, output=output, context=context, diff --git a/sdks/python/src/agent_control/integrations/google_adk/plugin.py b/sdks/python/src/agent_control/integrations/google_adk/plugin.py index 28e59698..870a08ce 100644 --- a/sdks/python/src/agent_control/integrations/google_adk/plugin.py +++ b/sdks/python/src/agent_control/integrations/google_adk/plugin.py @@ -268,6 +268,7 @@ async def before_tool_callback( return None step_name = self._resolve_tool_step_name(tool, tool_context=tool_context) + canonical_step_name = resolve_tool_name(tool) self._ensure_step_known(self._build_tool_step_schema(tool, step_name)) context = self._safe_context( step_type="tool", @@ -281,6 +282,7 @@ async def before_tool_callback( await _evaluate_and_enforce( self.agent_name, step_name, + canonical_step_name=canonical_step_name, input=tool_args, context=context, step_type="tool", @@ -311,6 +313,7 @@ async def after_tool_callback( return None step_name = self._resolve_tool_step_name(tool, tool_context=tool_context) + canonical_step_name = resolve_tool_name(tool) self._ensure_step_known(self._build_tool_step_schema(tool, step_name)) context = self._safe_context( step_type="tool", @@ -325,6 +328,7 @@ async def after_tool_callback( await _evaluate_and_enforce( self.agent_name, step_name, + canonical_step_name=canonical_step_name, input=tool_args, output=result, context=context, diff --git a/sdks/python/src/agent_control/integrations/strands/plugin.py b/sdks/python/src/agent_control/integrations/strands/plugin.py index 1aa503cd..9e9c51eb 100644 --- a/sdks/python/src/agent_control/integrations/strands/plugin.py +++ b/sdks/python/src/agent_control/integrations/strands/plugin.py @@ -115,6 +115,7 @@ async def _evaluate_and_enforce( ) -> None: result = await agent_control.evaluate_controls( step_name=step_name, + canonical_step_name=step_name if step_type == "tool" else None, input=input, output=output, context=context, diff --git a/sdks/python/tests/test_evaluation.py b/sdks/python/tests/test_evaluation.py index 2fb92555..cb3d24f1 100644 --- a/sdks/python/tests/test_evaluation.py +++ b/sdks/python/tests/test_evaluation.py @@ -57,9 +57,9 @@ def json(self) -> dict[str, object]: json={ "agent_name": "agent-example_01", "step": { - "type": "llm", - "name": "chat", - "input": "hello", + "type": "llm", + "name": "chat", + "input": "hello", "output": None, "context": None, }, diff --git a/sdks/python/tests/test_google_adk_plugin.py b/sdks/python/tests/test_google_adk_plugin.py index f68bd341..1a056f57 100644 --- a/sdks/python/tests/test_google_adk_plugin.py +++ b/sdks/python/tests/test_google_adk_plugin.py @@ -10,7 +10,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from agent_control import ControlSteerError, ControlViolationError from agent_control._state import state @@ -369,6 +368,7 @@ async def test_tool_callbacks_scope_step_name_by_agent(plugin_module): ) assert mock_eval.await_args.args[1] == "writer.get_weather" + assert mock_eval.await_args.kwargs["canonical_step_name"] == "get_weather" @pytest.mark.asyncio diff --git a/sdks/typescript/src/generated/models/control-selector.ts b/sdks/typescript/src/generated/models/control-selector.ts index 8144bb20..1e8283c9 100644 --- a/sdks/typescript/src/generated/models/control-selector.ts +++ b/sdks/typescript/src/generated/models/control-selector.ts @@ -18,7 +18,7 @@ import { SDKValidationError } from "./errors/sdk-validation-error.js"; */ export type ControlSelector = { /** - * Path to data using dot notation. Examples: 'input', 'output', 'context.user_id', 'name', 'type', '*' + * Path to data using dot notation. Examples: 'input', 'output', 'context.user_id', 'name', 'canonical_name', 'type', '*' */ path?: string | null | undefined; }; diff --git a/sdks/typescript/src/generated/models/step.ts b/sdks/typescript/src/generated/models/step.ts index 132cf9c9..db7d5746 100644 --- a/sdks/typescript/src/generated/models/step.ts +++ b/sdks/typescript/src/generated/models/step.ts @@ -3,11 +3,16 @@ */ import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; /** * Runtime payload for an agent step invocation. */ export type Step = { + /** + * Optional integration-independent identity for a qualified step name (for example, 'web_search' for 'writer.web_search'). + */ + canonicalName?: string | null | undefined; /** * Optional context (conversation history, metadata, etc.) */ @@ -32,6 +37,7 @@ export type Step = { /** @internal */ export type Step$Outbound = { + canonical_name?: string | null | undefined; context?: { [k: string]: any } | null | undefined; input: any; name: string; @@ -40,14 +46,20 @@ export type Step$Outbound = { }; /** @internal */ -export const Step$outboundSchema: z.ZodMiniType = z.object( - { +export const Step$outboundSchema: z.ZodMiniType = z.pipe( + z.object({ + canonicalName: z.optional(z.nullable(z.string())), context: z.optional(z.nullable(z.record(z.string(), z.any()))), input: z.any(), name: z.string(), output: z.optional(z.nullable(z.any())), type: z.string(), - }, + }), + z.transform((v) => { + return remap$(v, { + canonicalName: "canonical_name", + }); + }), ); export function stepToJSON(step: Step): string { diff --git a/server/src/agent_control_server/bootstrap/out_of_box_controls.py b/server/src/agent_control_server/bootstrap/out_of_box_controls.py index 6c523b65..1e28b2a1 100644 --- a/server/src/agent_control_server/bootstrap/out_of_box_controls.py +++ b/server/src/agent_control_server/bootstrap/out_of_box_controls.py @@ -1,9 +1,5 @@ """Startup bootstrap for out-of-box controls. -Phase 1 provides the tooling needed to seed controls safely, but does not -register the static out-of-box control catalog yet. Phase 2 should add those -definitions to ``OUT_OF_BOX_CONTROL_TEMPLATES``. - Namespace rule: - Standalone Agent Control seeds into ``DEFAULT_NAMESPACE_KEY``. - Galileo-integrated Agent Control should call the same helper with @@ -36,6 +32,7 @@ _CONTROL_SEED_UNIQUE_CONSTRAINT = "idx_controls_namespace_seed_source" _INITIAL_VERSION_NOTE = "Out-of-box control seed" _SLUG_NAME_ADAPTER = TypeAdapter(SlugName) +_OUT_OF_BOX_TAGS = ["out-of-box"] @dataclass(frozen=True, slots=True) @@ -109,7 +106,209 @@ def skipped_count(self) -> int: ) -OUT_OF_BOX_CONTROL_TEMPLATES: tuple[OutOfBoxControlTemplate, ...] = () +def _leaf_control_payload( + *, + description: str, + selector_path: str, + evaluator_name: str, + evaluator_config: Mapping[str, object], + step_types: list[str], + stages: list[str], + decision: str, + tags: list[str], + steering_message: str | None = None, +) -> dict[str, object]: + action: dict[str, object] = {"decision": decision} + if steering_message is not None: + action["steering_context"] = {"message": steering_message} + + return { + "description": description, + "enabled": True, + "execution": "server", + "scope": {"step_types": step_types, "stages": stages}, + "condition": { + "selector": {"path": selector_path}, + "evaluator": { + "name": evaluator_name, + "config": dict(evaluator_config), + }, + }, + "action": action, + "tags": [*_OUT_OF_BOX_TAGS, *tags], + } + + +OUT_OF_BOX_CONTROL_TEMPLATES: tuple[OutOfBoxControlTemplate, ...] = ( + OutOfBoxControlTemplate.from_payload( + source_id="oob-ssn-match", + name="oob-ssn-match", + data=_leaf_control_payload( + description="Block LLM output containing US Social Security Numbers.", + selector_path="output", + evaluator_name="regex", + evaluator_config={"pattern": r"\b\d{3}-\d{2}-\d{4}\b"}, + step_types=["llm"], + stages=["post"], + decision="deny", + tags=["pii", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-credit-card-number-match", + name="oob-credit-card-number-match", + data=_leaf_control_payload( + description="Block LLM output containing common credit-card-like numbers.", + selector_path="output", + evaluator_name="regex", + evaluator_config={"pattern": r"\b(?:\d[ -]?){13,19}\b"}, + step_types=["llm"], + stages=["post"], + decision="deny", + tags=["pii", "payment", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-phone-number-match", + name="oob-phone-number-match", + data=_leaf_control_payload( + description="Block LLM output containing common US phone number formats.", + selector_path="output", + evaluator_name="regex", + evaluator_config={ + "pattern": ( + r"\b(?:\+?1[-.\s]?)?(?:\(?[2-9]\d{2}\)?[-.\s]?)?" + r"[2-9]\d{2}[-.\s]?\d{4}\b" + ) + }, + step_types=["llm"], + stages=["post"], + decision="deny", + tags=["pii", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-dangerous-shell-command-match", + name="oob-dangerous-shell-command-match", + data=_leaf_control_payload( + description="Block tool commands matching common destructive shell operations.", + selector_path="input.command", + evaluator_name="regex", + evaluator_config={ + "pattern": ( + r"(?:\brm\s+-rf\s+(?:/|~|\$HOME)(?:\s|[|;&]|$)|" + r"\bsudo\s+rm\s+-rf(?:\s|[|;&]|$)|" + r"\bmkfs(?:\.[a-z0-9]+)?(?:\s|[|;&]|$)|" + r"\bdd\s+if=[^\s]+\s+of=/dev/[^\s]+(?:\s|[|;&]|$)|" + r"\bchmod\s+-R\s+777\s+/(?:\s|[|;&]|$)|" + r"\bchown\s+-R\s+[^|;&]*\s+/(?:\s|[|;&]|$)|" + r"\bshutdown\s+(?:-h\s+)?now(?:\s|[|;&]|$)|" + r"\breboot(?:\s|[|;&]|$))" + ), + "flags": ["IGNORECASE"], + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["tool", "shell", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-high-value-action-requires-approval", + name="oob-high-value-action-requires-approval", + data=_leaf_control_payload( + description=( + "Steer tool calls over the default amount threshold to collect approval." + ), + selector_path="input", + evaluator_name="json", + evaluator_config={ + "json_schema": { + "type": "object", + "anyOf": [ + {"not": {"required": ["amount"]}}, + { + "required": ["amount"], + "properties": { + "amount": {"type": "number", "maximum": 10000} + }, + }, + ], + } + }, + step_types=["tool"], + stages=["pre"], + decision="steer", + steering_message=( + "Pause this high-value action and submit its exact parameters to a trusted " + "host approval workflow. The host must bind any approval artifact to this " + "specific action; approval fields supplied in tool input are not evidence." + ), + tags=["tool", "approval", "json"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-outbound-communication-requires-approval", + name="oob-outbound-communication-requires-approval", + data=_leaf_control_payload( + description=( + "Steer outbound communication tool calls to collect approval before sending." + ), + selector_path="input", + evaluator_name="json", + evaluator_config={ + "json_schema": { + "type": "object", + "anyOf": [ + { + "not": { + "anyOf": [ + {"required": ["to"]}, + {"required": ["recipient"]}, + {"required": ["recipients"]}, + {"required": ["email"]}, + {"required": ["phone_number"]}, + {"required": ["channel"]}, + {"required": ["destination"]}, + ] + } + } + ], + } + }, + step_types=["tool"], + stages=["pre"], + decision="steer", + steering_message=( + "Pause this outbound communication and submit its exact recipients and " + "content to a trusted host approval workflow. The host must bind any " + "approval artifact to this specific action; approval fields supplied in " + "tool input are not evidence." + ), + tags=["tool", "approval", "exfiltration", "json"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-only-approved-tools-may-run", + name="oob-only-approved-tools-may-run", + data=_leaf_control_payload( + description="Deny tool calls whose step name is not in the approved tool list.", + selector_path="canonical_name", + evaluator_name="list", + evaluator_config={ + "values": ["search", "web_search", "retrieve", "calculator"], + "logic": "any", + "match_on": "no_match", + "match_mode": "exact", + "case_sensitive": False, + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["tool", "allowlist", "list"], + ), + ), +) def default_out_of_box_namespace_key() -> str: diff --git a/server/src/agent_control_server/config.py b/server/src/agent_control_server/config.py index 00335611..5e342222 100644 --- a/server/src/agent_control_server/config.py +++ b/server/src/agent_control_server/config.py @@ -214,7 +214,6 @@ class Settings(BaseSettings): "AGENT_CONTROL_ALLOW_HEADERS", "ALLOW_HEADERS", ) - def get_cors_origins(self) -> list[str]: """Parse CORS origins from string or list.""" return self._parse_list_setting(self.cors_origins) diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index d328c7f9..a2ff9f1a 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -43,7 +43,8 @@ from sqlalchemy.ext.asyncio import AsyncSession from ..auth_framework import Operation, Principal, get_authorizer, require_operation -from ..db import get_async_db +from ..bootstrap.out_of_box_controls import seed_out_of_box_controls +from ..db import AsyncSessionLocal, get_async_db from ..errors import ( APIError, APIValidationError, @@ -257,6 +258,56 @@ def _validate_attachment_filters( ) +async def _seed_out_of_box_controls_for_namespace( + *, + namespace_key: str, +) -> None: + """Best-effort idempotent namespace seeding for browse/list surfaces.""" + try: + await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=namespace_key, + available_evaluators=set(list_evaluators().keys()), + ) + except Exception: + _logger.warning( + "Out-of-box control seed failed for namespace '%s'; continuing request", + namespace_key, + exc_info=True, + ) + + +def _should_seed_out_of_box_controls_on_list( + *, + cursor: int | None, + name: str | None, + enabled: bool | None, + template_backed: bool | None, + cloned: bool | None, + step_type: str | None, + stage: str | None, + execution: str | None, + tag: str | None, + include_attachments: bool, + attachment_target_type: str | None, + attachment_target_id: str | None, +) -> bool: + return ( + cursor is None + and name is None + and enabled is None + and template_backed is None + and cloned is not True + and step_type is None + and stage is None + and execution is None + and tag is None + and not include_attachments + and attachment_target_type is None + and attachment_target_id is None + ) + + def _serialize_control_data( control_data: ControlDefinition | UnrenderedTemplateControl, ) -> dict[str, object]: @@ -1234,6 +1285,21 @@ async def list_controls( control_service = ControlService(db) namespace_key = principal.namespace_key + if _should_seed_out_of_box_controls_on_list( + cursor=cursor, + name=name, + enabled=enabled, + template_backed=template_backed, + cloned=cloned, + step_type=step_type, + stage=stage, + execution=execution, + tag=tag, + include_attachments=include_attachments, + attachment_target_type=attachment_target_type, + attachment_target_id=attachment_target_id, + ): + await _seed_out_of_box_controls_for_namespace(namespace_key=namespace_key) filter_by_attachment = target_principal is not None and ( attachment_target_type is not None or attachment_target_id is not None ) diff --git a/server/tests/test_controls_additional.py b/server/tests/test_controls_additional.py index cf7aa4b0..59884f55 100644 --- a/server/tests/test_controls_additional.py +++ b/server/tests/test_controls_additional.py @@ -12,6 +12,13 @@ from agent_control_evaluators import RegexEvaluatorConfig from agent_control_models import ConditionNode from agent_control_models.errors import ErrorCode, ErrorReason +from fastapi.testclient import TestClient +from sqlalchemy import select, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import Session +from starlette.requests import Request + from agent_control_server.auth_framework import Operation, Principal, set_authorizer from agent_control_server.db import get_async_db from agent_control_server.endpoints import controls as controls_module @@ -23,11 +30,6 @@ ControlBinding, ControlVersion, ) -from fastapi.testclient import TestClient -from sqlalchemy import select, text -from sqlalchemy.exc import IntegrityError -from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import Session from .conftest import engine from .utils import VALID_CONTROL_PAYLOAD @@ -40,6 +42,22 @@ def _make_integrity_error(constraint_name: str) -> IntegrityError: return IntegrityError("statement", {}, orig) +def _request(*, query: str = "", body: bytes = b"") -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "GET", + "path": "/", + "headers": [], + "query_string": query.encode(), + }, + receive, + ) + + def _create_control( client: TestClient, name: str | None = None, @@ -473,6 +491,75 @@ def test_clone_and_bind_context_tolerates_invalid_body_shapes( assert bad_target_resp.status_code == 422 +@pytest.mark.asyncio +async def test_clone_and_bind_context_returns_empty_for_malformed_json() -> None: + malformed_request = _request(body=b"{") + invalid_target_request = _request( + body=json.dumps( + { + "target_binding": { + "target_type": "log_stream", + "target_id": "", + } + } + ).encode() + ) + + assert await controls_module._clone_and_bind_context(malformed_request) == {} + assert await controls_module._clone_and_bind_context(invalid_target_request) == {} + + +def test_attachment_target_context_rejects_invalid_values() -> None: + invalid_type = _request(query="attachment_target_type=&attachment_target_id=target") + invalid_id = _request(query="attachment_target_type=log_stream&attachment_target_id=") + + assert controls_module._attachment_target_context(invalid_type) == {} + assert controls_module._attachment_target_context(invalid_id) == {} + + +@pytest.mark.asyncio +async def test_optional_attachment_authorization_skips_false_flag() -> None: + request = _request(query="include_attachments=false") + + assert await controls_module._optional_attachment_target_principal(request) is None + + +def test_enabled_from_stored_payload_defaults_for_non_mapping() -> None: + assert controls_module._enabled_from_stored_payload("invalid") is True + + +@pytest.mark.asyncio +async def test_seed_out_of_box_controls_failure_is_best_effort( + monkeypatch: pytest.MonkeyPatch, +) -> None: + seed = AsyncMock(side_effect=RuntimeError("database unavailable")) + monkeypatch.setattr(controls_module, "seed_out_of_box_controls", seed) + + await controls_module._seed_out_of_box_controls_for_namespace( + namespace_key="seed-failure-namespace" + ) + + seed.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_resolve_clone_name_reports_generated_name_exhaustion() -> None: + control_service = MagicMock() + control_service.active_control_name_exists = AsyncMock(return_value=True) + + with pytest.raises(APIError) as exc_info: + await controls_module._resolve_clone_name( + control_service, + namespace_key=DEFAULT_NAMESPACE_KEY, + source_id=1, + source_name="source-control", + requested_name=None, + ) + + assert exc_info.value.error_code == ErrorCode.CONTROL_NAME_CONFLICT + assert control_service.active_control_name_exists.await_count == 5 + + def test_clone_and_bind_context_drops_invalid_target_fields( client: TestClient, ) -> None: diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 59135451..2975dbd2 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -6,7 +6,15 @@ from typing import cast import pytest +from agent_control_evaluators.json.config import JSONEvaluatorConfig +from agent_control_evaluators.json.evaluator import JSONEvaluator +from agent_control_evaluators.list.config import ListEvaluatorConfig +from agent_control_evaluators.list.evaluator import ListEvaluator +from agent_control_evaluators.regex.config import RegexEvaluatorConfig +from agent_control_evaluators.regex.evaluator import RegexEvaluator +from agent_control_models import EvaluatorSpec from agent_control_server.bootstrap.out_of_box_controls import ( + OUT_OF_BOX_CONTROL_TEMPLATES, OutOfBoxControlTemplate, default_out_of_box_namespace_key, missing_required_evaluators, @@ -27,6 +35,17 @@ from .conftest import AsyncSessionTest, engine +_EXPECTED_OOB_CONTROL_NAMES = ( + "oob-ssn-match", + "oob-credit-card-number-match", + "oob-phone-number-match", + "oob-dangerous-shell-command-match", + "oob-high-value-action-requires-approval", + "oob-outbound-communication-requires-approval", + "oob-only-approved-tools-may-run", +) +_AVAILABLE_PHASE_2_EVALUATORS = {"regex", "json", "list"} + def _control_payload(*, evaluator_name: str = "regex") -> dict[str, object]: return { @@ -75,10 +94,39 @@ def _count_table_rows(table: Table) -> int: return cast(int, session.scalar(select(func.count()).select_from(table))) +def _oob_evaluator_spec(name: str) -> EvaluatorSpec: + template = next(template for template in OUT_OF_BOX_CONTROL_TEMPLATES if template.name == name) + leaf = template.control.primary_leaf() + assert leaf is not None + leaf_parts = leaf.leaf_parts() + assert leaf_parts is not None + _, evaluator = leaf_parts + return evaluator + + def test_default_namespace_key_uses_standalone_namespace() -> None: assert default_out_of_box_namespace_key() == DEFAULT_NAMESPACE_KEY +def test_out_of_box_catalog_contains_phase_2_templates() -> None: + assert tuple(template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES) == ( + _EXPECTED_OOB_CONTROL_NAMES + ) + assert { + evaluator + for template in OUT_OF_BOX_CONTROL_TEMPLATES + for evaluator in template.required_evaluators + } == _AVAILABLE_PHASE_2_EVALUATORS + approved_tools = next( + template + for template in OUT_OF_BOX_CONTROL_TEMPLATES + if template.name == "oob-only-approved-tools-may-run" + ) + approved_tools_leaf = approved_tools.control.primary_leaf() + assert approved_tools_leaf is not None + assert approved_tools_leaf.selector.path == "canonical_name" + + def test_missing_required_evaluators_returns_sorted_names() -> None: missing = missing_required_evaluators( {"galileo.luna", "regex", "json"}, @@ -158,6 +206,48 @@ async def test_seed_creates_control_version_in_namespace_without_bindings() -> N assert _count_table_rows(ControlBinding.__table__) == 0 +@pytest.mark.asyncio +async def test_seed_default_catalog_creates_all_controls_without_bindings() -> None: + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators=_AVAILABLE_PHASE_2_EVALUATORS, + ) + + assert result.created == _EXPECTED_OOB_CONTROL_NAMES + assert result.skipped_existing == () + assert result.skipped_missing_evaluator == () + assert result.skipped_conflict == () + + controls = _fetch_controls() + assert tuple(control.name for control in controls) == _EXPECTED_OOB_CONTROL_NAMES + assert {control.namespace_key for control in controls} == {DEFAULT_NAMESPACE_KEY} + assert len(_fetch_versions()) == len(_EXPECTED_OOB_CONTROL_NAMES) + assert _count_table_rows(policy_controls) == 0 + assert _count_table_rows(agent_controls) == 0 + assert _count_table_rows(ControlBinding.__table__) == 0 + + +@pytest.mark.asyncio +async def test_seed_default_catalog_is_idempotent() -> None: + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators=_AVAILABLE_PHASE_2_EVALUATORS, + ) + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators=_AVAILABLE_PHASE_2_EVALUATORS, + ) + + assert result.created == () + assert result.skipped_existing == _EXPECTED_OOB_CONTROL_NAMES + assert len(_fetch_controls()) == len(_EXPECTED_OOB_CONTROL_NAMES) + assert len(_fetch_versions()) == len(_EXPECTED_OOB_CONTROL_NAMES) + + @pytest.mark.asyncio async def test_seed_is_idempotent_for_existing_active_control_names() -> None: template = _template(name="oob-idempotent-control") @@ -282,3 +372,80 @@ async def seed_source_exists( assert result.skipped_conflict == ("oob-race-control",) assert len(_fetch_controls()) == 1 assert len(_fetch_versions()) == 1 + + +@pytest.mark.asyncio +async def test_regex_out_of_box_controls_match_representative_payloads() -> None: + ssn_spec = _oob_evaluator_spec("oob-ssn-match") + ssn_evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(ssn_spec.config)) + ssn_result = await ssn_evaluator.evaluate("Customer SSN is 123-45-6789.") + assert ssn_result.matched is True + + shell_spec = _oob_evaluator_spec("oob-dangerous-shell-command-match") + shell_evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(shell_spec.config)) + for command in ( + "sudo rm -rf /", + "rm -rf /", + "rm -rf ~", + "chmod -R 777 /", + "chown -R root /", + ): + shell_result = await shell_evaluator.evaluate(command) + assert shell_result.matched is True, command + + +@pytest.mark.asyncio +async def test_json_out_of_box_controls_ignore_caller_controlled_approval_flags() -> None: + high_value_spec = _oob_evaluator_spec("oob-high-value-action-requires-approval") + high_value_evaluator = JSONEvaluator( + JSONEvaluatorConfig.model_validate(high_value_spec.config) + ) + + high_value_result = await high_value_evaluator.evaluate({"amount": 25000}) + low_value_result = await high_value_evaluator.evaluate({"amount": 250}) + caller_approved_results = [ + await high_value_evaluator.evaluate({"amount": 25000, "approved": True}), + await high_value_evaluator.evaluate( + {"amount": 25000, "approval": {"approved": True}} + ), + ] + + assert high_value_result.matched is True + assert low_value_result.matched is False + assert all(result.matched is True for result in caller_approved_results) + + outbound_spec = _oob_evaluator_spec("oob-outbound-communication-requires-approval") + outbound_evaluator = JSONEvaluator(JSONEvaluatorConfig.model_validate(outbound_spec.config)) + + outbound_result = await outbound_evaluator.evaluate( + {"to": "customer@example.com", "message": "Hello"} + ) + internal_result = await outbound_evaluator.evaluate({"query": "customer history"}) + caller_approved_outbound_results = [ + await outbound_evaluator.evaluate( + {"to": "customer@example.com", "message": "Hello", "approved": True} + ), + await outbound_evaluator.evaluate( + { + "to": "customer@example.com", + "message": "Hello", + "approval": {"approved": True}, + } + ), + ] + + assert outbound_result.matched is True + assert internal_result.matched is False + assert all(result.matched is True for result in caller_approved_outbound_results) + + +@pytest.mark.asyncio +async def test_list_out_of_box_control_matches_unapproved_tools() -> None: + tool_spec = _oob_evaluator_spec("oob-only-approved-tools-may-run") + tool_evaluator = ListEvaluator(ListEvaluatorConfig.model_validate(tool_spec.config)) + + delete_result = await tool_evaluator.evaluate("delete_user") + search_result = await tool_evaluator.evaluate("web_search") + + assert delete_result.matched is True + assert search_result.matched is False diff --git a/server/tests/test_principal_namespace_flow.py b/server/tests/test_principal_namespace_flow.py index 8f16a795..f02cf61b 100644 --- a/server/tests/test_principal_namespace_flow.py +++ b/server/tests/test_principal_namespace_flow.py @@ -6,14 +6,19 @@ from copy import deepcopy from typing import Any +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from sqlalchemy.orm import Session + from agent_control_server.auth_framework import ( Operation, Principal, set_authorizer, ) -from fastapi import FastAPI, Request -from fastapi.testclient import TestClient +from agent_control_server.bootstrap.out_of_box_controls import OUT_OF_BOX_CONTROL_TEMPLATES +from agent_control_server.models import Control +from .conftest import engine from .utils import VALID_CONTROL_PAYLOAD @@ -39,6 +44,24 @@ async def authorize( ) +class ControlsReadOnlyAuthorizer(HeaderNamespaceAuthorizer): + """Allow the controls list read and record every authorization operation.""" + + def __init__(self) -> None: + self.operations: list[Operation] = [] + + async def authorize( + self, + request: Request, + operation: Operation, + context: dict[str, Any] | None = None, + ) -> Principal: + self.operations.append(operation) + if operation is not Operation.CONTROLS_READ: + raise AssertionError(f"Unexpected authorization operation: {operation}") + return await super().authorize(request, operation, context) + + def _client(app: FastAPI, namespace_key: str) -> TestClient: return TestClient( app, @@ -73,6 +96,100 @@ def _evaluation_payload(agent_name: str) -> dict[str, Any]: } +def test_controls_list_seeds_out_of_box_controls_for_principal_namespace( + app: FastAPI, +) -> None: + authorizer = ControlsReadOnlyAuthorizer() + set_authorizer(authorizer) + + namespace_client = _client(app, "org-oob-controls") + filtered = namespace_client.get("/api/v1/controls", params={"name": "oob"}) + assert filtered.status_code == 200, filtered.text + assert filtered.json()["controls"] == [] + + resp = namespace_client.get("/api/v1/controls", params={"limit": 10}) + assert resp.status_code == 200, resp.text + + expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} + returned_names = {control["name"] for control in resp.json()["controls"]} + assert expected_names.issubset(returned_names) + assert authorizer.operations == [ + Operation.CONTROLS_READ, + Operation.CONTROLS_READ, + ] + + +def test_controls_list_seeds_out_of_box_controls_alongside_custom_control( + app: FastAPI, +) -> None: + set_authorizer(HeaderNamespaceAuthorizer()) + namespace_client = _client(app, "org-with-custom-control") + custom_name = f"custom-{uuid.uuid4().hex[:12]}" + + created = namespace_client.put( + "/api/v1/controls", + json={"name": custom_name, "data": VALID_CONTROL_PAYLOAD}, + ) + assert created.status_code == 200, created.text + + response = namespace_client.get( + "/api/v1/controls", + params={"limit": 20, "cloned": "false"}, + ) + assert response.status_code == 200, response.text + + returned_names = {control["name"] for control in response.json()["controls"]} + expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} + assert returned_names == {*expected_names, custom_name} + + +def test_controls_list_completes_partially_seeded_namespace(app: FastAPI) -> None: + set_authorizer(HeaderNamespaceAuthorizer()) + namespace_key = "org-partially-seeded" + first_template = OUT_OF_BOX_CONTROL_TEMPLATES[0] + with Session(engine) as session: + session.add( + Control( + namespace_key=namespace_key, + name=first_template.name, + data=first_template.control.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + exclude_unset=True, + ), + seed_source_id=first_template.source_id, + ) + ) + session.commit() + + response = _client(app, namespace_key).get( + "/api/v1/controls", + params={"limit": 20, "cloned": "false"}, + ) + assert response.status_code == 200, response.text + + returned_names = {control["name"] for control in response.json()["controls"]} + expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} + assert returned_names == expected_names + + +def test_controls_list_retries_out_of_box_seeding_for_default_namespace( + app: FastAPI, +) -> None: + set_authorizer(HeaderNamespaceAuthorizer()) + + response = _client(app, "default").get( + "/api/v1/controls", + params={"limit": 20, "cloned": "false"}, + ) + assert response.status_code == 200, response.text + + returned_names = {control["name"] for control in response.json()["controls"]} + expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} + assert returned_names == expected_names + + def test_principal_namespace_scopes_management_and_runtime(app: FastAPI) -> None: set_authorizer(HeaderNamespaceAuthorizer())