From c99a3850c39bdb9a564df45cab00f580ae3b0090 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 12:18:51 -0700 Subject: [PATCH 01/18] add out of hte box controls --- .../bootstrap/__init__.py | 2 + .../bootstrap/out_of_box_controls.py | 235 ++++++++++++++++++ server/src/agent_control_server/main.py | 22 ++ server/tests/test_main_lifespan.py | 16 ++ .../test_out_of_box_controls_bootstrap.py | 210 ++++++++++++++++ 5 files changed, 485 insertions(+) create mode 100644 server/src/agent_control_server/bootstrap/__init__.py create mode 100644 server/src/agent_control_server/bootstrap/out_of_box_controls.py create mode 100644 server/tests/test_out_of_box_controls_bootstrap.py diff --git a/server/src/agent_control_server/bootstrap/__init__.py b/server/src/agent_control_server/bootstrap/__init__.py new file mode 100644 index 00000000..9f37e7a4 --- /dev/null +++ b/server/src/agent_control_server/bootstrap/__init__.py @@ -0,0 +1,2 @@ +"""Startup bootstrap helpers for server-managed defaults.""" + 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 new file mode 100644 index 00000000..3ba30b1d --- /dev/null +++ b/server/src/agent_control_server/bootstrap/out_of_box_controls.py @@ -0,0 +1,235 @@ +"""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 + ``namespace_key`` set to the Galileo ``organization_id`` carried by the + upstream auth bridge. +""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Self, cast + +from agent_control_models import ControlDefinition +from agent_control_models.server import SlugName +from pydantic import TypeAdapter +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from ..models import DEFAULT_NAMESPACE_KEY +from ..services.controls import ControlService + +_CONTROL_NAME_UNIQUE_CONSTRAINTS = frozenset( + { + "controls_name_key", + "idx_controls_name_active", + "idx_controls_namespace_name_active", + } +) +_INITIAL_VERSION_NOTE = "Out-of-box control seed" +_SLUG_NAME_ADAPTER = TypeAdapter(SlugName) + + +@dataclass(frozen=True, slots=True) +class OutOfBoxControlTemplate: + """Validated control definition plus the evaluator names it needs.""" + + name: str + control: ControlDefinition + required_evaluators: frozenset[str] = field(default_factory=frozenset) + + def __post_init__(self) -> None: + object.__setattr__(self, "name", _SLUG_NAME_ADAPTER.validate_python(self.name)) + if not self.required_evaluators: + required_evaluators = { + evaluator.name for _, evaluator in self.control.iter_condition_leaf_parts() + } + object.__setattr__(self, "required_evaluators", frozenset(required_evaluators)) + return + + object.__setattr__(self, "required_evaluators", frozenset(self.required_evaluators)) + + @classmethod + def from_payload( + cls, + *, + name: str, + data: Mapping[str, object], + required_evaluators: Collection[str] = frozenset(), + ) -> Self: + """Build a template from raw JSON-like data and validate it immediately.""" + return cls( + name=name, + control=ControlDefinition.model_validate(data), + required_evaluators=frozenset(required_evaluators), + ) + + +@dataclass(frozen=True, slots=True) +class SkippedOutOfBoxControl: + """A control skipped because the current pod cannot evaluate it.""" + + name: str + missing_evaluators: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class OutOfBoxSeedResult: + """Summary of one bootstrap seed pass.""" + + created: tuple[str, ...] = () + skipped_existing: tuple[str, ...] = () + skipped_missing_evaluator: tuple[SkippedOutOfBoxControl, ...] = () + skipped_conflict: tuple[str, ...] = () + + @property + def created_count(self) -> int: + """Number of controls inserted by this seed pass.""" + return len(self.created) + + @property + def skipped_count(self) -> int: + """Number of controls skipped by this seed pass.""" + return ( + len(self.skipped_existing) + + len(self.skipped_missing_evaluator) + + len(self.skipped_conflict) + ) + + +OUT_OF_BOX_CONTROL_TEMPLATES: tuple[OutOfBoxControlTemplate, ...] = () + + +def default_out_of_box_namespace_key() -> str: + """Return the standalone namespace used for server startup seeding.""" + return DEFAULT_NAMESPACE_KEY + + +def missing_required_evaluators( + required_evaluators: Collection[str], + available_evaluators: Collection[str], +) -> tuple[str, ...]: + """Return required evaluator names absent from the current pod.""" + missing = set(required_evaluators) - set(available_evaluators) + return tuple(sorted(missing)) + + +async def seed_out_of_box_controls( + *, + session_factory: async_sessionmaker[AsyncSession], + namespace_key: str, + available_evaluators: Collection[str], + templates: Sequence[OutOfBoxControlTemplate] = OUT_OF_BOX_CONTROL_TEMPLATES, +) -> OutOfBoxSeedResult: + """Create missing out-of-box controls in a namespace. + + Existing active controls are left untouched so customer edits survive + restarts and upgrades. Duplicate-name integrity errors are treated as + benign races with another pod and are reported as ``skipped_conflict``. + """ + if not templates: + return OutOfBoxSeedResult() + + created: list[str] = [] + skipped_existing: list[str] = [] + skipped_missing_evaluator: list[SkippedOutOfBoxControl] = [] + skipped_conflict: list[str] = [] + + available_evaluator_names = set(available_evaluators) + async with session_factory() as session: + for template in templates: + missing = missing_required_evaluators( + template.required_evaluators, + available_evaluator_names, + ) + if missing: + skipped_missing_evaluator.append( + SkippedOutOfBoxControl( + name=template.name, + missing_evaluators=missing, + ) + ) + continue + + outcome = await _seed_one_control( + session, + namespace_key=namespace_key, + template=template, + ) + if outcome == "created": + created.append(template.name) + elif outcome == "conflict": + skipped_conflict.append(template.name) + else: + skipped_existing.append(template.name) + + return OutOfBoxSeedResult( + created=tuple(created), + skipped_existing=tuple(skipped_existing), + skipped_missing_evaluator=tuple(skipped_missing_evaluator), + skipped_conflict=tuple(skipped_conflict), + ) + + +async def _seed_one_control( + session: AsyncSession, + *, + namespace_key: str, + template: OutOfBoxControlTemplate, +) -> str: + control_service = ControlService(session) + if await control_service.active_control_name_exists(template.name, namespace_key=namespace_key): + return "existing" + + control = control_service.create_control( + namespace_key=namespace_key, + name=template.name, + data=_serialize_control_data(template.control), + ) + try: + await control_service.create_version( + control, + event_type="created", + note=_INITIAL_VERSION_NOTE, + ) + await session.commit() + except IntegrityError as exc: + await session.rollback() + if _is_control_name_conflict(exc): + return "conflict" + raise + return "created" + + +def _serialize_control_data(control_data: ControlDefinition) -> dict[str, object]: + data_json = control_data.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + exclude_unset=True, + ) + if "scope" in data_json and isinstance(data_json["scope"], dict): + data_json["scope"] = { + key: value for key, value in data_json["scope"].items() if value is not None + } + if "enabled" not in data_json: + data_json["enabled"] = control_data.enabled + return cast(dict[str, object], data_json) + + +def _is_control_name_conflict(error: IntegrityError) -> bool: + diag = getattr(getattr(error.orig, "diag", None), "constraint_name", None) + if diag in _CONTROL_NAME_UNIQUE_CONSTRAINTS: + return True + + error_text = " ".join( + part for part in (str(error.orig), str(error)) if part and part != "None" + ) + return any(name in error_text for name in _CONTROL_NAME_UNIQUE_CONSTRAINTS) diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index 16152824..bed2bf9c 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -18,6 +18,10 @@ from . import __version__ as server_version from .auth import get_api_key_from_header +from .bootstrap.out_of_box_controls import ( + default_out_of_box_namespace_key, + seed_out_of_box_controls, +) from .config import observability_settings, settings from .db import AsyncSessionLocal, async_engine from .endpoints.agents import router as agent_router @@ -142,6 +146,24 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: available = list(list_evaluators().keys()) logger.info(f"Evaluator discovery complete. Available evaluators: {available}") + try: + seed_result = await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=default_out_of_box_namespace_key(), + available_evaluators=set(available), + ) + if seed_result.created_count or seed_result.skipped_count: + logger.info( + "Out-of-box control bootstrap complete: created=%s " + "skipped_existing=%s skipped_missing_evaluator=%s skipped_conflict=%s", + seed_result.created_count, + len(seed_result.skipped_existing), + len(seed_result.skipped_missing_evaluator), + len(seed_result.skipped_conflict), + ) + except Exception: + logger.warning("Out-of-box control bootstrap failed; continuing startup", exc_info=True) + # Initialize observability components (stored on app.state) if observability_settings.enabled: logger.info("Initializing observability components...") diff --git a/server/tests/test_main_lifespan.py b/server/tests/test_main_lifespan.py index 293fb957..1d464e68 100644 --- a/server/tests/test_main_lifespan.py +++ b/server/tests/test_main_lifespan.py @@ -217,6 +217,22 @@ def test_lifespan_skips_observability_when_disabled(monkeypatch) -> None: assert not hasattr(app.state, "event_ingestor") +def test_lifespan_fails_open_when_out_of_box_bootstrap_fails(monkeypatch, caplog) -> None: + async def fail_seed_out_of_box_controls(**kwargs: object) -> None: + raise RuntimeError("boom") + + monkeypatch.setattr(observability_settings, "enabled", False) + monkeypatch.setattr(main_module, "seed_out_of_box_controls", fail_seed_out_of_box_controls) + + app = FastAPI(lifespan=lifespan) + + with caplog.at_level("WARNING"): + with TestClient(app): + pass + + assert "Out-of-box control bootstrap failed; continuing startup" in caplog.text + + def test_custom_openapi_replaces_jsonvalue_variants(monkeypatch) -> None: # Given: a custom openapi generator that includes Pydantic JSONValue schemas json_value_schema_names = ( diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py new file mode 100644 index 00000000..5ba30577 --- /dev/null +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +import uuid +from copy import deepcopy +from typing import cast + +import pytest +from agent_control_server.bootstrap.out_of_box_controls import ( + OutOfBoxControlTemplate, + default_out_of_box_namespace_key, + missing_required_evaluators, + seed_out_of_box_controls, +) +from agent_control_server.models import ( + DEFAULT_NAMESPACE_KEY, + Control, + ControlBinding, + ControlVersion, + agent_controls, + policy_controls, +) +from agent_control_server.services.controls import ControlService +from pydantic import ValidationError +from sqlalchemy import Table, func, select +from sqlalchemy.orm import Session + +from .conftest import AsyncSessionTest, engine + + +def _control_payload(*, evaluator_name: str = "regex") -> dict[str, object]: + return { + "description": "Synthetic out-of-box control", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["post"]}, + "condition": { + "selector": {"path": "output"}, + "evaluator": { + "name": evaluator_name, + "config": {"pattern": r"\bsecret\b"}, + }, + }, + "action": {"decision": "deny"}, + "tags": ["out-of-box"], + } + + +def _template( + *, + name: str | None = None, + evaluator_name: str = "regex", +) -> OutOfBoxControlTemplate: + return OutOfBoxControlTemplate.from_payload( + name=name or f"oob-test-{uuid.uuid4().hex}", + data=_control_payload(evaluator_name=evaluator_name), + ) + + +def _fetch_controls() -> list[Control]: + with Session(engine) as session: + return list(session.scalars(select(Control).order_by(Control.id)).all()) + + +def _fetch_versions() -> list[ControlVersion]: + with Session(engine) as session: + return list(session.scalars(select(ControlVersion).order_by(ControlVersion.id)).all()) + + +def _count_table_rows(table: Table) -> int: + with Session(engine) as session: + return cast(int, session.scalar(select(func.count()).select_from(table))) + + +def test_default_namespace_key_uses_standalone_namespace() -> None: + assert default_out_of_box_namespace_key() == DEFAULT_NAMESPACE_KEY + + +def test_missing_required_evaluators_returns_sorted_names() -> None: + missing = missing_required_evaluators( + {"galileo.luna", "regex", "json"}, + {"json"}, + ) + + assert missing == ("galileo.luna", "regex") + + +def test_template_from_payload_validates_control_definition() -> None: + payload = deepcopy(_control_payload()) + payload["condition"] = { + "selector": {"path": "invalid_root.value"}, + "evaluator": {"name": "regex", "config": {"pattern": "x"}}, + } + + with pytest.raises(ValidationError): + OutOfBoxControlTemplate.from_payload(name="invalid-oob-control", data=payload) + + +@pytest.mark.asyncio +async def test_seed_skips_template_when_required_evaluator_is_missing() -> None: + template = _template(name="oob-missing-evaluator") + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"json"}, + templates=(template,), + ) + + assert result.created == () + assert result.skipped_existing == () + assert result.skipped_conflict == () + assert len(result.skipped_missing_evaluator) == 1 + assert result.skipped_missing_evaluator[0].name == "oob-missing-evaluator" + assert result.skipped_missing_evaluator[0].missing_evaluators == ("regex",) + assert _fetch_controls() == [] + + +@pytest.mark.asyncio +async def test_seed_creates_control_version_in_namespace_without_bindings() -> None: + template = _template(name="oob-create-control") + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key="galileo-org-123", + available_evaluators={"regex"}, + templates=(template,), + ) + + assert result.created == ("oob-create-control",) + controls = _fetch_controls() + assert len(controls) == 1 + control = controls[0] + assert control.namespace_key == "galileo-org-123" + assert control.name == "oob-create-control" + assert control.data["enabled"] is True + assert control.data["condition"]["evaluator"]["name"] == "regex" + + versions = _fetch_versions() + assert len(versions) == 1 + assert versions[0].control_id == control.id + assert versions[0].version_num == 1 + assert versions[0].event_type == "created" + assert versions[0].note == "Out-of-box control seed" + assert versions[0].snapshot["name"] == "oob-create-control" + + 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_is_idempotent_for_existing_active_control_names() -> None: + template = _template(name="oob-idempotent-control") + + first_result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + second_result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + assert first_result.created == ("oob-idempotent-control",) + assert second_result.created == () + assert second_result.skipped_existing == ("oob-idempotent-control",) + assert len(_fetch_controls()) == 1 + assert len(_fetch_versions()) == 1 + + +@pytest.mark.asyncio +async def test_seed_treats_duplicate_insert_integrity_error_as_skip( + monkeypatch: pytest.MonkeyPatch, +) -> None: + template = _template(name="oob-race-control") + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + async def active_control_name_exists( + self: ControlService, + name: str, + *, + namespace_key: str, + exclude_control_id: int | None = None, + ) -> bool: + return False + + monkeypatch.setattr(ControlService, "active_control_name_exists", active_control_name_exists) + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + assert result.created == () + assert result.skipped_existing == () + assert result.skipped_conflict == ("oob-race-control",) + assert len(_fetch_controls()) == 1 + assert len(_fetch_versions()) == 1 + From 51d1a1282076a79f60b538e39608af13a28b5a78 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 13:25:10 -0700 Subject: [PATCH 02/18] phase 2 add controls --- .../bootstrap/out_of_box_controls.py | 250 +++++++++++++++++- .../test_out_of_box_controls_bootstrap.py | 149 +++++++++++ 2 files changed, 394 insertions(+), 5 deletions(-) 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 3ba30b1d..96ca5d65 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 @@ -35,6 +31,7 @@ ) _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) @@ -104,7 +101,250 @@ 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( + 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( + 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( + 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( + 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"\b(?:rm\s+-rf\s+(?:/|~|\$HOME)|sudo\s+rm\s+-rf|" + r"mkfs(?:\.[a-z0-9]+)?|dd\s+if=[^\s]+\s+of=/dev/[^\s]+|" + r"chmod\s+-R\s+777\s+/|chown\s+-R\s+[^|;&]*\s+/|" + r"shutdown\s+(?:-h\s+)?now|reboot)\b" + ), + "flags": ["IGNORECASE"], + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["tool", "shell", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + 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} + }, + }, + { + "required": ["amount"], + "properties": { + "amount": {"type": "number", "exclusiveMinimum": 10000} + }, + "anyOf": [ + { + "required": ["approved"], + "properties": {"approved": {"const": True}}, + }, + { + "required": ["approval"], + "properties": { + "approval": { + "type": "object", + "required": ["approved"], + "properties": {"approved": {"const": True}}, + } + }, + }, + ], + }, + ], + } + }, + step_types=["tool"], + stages=["pre"], + decision="steer", + steering_message=( + "This high-value action requires approval. Ask for approval, record it " + "in the tool input, then retry." + ), + tags=["tool", "approval", "json"], + ), + ), + OutOfBoxControlTemplate.from_payload( + 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"]}, + ] + } + }, + { + "required": ["approved"], + "properties": {"approved": {"const": True}}, + }, + { + "required": ["approval"], + "properties": { + "approval": { + "type": "object", + "required": ["approved"], + "properties": {"approved": {"const": True}}, + } + }, + }, + ], + } + }, + step_types=["tool"], + stages=["pre"], + decision="steer", + steering_message=( + "Outbound communication requires approval. Ask the user to approve the " + "recipient and message before sending." + ), + tags=["tool", "approval", "exfiltration", "json"], + ), + ), + OutOfBoxControlTemplate.from_payload( + name="oob-sensitive-tool-requires-approved-role", + data=_leaf_control_payload( + description="Deny sensitive tool use when runtime context has an unapproved role.", + selector_path="context.user.role", + evaluator_name="list", + evaluator_config={ + "values": ["admin", "security", "compliance", "manager"], + "logic": "any", + "match_on": "no_match", + "match_mode": "exact", + "case_sensitive": False, + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["tool", "rbac", "list"], + ), + ), + OutOfBoxControlTemplate.from_payload( + 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="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/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 5ba30577..99845e9b 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -5,7 +5,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, @@ -26,6 +34,18 @@ 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-sensitive-tool-requires-approved-role", + "oob-only-approved-tools-may-run", +) +_AVAILABLE_PHASE_2_EVALUATORS = {"regex", "json", "list"} + def _control_payload(*, evaluator_name: str = "regex") -> dict[str, object]: return { @@ -71,10 +91,31 @@ 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 + + def test_missing_required_evaluators_returns_sorted_names() -> None: missing = missing_required_evaluators( {"galileo.luna", "regex", "json"}, @@ -148,6 +189,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") @@ -208,3 +291,69 @@ async def active_control_name_exists( 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)) + shell_result = await shell_evaluator.evaluate("sudo rm -rf /") + assert shell_result.matched is True + + +@pytest.mark.asyncio +async def test_json_out_of_box_controls_match_missing_approval_only() -> 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}) + approved_result = 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 approved_result.matched is False + + 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"}) + approved_outbound_result = await outbound_evaluator.evaluate( + {"to": "customer@example.com", "message": "Hello", "approved": True} + ) + + assert outbound_result.matched is True + assert internal_result.matched is False + assert approved_outbound_result.matched is False + + +@pytest.mark.asyncio +async def test_list_out_of_box_controls_match_unapproved_values() -> None: + role_spec = _oob_evaluator_spec("oob-sensitive-tool-requires-approved-role") + role_evaluator = ListEvaluator(ListEvaluatorConfig.model_validate(role_spec.config)) + + viewer_result = await role_evaluator.evaluate("viewer") + admin_result = await role_evaluator.evaluate("admin") + + assert viewer_result.matched is True + assert admin_result.matched is False + + 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 From e6de28e2b6b68949f79a92269940bcdb9cc005d0 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 16:54:48 -0700 Subject: [PATCH 03/18] fix tests --- .../endpoints/controls.py | 95 ++++++++++++++++++- server/tests/test_principal_namespace_flow.py | 19 ++++ 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index d328c7f9..39a525ea 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -43,7 +43,11 @@ 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 ( + default_out_of_box_namespace_key, + seed_out_of_box_controls, +) +from ..db import AsyncSessionLocal, get_async_db from ..errors import ( APIError, APIValidationError, @@ -54,7 +58,7 @@ NotFoundError, ) from ..logging_utils import get_logger -from ..models import Agent, AgentData +from ..models import Agent, AgentData, Control from ..services.condition_traversal import iter_condition_leaves_with_paths from ..services.control_bindings import ControlBindingsService from ..services.control_definitions import parse_control_definition_or_api_error @@ -257,6 +261,78 @@ def _validate_attachment_filters( ) +async def _seed_out_of_box_controls_for_namespace( + db: AsyncSession, + *, + namespace_key: str, +) -> None: + """Best-effort namespace seeding for browse/list surfaces.""" + try: + if namespace_key == default_out_of_box_namespace_key(): + return + if await _namespace_has_active_controls(db, namespace_key=namespace_key): + return + await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=namespace_key, + available_evaluators=set(list_evaluators().keys()), + ) + except Exception: + await db.rollback() + _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 None + 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 + ) + + +async def _namespace_has_active_controls( + db: AsyncSession, + *, + namespace_key: str, +) -> bool: + result = await db.execute( + select(Control.id) + .where( + Control.namespace_key == namespace_key, + Control.deleted_at.is_(None), + ) + .limit(1) + ) + return result.first() is not None + + def _serialize_control_data( control_data: ControlDefinition | UnrenderedTemplateControl, ) -> dict[str, object]: @@ -1234,6 +1310,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(db, 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_principal_namespace_flow.py b/server/tests/test_principal_namespace_flow.py index 8f16a795..746bc620 100644 --- a/server/tests/test_principal_namespace_flow.py +++ b/server/tests/test_principal_namespace_flow.py @@ -11,6 +11,7 @@ Principal, set_authorizer, ) +from agent_control_server.bootstrap.out_of_box_controls import OUT_OF_BOX_CONTROL_TEMPLATES from fastapi import FastAPI, Request from fastapi.testclient import TestClient @@ -73,6 +74,24 @@ def _evaluation_payload(agent_name: str) -> dict[str, Any]: } +def test_controls_list_seeds_out_of_box_controls_for_principal_namespace( + app: FastAPI, +) -> None: + set_authorizer(HeaderNamespaceAuthorizer()) + + 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) + + def test_principal_namespace_scopes_management_and_runtime(app: FastAPI) -> None: set_authorizer(HeaderNamespaceAuthorizer()) From e6eba19cd71ce16cc44c91ded08b833a3f63d756 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 17:50:45 -0700 Subject: [PATCH 04/18] fix lazy controls attachment to namespace --- server/src/agent_control_server/config.py | 11 +++++ server/src/agent_control_server/main.py | 55 ++++++++++++++++------- server/tests/test_config.py | 25 +++++++++++ server/tests/test_main_lifespan.py | 29 +++++++++++- 4 files changed, 103 insertions(+), 17 deletions(-) diff --git a/server/src/agent_control_server/config.py b/server/src/agent_control_server/config.py index fe481881..13c92388 100644 --- a/server/src/agent_control_server/config.py +++ b/server/src/agent_control_server/config.py @@ -207,6 +207,13 @@ class Settings(BaseSettings): "AGENT_CONTROL_ALLOW_HEADERS", "ALLOW_HEADERS", ) + out_of_box_namespace_keys: list[str] | str = _env_alias_field( + [], + "AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", + "OUT_OF_BOX_NAMESPACE_KEYS", + "GALILEO_ORGANIZATION_ID", + "GALILEO_ORGANIZATION_IDS", + ) def get_cors_origins(self) -> list[str]: """Parse CORS origins from string or list.""" @@ -220,6 +227,10 @@ def get_allow_headers(self) -> list[str]: """Parse allow_headers from string or list.""" return self._parse_list_setting(self.allow_headers) + def get_out_of_box_namespace_keys(self) -> list[str]: + """Parse namespace keys that should receive startup out-of-box controls.""" + return self._parse_list_setting(self.out_of_box_namespace_keys) + @staticmethod def _parse_list_setting(value: list[str] | str) -> list[str]: """Parse wildcard/comma-separated settings from string or list.""" diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index bed2bf9c..8fd88074 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -89,6 +89,22 @@ def _configure_logging_once() -> None: _logging_configured = True +def _out_of_box_bootstrap_namespace_keys() -> tuple[str, ...]: + """Return startup namespaces that should receive out-of-box controls.""" + namespace_keys: list[str] = [] + seen: set[str] = set() + for namespace_key in ( + default_out_of_box_namespace_key(), + *settings.get_out_of_box_namespace_keys(), + ): + normalized = namespace_key.strip() + if not normalized or normalized in seen: + continue + namespace_keys.append(normalized) + seen.add(normalized) + return tuple(namespace_keys) + + def add_prometheus_metrics(app: FastAPI, metrics_prefix: str) -> None: """Configure Prometheus metrics for the FastAPI app.""" app.add_middleware( @@ -146,23 +162,30 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: available = list(list_evaluators().keys()) logger.info(f"Evaluator discovery complete. Available evaluators: {available}") - try: - seed_result = await seed_out_of_box_controls( - session_factory=AsyncSessionLocal, - namespace_key=default_out_of_box_namespace_key(), - available_evaluators=set(available), - ) - if seed_result.created_count or seed_result.skipped_count: - logger.info( - "Out-of-box control bootstrap complete: created=%s " - "skipped_existing=%s skipped_missing_evaluator=%s skipped_conflict=%s", - seed_result.created_count, - len(seed_result.skipped_existing), - len(seed_result.skipped_missing_evaluator), - len(seed_result.skipped_conflict), + for namespace_key in _out_of_box_bootstrap_namespace_keys(): + try: + seed_result = await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=namespace_key, + available_evaluators=set(available), + ) + if seed_result.created_count or seed_result.skipped_count: + logger.info( + "Out-of-box control bootstrap complete for namespace '%s': " + "created=%s skipped_existing=%s skipped_missing_evaluator=%s " + "skipped_conflict=%s", + namespace_key, + seed_result.created_count, + len(seed_result.skipped_existing), + len(seed_result.skipped_missing_evaluator), + len(seed_result.skipped_conflict), + ) + except Exception: + logger.warning( + "Out-of-box control bootstrap failed for namespace '%s'; continuing startup", + namespace_key, + exc_info=True, ) - except Exception: - logger.warning("Out-of-box control bootstrap failed; continuing startup", exc_info=True) # Initialize observability components (stored on app.state) if observability_settings.enabled: diff --git a/server/tests/test_config.py b/server/tests/test_config.py index 2a9bd472..b9a18fd8 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -108,6 +108,31 @@ def test_db_config_reads_pool_settings_from_env(monkeypatch) -> None: assert config.statement_timeout_seconds == 2.5 +def test_settings_parses_out_of_box_namespace_keys_from_env(monkeypatch) -> None: + # Given: startup OOTB namespace keys are configured as a comma-separated env var + monkeypatch.setenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", "org-a, org-b,, org-c") + + # When: loading server settings + config = Settings() + + # Then: the namespace keys are parsed for startup bootstrap seeding + assert config.get_out_of_box_namespace_keys() == ["org-a", "org-b", "org-c"] + + +def test_settings_reads_galileo_organization_id_for_out_of_box_namespace( + monkeypatch, +) -> None: + # Given: Galileo provides a single organization id in the environment + monkeypatch.delenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", raising=False) + monkeypatch.setenv("GALILEO_ORGANIZATION_ID", "org-devstack") + + # When: loading server settings + config = Settings() + + # Then: the organization id is treated as a startup OOTB namespace key + assert config.get_out_of_box_namespace_keys() == ["org-devstack"] + + def test_db_config_pool_defaults(monkeypatch) -> None: # Given: no pool or timeout settings in the environment for name in ( diff --git a/server/tests/test_main_lifespan.py b/server/tests/test_main_lifespan.py index 1d464e68..7e0b823e 100644 --- a/server/tests/test_main_lifespan.py +++ b/server/tests/test_main_lifespan.py @@ -7,6 +7,7 @@ import textwrap from agent_control_server import main as main_module +from agent_control_server.bootstrap.out_of_box_controls import OutOfBoxSeedResult from agent_control_server.config import observability_settings, settings from agent_control_server.main import lifespan from agent_control_server.observability.sinks import ( @@ -217,6 +218,29 @@ def test_lifespan_skips_observability_when_disabled(monkeypatch) -> None: assert not hasattr(app.state, "event_ingestor") +def test_lifespan_seeds_configured_out_of_box_namespaces(monkeypatch) -> None: + calls: list[str] = [] + + async def fake_seed_out_of_box_controls(**kwargs: object) -> OutOfBoxSeedResult: + calls.append(str(kwargs["namespace_key"])) + return OutOfBoxSeedResult() + + monkeypatch.setattr(observability_settings, "enabled", False) + monkeypatch.setattr( + type(settings), + "get_out_of_box_namespace_keys", + lambda self: ["org-a", "default", "org-b"], + ) + monkeypatch.setattr(main_module, "seed_out_of_box_controls", fake_seed_out_of_box_controls) + + app = FastAPI(lifespan=lifespan) + + with TestClient(app): + pass + + assert calls == ["default", "org-a", "org-b"] + + def test_lifespan_fails_open_when_out_of_box_bootstrap_fails(monkeypatch, caplog) -> None: async def fail_seed_out_of_box_controls(**kwargs: object) -> None: raise RuntimeError("boom") @@ -230,7 +254,10 @@ async def fail_seed_out_of_box_controls(**kwargs: object) -> None: with TestClient(app): pass - assert "Out-of-box control bootstrap failed; continuing startup" in caplog.text + assert ( + "Out-of-box control bootstrap failed for namespace 'default'; continuing startup" + in caplog.text + ) def test_custom_openapi_replaces_jsonvalue_variants(monkeypatch) -> None: From 9ecd541d1bdebd5faf688b1a39b22e5c38f9fddd Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 17:51:44 -0700 Subject: [PATCH 05/18] fix tests --- server/tests/test_config.py | 34 ++++++++++++++++++++++++------ server/tests/test_main_lifespan.py | 4 ++-- 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/server/tests/test_config.py b/server/tests/test_config.py index b9a18fd8..7d1ef43f 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -110,27 +110,49 @@ def test_db_config_reads_pool_settings_from_env(monkeypatch) -> None: def test_settings_parses_out_of_box_namespace_keys_from_env(monkeypatch) -> None: # Given: startup OOTB namespace keys are configured as a comma-separated env var - monkeypatch.setenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", "org-a, org-b,, org-c") + monkeypatch.setenv( + "AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", + "namespace-alpha, namespace-beta,, namespace-gamma", + ) # When: loading server settings config = Settings() # Then: the namespace keys are parsed for startup bootstrap seeding - assert config.get_out_of_box_namespace_keys() == ["org-a", "org-b", "org-c"] + assert config.get_out_of_box_namespace_keys() == [ + "namespace-alpha", + "namespace-beta", + "namespace-gamma", + ] def test_settings_reads_galileo_organization_id_for_out_of_box_namespace( monkeypatch, ) -> None: - # Given: Galileo provides a single organization id in the environment + # Given: Galileo provides a deployment-specific organization id in the environment monkeypatch.delenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", raising=False) - monkeypatch.setenv("GALILEO_ORGANIZATION_ID", "org-devstack") + monkeypatch.setenv("GALILEO_ORGANIZATION_ID", "namespace-7f3c9a") # When: loading server settings config = Settings() - # Then: the organization id is treated as a startup OOTB namespace key - assert config.get_out_of_box_namespace_keys() == ["org-devstack"] + # Then: the arbitrary organization id is treated as a startup OOTB namespace key + assert config.get_out_of_box_namespace_keys() == ["namespace-7f3c9a"] + + +def test_settings_reads_galileo_organization_ids_for_out_of_box_namespaces( + monkeypatch, +) -> None: + # Given: Galileo provides deployment-specific organization ids in the environment + monkeypatch.delenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", raising=False) + monkeypatch.delenv("GALILEO_ORGANIZATION_ID", raising=False) + monkeypatch.setenv("GALILEO_ORGANIZATION_IDS", "namespace-one, namespace-two") + + # When: loading server settings + config = Settings() + + # Then: each arbitrary organization id is treated as a startup OOTB namespace key + assert config.get_out_of_box_namespace_keys() == ["namespace-one", "namespace-two"] def test_db_config_pool_defaults(monkeypatch) -> None: diff --git a/server/tests/test_main_lifespan.py b/server/tests/test_main_lifespan.py index 7e0b823e..feb3e62a 100644 --- a/server/tests/test_main_lifespan.py +++ b/server/tests/test_main_lifespan.py @@ -229,7 +229,7 @@ async def fake_seed_out_of_box_controls(**kwargs: object) -> OutOfBoxSeedResult: monkeypatch.setattr( type(settings), "get_out_of_box_namespace_keys", - lambda self: ["org-a", "default", "org-b"], + lambda self: ["namespace-alpha", "default", "namespace-beta"], ) monkeypatch.setattr(main_module, "seed_out_of_box_controls", fake_seed_out_of_box_controls) @@ -238,7 +238,7 @@ async def fake_seed_out_of_box_controls(**kwargs: object) -> OutOfBoxSeedResult: with TestClient(app): pass - assert calls == ["default", "org-a", "org-b"] + assert calls == ["default", "namespace-alpha", "namespace-beta"] def test_lifespan_fails_open_when_out_of_box_bootstrap_fails(monkeypatch, caplog) -> None: From 37e6f5961882c4ddadd5c3f27d670a0e8fc41a72 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 18:06:27 -0700 Subject: [PATCH 06/18] create controls when controls tab is clicked --- server/src/agent_control_server/config.py | 12 ----- server/src/agent_control_server/main.py | 55 +++++++---------------- server/tests/test_config.py | 47 ------------------- server/tests/test_main_lifespan.py | 29 +----------- 4 files changed, 17 insertions(+), 126 deletions(-) diff --git a/server/src/agent_control_server/config.py b/server/src/agent_control_server/config.py index 13c92388..598a4bde 100644 --- a/server/src/agent_control_server/config.py +++ b/server/src/agent_control_server/config.py @@ -207,14 +207,6 @@ class Settings(BaseSettings): "AGENT_CONTROL_ALLOW_HEADERS", "ALLOW_HEADERS", ) - out_of_box_namespace_keys: list[str] | str = _env_alias_field( - [], - "AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", - "OUT_OF_BOX_NAMESPACE_KEYS", - "GALILEO_ORGANIZATION_ID", - "GALILEO_ORGANIZATION_IDS", - ) - def get_cors_origins(self) -> list[str]: """Parse CORS origins from string or list.""" return self._parse_list_setting(self.cors_origins) @@ -227,10 +219,6 @@ def get_allow_headers(self) -> list[str]: """Parse allow_headers from string or list.""" return self._parse_list_setting(self.allow_headers) - def get_out_of_box_namespace_keys(self) -> list[str]: - """Parse namespace keys that should receive startup out-of-box controls.""" - return self._parse_list_setting(self.out_of_box_namespace_keys) - @staticmethod def _parse_list_setting(value: list[str] | str) -> list[str]: """Parse wildcard/comma-separated settings from string or list.""" diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index 8fd88074..bed2bf9c 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -89,22 +89,6 @@ def _configure_logging_once() -> None: _logging_configured = True -def _out_of_box_bootstrap_namespace_keys() -> tuple[str, ...]: - """Return startup namespaces that should receive out-of-box controls.""" - namespace_keys: list[str] = [] - seen: set[str] = set() - for namespace_key in ( - default_out_of_box_namespace_key(), - *settings.get_out_of_box_namespace_keys(), - ): - normalized = namespace_key.strip() - if not normalized or normalized in seen: - continue - namespace_keys.append(normalized) - seen.add(normalized) - return tuple(namespace_keys) - - def add_prometheus_metrics(app: FastAPI, metrics_prefix: str) -> None: """Configure Prometheus metrics for the FastAPI app.""" app.add_middleware( @@ -162,30 +146,23 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: available = list(list_evaluators().keys()) logger.info(f"Evaluator discovery complete. Available evaluators: {available}") - for namespace_key in _out_of_box_bootstrap_namespace_keys(): - try: - seed_result = await seed_out_of_box_controls( - session_factory=AsyncSessionLocal, - namespace_key=namespace_key, - available_evaluators=set(available), - ) - if seed_result.created_count or seed_result.skipped_count: - logger.info( - "Out-of-box control bootstrap complete for namespace '%s': " - "created=%s skipped_existing=%s skipped_missing_evaluator=%s " - "skipped_conflict=%s", - namespace_key, - seed_result.created_count, - len(seed_result.skipped_existing), - len(seed_result.skipped_missing_evaluator), - len(seed_result.skipped_conflict), - ) - except Exception: - logger.warning( - "Out-of-box control bootstrap failed for namespace '%s'; continuing startup", - namespace_key, - exc_info=True, + try: + seed_result = await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=default_out_of_box_namespace_key(), + available_evaluators=set(available), + ) + if seed_result.created_count or seed_result.skipped_count: + logger.info( + "Out-of-box control bootstrap complete: created=%s " + "skipped_existing=%s skipped_missing_evaluator=%s skipped_conflict=%s", + seed_result.created_count, + len(seed_result.skipped_existing), + len(seed_result.skipped_missing_evaluator), + len(seed_result.skipped_conflict), ) + except Exception: + logger.warning("Out-of-box control bootstrap failed; continuing startup", exc_info=True) # Initialize observability components (stored on app.state) if observability_settings.enabled: diff --git a/server/tests/test_config.py b/server/tests/test_config.py index 7d1ef43f..2a9bd472 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -108,53 +108,6 @@ def test_db_config_reads_pool_settings_from_env(monkeypatch) -> None: assert config.statement_timeout_seconds == 2.5 -def test_settings_parses_out_of_box_namespace_keys_from_env(monkeypatch) -> None: - # Given: startup OOTB namespace keys are configured as a comma-separated env var - monkeypatch.setenv( - "AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", - "namespace-alpha, namespace-beta,, namespace-gamma", - ) - - # When: loading server settings - config = Settings() - - # Then: the namespace keys are parsed for startup bootstrap seeding - assert config.get_out_of_box_namespace_keys() == [ - "namespace-alpha", - "namespace-beta", - "namespace-gamma", - ] - - -def test_settings_reads_galileo_organization_id_for_out_of_box_namespace( - monkeypatch, -) -> None: - # Given: Galileo provides a deployment-specific organization id in the environment - monkeypatch.delenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", raising=False) - monkeypatch.setenv("GALILEO_ORGANIZATION_ID", "namespace-7f3c9a") - - # When: loading server settings - config = Settings() - - # Then: the arbitrary organization id is treated as a startup OOTB namespace key - assert config.get_out_of_box_namespace_keys() == ["namespace-7f3c9a"] - - -def test_settings_reads_galileo_organization_ids_for_out_of_box_namespaces( - monkeypatch, -) -> None: - # Given: Galileo provides deployment-specific organization ids in the environment - monkeypatch.delenv("AGENT_CONTROL_OUT_OF_BOX_NAMESPACE_KEYS", raising=False) - monkeypatch.delenv("GALILEO_ORGANIZATION_ID", raising=False) - monkeypatch.setenv("GALILEO_ORGANIZATION_IDS", "namespace-one, namespace-two") - - # When: loading server settings - config = Settings() - - # Then: each arbitrary organization id is treated as a startup OOTB namespace key - assert config.get_out_of_box_namespace_keys() == ["namespace-one", "namespace-two"] - - def test_db_config_pool_defaults(monkeypatch) -> None: # Given: no pool or timeout settings in the environment for name in ( diff --git a/server/tests/test_main_lifespan.py b/server/tests/test_main_lifespan.py index feb3e62a..1d464e68 100644 --- a/server/tests/test_main_lifespan.py +++ b/server/tests/test_main_lifespan.py @@ -7,7 +7,6 @@ import textwrap from agent_control_server import main as main_module -from agent_control_server.bootstrap.out_of_box_controls import OutOfBoxSeedResult from agent_control_server.config import observability_settings, settings from agent_control_server.main import lifespan from agent_control_server.observability.sinks import ( @@ -218,29 +217,6 @@ def test_lifespan_skips_observability_when_disabled(monkeypatch) -> None: assert not hasattr(app.state, "event_ingestor") -def test_lifespan_seeds_configured_out_of_box_namespaces(monkeypatch) -> None: - calls: list[str] = [] - - async def fake_seed_out_of_box_controls(**kwargs: object) -> OutOfBoxSeedResult: - calls.append(str(kwargs["namespace_key"])) - return OutOfBoxSeedResult() - - monkeypatch.setattr(observability_settings, "enabled", False) - monkeypatch.setattr( - type(settings), - "get_out_of_box_namespace_keys", - lambda self: ["namespace-alpha", "default", "namespace-beta"], - ) - monkeypatch.setattr(main_module, "seed_out_of_box_controls", fake_seed_out_of_box_controls) - - app = FastAPI(lifespan=lifespan) - - with TestClient(app): - pass - - assert calls == ["default", "namespace-alpha", "namespace-beta"] - - def test_lifespan_fails_open_when_out_of_box_bootstrap_fails(monkeypatch, caplog) -> None: async def fail_seed_out_of_box_controls(**kwargs: object) -> None: raise RuntimeError("boom") @@ -254,10 +230,7 @@ async def fail_seed_out_of_box_controls(**kwargs: object) -> None: with TestClient(app): pass - assert ( - "Out-of-box control bootstrap failed for namespace 'default'; continuing startup" - in caplog.text - ) + assert "Out-of-box control bootstrap failed; continuing startup" in caplog.text def test_custom_openapi_replaces_jsonvalue_variants(monkeypatch) -> None: From c6aae0a936272d8168816013d86fe4979fdb732c Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 18:45:17 -0700 Subject: [PATCH 07/18] show for clone=false --- server/src/agent_control_server/endpoints/controls.py | 2 +- server/tests/test_principal_namespace_flow.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index 39a525ea..39be7f72 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -306,7 +306,7 @@ def _should_seed_out_of_box_controls_on_list( and name is None and enabled is None and template_backed is None - and cloned is None + and cloned is not True and step_type is None and stage is None and execution is None diff --git a/server/tests/test_principal_namespace_flow.py b/server/tests/test_principal_namespace_flow.py index 746bc620..af8b0dfd 100644 --- a/server/tests/test_principal_namespace_flow.py +++ b/server/tests/test_principal_namespace_flow.py @@ -84,7 +84,7 @@ def test_controls_list_seeds_out_of_box_controls_for_principal_namespace( assert filtered.status_code == 200, filtered.text assert filtered.json()["controls"] == [] - resp = namespace_client.get("/api/v1/controls", params={"limit": 10}) + resp = namespace_client.get("/api/v1/controls", params={"limit": 10, "cloned": "false"}) assert resp.status_code == 200, resp.text expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} From 54c06d8358c9db7766e69a702125121cc310b418 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 30 Jul 2026 08:41:24 -0700 Subject: [PATCH 08/18] address comments --- ...d7e2b4_out_of_box_control_seed_identity.py | 39 +++++++++ .../bootstrap/out_of_box_controls.py | 33 +++++++- server/src/agent_control_server/config.py | 7 ++ server/src/agent_control_server/main.py | 17 ++-- server/src/agent_control_server/models.py | 12 +++ .../agent_control_server/services/controls.py | 18 +++++ server/tests/test_config.py | 2 + server/tests/test_main_lifespan.py | 28 +++++++ .../test_out_of_box_controls_bootstrap.py | 80 ++++++++++++++++++- 9 files changed, 224 insertions(+), 12 deletions(-) create mode 100644 server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py diff --git a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py new file mode 100644 index 00000000..63924de5 --- /dev/null +++ b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py @@ -0,0 +1,39 @@ +"""add immutable out-of-box control seed identity + +Revision ID: f3a1c8d7e2b4 +Revises: e2b7f4a9c6d1 +Create Date: 2026-07-30 12:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "f3a1c8d7e2b4" +down_revision = "e2b7f4a9c6d1" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("controls", sa.Column("seed_source_id", sa.String(length=255), nullable=True)) + op.add_column( + "controls", + sa.Column("seed_opted_out_at", sa.DateTime(timezone=True), nullable=True), + ) + op.create_index( + "idx_controls_namespace_seed_source", + "controls", + ["namespace_key", "seed_source_id"], + unique=True, + postgresql_where=sa.text("seed_source_id IS NOT NULL"), + ) + + +def downgrade() -> None: + op.drop_index("idx_controls_namespace_seed_source", table_name="controls") + op.drop_column("controls", "seed_opted_out_at") + op.drop_column("controls", "seed_source_id") 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 3ba30b1d..6c523b65 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 @@ -33,6 +33,7 @@ "idx_controls_namespace_name_active", } ) +_CONTROL_SEED_UNIQUE_CONSTRAINT = "idx_controls_namespace_seed_source" _INITIAL_VERSION_NOTE = "Out-of-box control seed" _SLUG_NAME_ADAPTER = TypeAdapter(SlugName) @@ -41,11 +42,13 @@ class OutOfBoxControlTemplate: """Validated control definition plus the evaluator names it needs.""" + source_id: str name: str control: ControlDefinition required_evaluators: frozenset[str] = field(default_factory=frozenset) def __post_init__(self) -> None: + object.__setattr__(self, "source_id", _SLUG_NAME_ADAPTER.validate_python(self.source_id)) object.__setattr__(self, "name", _SLUG_NAME_ADAPTER.validate_python(self.name)) if not self.required_evaluators: required_evaluators = { @@ -60,12 +63,14 @@ def __post_init__(self) -> None: def from_payload( cls, *, + source_id: str, name: str, data: Mapping[str, object], required_evaluators: Collection[str] = frozenset(), ) -> Self: """Build a template from raw JSON-like data and validate it immediately.""" return cls( + source_id=source_id, name=name, control=ControlDefinition.model_validate(data), required_evaluators=frozenset(required_evaluators), @@ -130,9 +135,10 @@ async def seed_out_of_box_controls( ) -> OutOfBoxSeedResult: """Create missing out-of-box controls in a namespace. - Existing active controls are left untouched so customer edits survive - restarts and upgrades. Duplicate-name integrity errors are treated as - benign races with another pod and are reported as ``skipped_conflict``. + Existing seeded controls are found by immutable source ID, so customer + renames and explicit deletion opt-outs survive restarts and upgrades. + Duplicate-name and duplicate-source integrity errors are treated as benign + races with another pod and are reported as ``skipped_conflict``. """ if not templates: return OutOfBoxSeedResult() @@ -185,6 +191,11 @@ async def _seed_one_control( template: OutOfBoxControlTemplate, ) -> str: control_service = ControlService(session) + if await control_service.seed_source_exists( + template.source_id, + namespace_key=namespace_key, + ): + return "existing" if await control_service.active_control_name_exists(template.name, namespace_key=namespace_key): return "existing" @@ -192,6 +203,7 @@ async def _seed_one_control( namespace_key=namespace_key, name=template.name, data=_serialize_control_data(template.control), + seed_source_id=template.source_id, ) try: await control_service.create_version( @@ -202,7 +214,7 @@ async def _seed_one_control( await session.commit() except IntegrityError as exc: await session.rollback() - if _is_control_name_conflict(exc): + if _is_control_seed_conflict(exc): return "conflict" raise return "created" @@ -233,3 +245,16 @@ def _is_control_name_conflict(error: IntegrityError) -> bool: part for part in (str(error.orig), str(error)) if part and part != "None" ) return any(name in error_text for name in _CONTROL_NAME_UNIQUE_CONSTRAINTS) + + +def _is_control_seed_conflict(error: IntegrityError) -> bool: + diag = getattr(getattr(error.orig, "diag", None), "constraint_name", None) + if diag == _CONTROL_SEED_UNIQUE_CONSTRAINT: + return True + if _is_control_name_conflict(error): + return True + + error_text = " ".join( + part for part in (str(error.orig), str(error)) if part and part != "None" + ) + return _CONTROL_SEED_UNIQUE_CONSTRAINT in error_text diff --git a/server/src/agent_control_server/config.py b/server/src/agent_control_server/config.py index e9e62f9e..00335611 100644 --- a/server/src/agent_control_server/config.py +++ b/server/src/agent_control_server/config.py @@ -183,6 +183,13 @@ class Settings(BaseSettings): # API settings api_version: str = _env_alias_field("v1", "AGENT_CONTROL_API_VERSION", "API_VERSION") api_prefix: str = _env_alias_field("/api", "AGENT_CONTROL_API_PREFIX", "API_PREFIX") + out_of_box_bootstrap_timeout_seconds: float = Field( + default=10.0, + gt=0, + validation_alias=AliasChoices( + "AGENT_CONTROL_OUT_OF_BOX_BOOTSTRAP_TIMEOUT_SECONDS", + ), + ) # Prometheus metrics settings prometheus_metrics_prefix: str = _env_alias_field( diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index 1085c47d..bd0b9efd 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -1,5 +1,6 @@ """Main server application entry point.""" +import asyncio import inspect import logging from collections.abc import AsyncGenerator @@ -152,11 +153,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: logger.info(f"Evaluator discovery complete. Available evaluators: {available}") try: - seed_result = await seed_out_of_box_controls( - session_factory=AsyncSessionLocal, - namespace_key=default_out_of_box_namespace_key(), - available_evaluators=set(available), - ) + async with asyncio.timeout(settings.out_of_box_bootstrap_timeout_seconds): + seed_result = await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=default_out_of_box_namespace_key(), + available_evaluators=set(available), + ) if seed_result.created_count or seed_result.skipped_count: logger.info( "Out-of-box control bootstrap complete: created=%s " @@ -166,6 +168,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: len(seed_result.skipped_missing_evaluator), len(seed_result.skipped_conflict), ) + except TimeoutError: + logger.warning( + "Out-of-box control bootstrap timed out after %s seconds; continuing startup", + settings.out_of_box_bootstrap_timeout_seconds, + ) except Exception: logger.warning("Out-of-box control bootstrap failed; continuing startup", exc_info=True) diff --git a/server/src/agent_control_server/models.py b/server/src/agent_control_server/models.py index c31ccddf..3b1dd4e5 100644 --- a/server/src/agent_control_server/models.py +++ b/server/src/agent_control_server/models.py @@ -181,6 +181,14 @@ class Control(Base): postgresql_where=text("cloned_from_control_id IS NOT NULL"), sqlite_where=text("cloned_from_control_id IS NOT NULL"), ), + Index( + "idx_controls_namespace_seed_source", + "namespace_key", + "seed_source_id", + unique=True, + postgresql_where=text("seed_source_id IS NOT NULL"), + sqlite_where=text("seed_source_id IS NOT NULL"), + ), ) id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) @@ -195,6 +203,10 @@ class Control(Base): cloned_from_control_id: Mapped[int | None] = mapped_column( Integer, nullable=True ) + seed_source_id: Mapped[str | None] = mapped_column(String(255), nullable=True) + seed_opted_out_at: Mapped[dt.datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) deleted_at: Mapped[dt.datetime | None] = mapped_column( DateTime(timezone=True), nullable=True ) diff --git a/server/src/agent_control_server/services/controls.py b/server/src/agent_control_server/services/controls.py index 293fc130..619cbed1 100644 --- a/server/src/agent_control_server/services/controls.py +++ b/server/src/agent_control_server/services/controls.py @@ -126,6 +126,7 @@ def create_control( name: str, data: dict[str, Any], cloned_from_control_id: int | None = None, + seed_source_id: str | None = None, ) -> Control: """Create a new pending control row.""" control = Control( @@ -133,6 +134,7 @@ def create_control( name=name, data=data, cloned_from_control_id=cloned_from_control_id, + seed_source_id=seed_source_id, ) self._db.add(control) return control @@ -158,6 +160,22 @@ def set_control_enabled(control: Control, *, enabled: bool) -> None: def mark_control_deleted(control: Control, *, deleted_at: dt.datetime) -> None: """Mark a control as soft-deleted.""" control.deleted_at = deleted_at + if control.seed_source_id is not None: + control.seed_opted_out_at = deleted_at + + async def seed_source_exists( + self, + seed_source_id: str, + *, + namespace_key: str, + ) -> bool: + """Return whether a control has claimed an immutable seed identity.""" + stmt = select(Control.id).where( + Control.namespace_key == namespace_key, + Control.seed_source_id == seed_source_id, + ) + result = await self._db.execute(stmt) + return result.first() is not None async def get_control_or_404( self, diff --git a/server/tests/test_config.py b/server/tests/test_config.py index 62157fae..7cbb3186 100644 --- a/server/tests/test_config.py +++ b/server/tests/test_config.py @@ -148,6 +148,7 @@ def test_settings_reads_agent_control_prefixed_env_vars(monkeypatch) -> None: monkeypatch.setenv("AGENT_CONTROL_CORS_ORIGINS", "https://a.example, https://b.example") monkeypatch.setenv("AGENT_CONTROL_ALLOW_METHODS", "GET, POST") monkeypatch.setenv("AGENT_CONTROL_ALLOW_HEADERS", "Authorization, Content-Type") + monkeypatch.setenv("AGENT_CONTROL_OUT_OF_BOX_BOOTSTRAP_TIMEOUT_SECONDS", "3.5") # When: loading settings from the environment config = Settings() @@ -157,6 +158,7 @@ def test_settings_reads_agent_control_prefixed_env_vars(monkeypatch) -> None: assert config.get_cors_origins() == ["https://a.example", "https://b.example"] assert config.get_allow_methods() == ["GET", "POST"] assert config.get_allow_headers() == ["Authorization", "Content-Type"] + assert config.out_of_box_bootstrap_timeout_seconds == 3.5 def test_settings_reads_legacy_env_vars(monkeypatch) -> None: diff --git a/server/tests/test_main_lifespan.py b/server/tests/test_main_lifespan.py index eee44ac3..8a2b818e 100644 --- a/server/tests/test_main_lifespan.py +++ b/server/tests/test_main_lifespan.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import os import subprocess @@ -233,6 +234,33 @@ async def fail_seed_out_of_box_controls(**kwargs: object) -> None: assert "Out-of-box control bootstrap failed; continuing startup" in caplog.text +def test_lifespan_times_out_blocked_out_of_box_bootstrap(monkeypatch, caplog) -> None: + cancelled = False + + async def block_seed_out_of_box_controls(**kwargs: object) -> None: + nonlocal cancelled + try: + await asyncio.Event().wait() + finally: + cancelled = True + + monkeypatch.setattr(observability_settings, "enabled", False) + monkeypatch.setattr(settings, "out_of_box_bootstrap_timeout_seconds", 0.01) + monkeypatch.setattr(main_module, "seed_out_of_box_controls", block_seed_out_of_box_controls) + + app = FastAPI(lifespan=lifespan) + + with caplog.at_level("WARNING"): + with TestClient(app): + pass + + assert cancelled is True + assert ( + "Out-of-box control bootstrap timed out after 0.01 seconds; continuing startup" + in caplog.text + ) + + def test_custom_openapi_replaces_jsonvalue_variants(monkeypatch) -> None: # Given: a custom openapi generator that includes Pydantic JSONValue schemas json_value_schema_names = ( diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 5ba30577..59135451 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime as dt import uuid from copy import deepcopy from typing import cast @@ -48,10 +49,13 @@ def _control_payload(*, evaluator_name: str = "regex") -> dict[str, object]: def _template( *, name: str | None = None, + source_id: str | None = None, evaluator_name: str = "regex", ) -> OutOfBoxControlTemplate: + template_name = name or f"oob-test-{uuid.uuid4().hex}" return OutOfBoxControlTemplate.from_payload( - name=name or f"oob-test-{uuid.uuid4().hex}", + source_id=source_id or template_name, + name=template_name, data=_control_payload(evaluator_name=evaluator_name), ) @@ -92,7 +96,11 @@ def test_template_from_payload_validates_control_definition() -> None: } with pytest.raises(ValidationError): - OutOfBoxControlTemplate.from_payload(name="invalid-oob-control", data=payload) + OutOfBoxControlTemplate.from_payload( + source_id="invalid-oob-control", + name="invalid-oob-control", + data=payload, + ) @pytest.mark.asyncio @@ -132,6 +140,8 @@ async def test_seed_creates_control_version_in_namespace_without_bindings() -> N control = controls[0] assert control.namespace_key == "galileo-org-123" assert control.name == "oob-create-control" + assert control.seed_source_id == "oob-create-control" + assert control.seed_opted_out_at is None assert control.data["enabled"] is True assert control.data["condition"]["evaluator"]["name"] == "regex" @@ -172,6 +182,62 @@ async def test_seed_is_idempotent_for_existing_active_control_names() -> None: assert len(_fetch_versions()) == 1 +@pytest.mark.asyncio +async def test_seed_does_not_duplicate_a_renamed_seeded_control() -> None: + template = _template(name="oob-original-name", source_id="stable-seed-id") + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + with Session(engine) as session: + control = session.scalar(select(Control)) + assert control is not None + control.name = "customer-renamed-control" + session.commit() + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + assert result.skipped_existing == ("oob-original-name",) + assert [control.name for control in _fetch_controls()] == ["customer-renamed-control"] + + +@pytest.mark.asyncio +async def test_seed_respects_deleted_control_opt_out_tombstone() -> None: + template = _template(name="oob-deleted-control", source_id="stable-seed-id") + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + deleted_at = dt.datetime.now(dt.UTC) + with Session(engine) as session: + control = session.scalar(select(Control)) + assert control is not None + ControlService.mark_control_deleted(control, deleted_at=deleted_at) + session.commit() + + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + assert result.skipped_existing == ("oob-deleted-control",) + controls = _fetch_controls() + assert len(controls) == 1 + assert controls[0].deleted_at == deleted_at + assert controls[0].seed_opted_out_at == deleted_at + + @pytest.mark.asyncio async def test_seed_treats_duplicate_insert_integrity_error_as_skip( monkeypatch: pytest.MonkeyPatch, @@ -193,7 +259,16 @@ async def active_control_name_exists( ) -> bool: return False + async def seed_source_exists( + self: ControlService, + seed_source_id: str, + *, + namespace_key: str, + ) -> bool: + return False + monkeypatch.setattr(ControlService, "active_control_name_exists", active_control_name_exists) + monkeypatch.setattr(ControlService, "seed_source_exists", seed_source_exists) result = await seed_out_of_box_controls( session_factory=AsyncSessionTest, @@ -207,4 +282,3 @@ async def active_control_name_exists( assert result.skipped_conflict == ("oob-race-control",) assert len(_fetch_controls()) == 1 assert len(_fetch_versions()) == 1 - From e954f9ee07dd80e7ee53441efa80b286a53456d1 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 30 Jul 2026 11:01:02 -0700 Subject: [PATCH 09/18] address comments and coverage --- engine/src/agent_control_engine/selectors.py | 2 + engine/tests/test_selectors.py | 12 ++ models/src/agent_control_models/agent.py | 9 ++ models/src/agent_control_models/controls.py | 14 ++- sdks/python/src/agent_control/evaluation.py | 2 + .../src/agent_control/integrations/_core.py | 2 + .../integrations/google_adk/plugin.py | 4 + .../integrations/strands/plugin.py | 1 + sdks/python/tests/test_evaluation.py | 6 +- sdks/python/tests/test_google_adk_plugin.py | 2 +- .../bootstrap/out_of_box_controls.py | 83 +++----------- .../endpoints/controls.py | 33 +----- server/tests/test_controls_additional.py | 97 +++++++++++++++- .../test_out_of_box_controls_bootstrap.py | 61 ++++++---- server/tests/test_principal_namespace_flow.py | 106 +++++++++++++++++- 15 files changed, 302 insertions(+), 132 deletions(-) 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/server/src/agent_control_server/bootstrap/out_of_box_controls.py b/server/src/agent_control_server/bootstrap/out_of_box_controls.py index 924c2b00..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 @@ -196,10 +196,14 @@ def _leaf_control_payload( evaluator_name="regex", evaluator_config={ "pattern": ( - r"\b(?:rm\s+-rf\s+(?:/|~|\$HOME)|sudo\s+rm\s+-rf|" - r"mkfs(?:\.[a-z0-9]+)?|dd\s+if=[^\s]+\s+of=/dev/[^\s]+|" - r"chmod\s+-R\s+777\s+/|chown\s+-R\s+[^|;&]*\s+/|" - r"shutdown\s+(?:-h\s+)?now|reboot)\b" + 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"], }, @@ -229,28 +233,6 @@ def _leaf_control_payload( "amount": {"type": "number", "maximum": 10000} }, }, - { - "required": ["amount"], - "properties": { - "amount": {"type": "number", "exclusiveMinimum": 10000} - }, - "anyOf": [ - { - "required": ["approved"], - "properties": {"approved": {"const": True}}, - }, - { - "required": ["approval"], - "properties": { - "approval": { - "type": "object", - "required": ["approved"], - "properties": {"approved": {"const": True}}, - } - }, - }, - ], - }, ], } }, @@ -258,8 +240,9 @@ def _leaf_control_payload( stages=["pre"], decision="steer", steering_message=( - "This high-value action requires approval. Ask for approval, record it " - "in the tool input, then retry." + "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"], ), @@ -289,21 +272,7 @@ def _leaf_control_payload( {"required": ["destination"]}, ] } - }, - { - "required": ["approved"], - "properties": {"approved": {"const": True}}, - }, - { - "required": ["approval"], - "properties": { - "approval": { - "type": "object", - "required": ["approved"], - "properties": {"approved": {"const": True}}, - } - }, - }, + } ], } }, @@ -311,38 +280,20 @@ def _leaf_control_payload( stages=["pre"], decision="steer", steering_message=( - "Outbound communication requires approval. Ask the user to approve the " - "recipient and message before sending." + "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-sensitive-tool-requires-approved-role", - name="oob-sensitive-tool-requires-approved-role", - data=_leaf_control_payload( - description="Deny sensitive tool use when runtime context has an unapproved role.", - selector_path="context.user.role", - evaluator_name="list", - evaluator_config={ - "values": ["admin", "security", "compliance", "manager"], - "logic": "any", - "match_on": "no_match", - "match_mode": "exact", - "case_sensitive": False, - }, - step_types=["tool"], - stages=["pre"], - decision="deny", - tags=["tool", "rbac", "list"], - ), - ), 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="name", + selector_path="canonical_name", evaluator_name="list", evaluator_config={ "values": ["search", "web_search", "retrieve", "calculator"], diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index 39be7f72..a2ff9f1a 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -43,10 +43,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from ..auth_framework import Operation, Principal, get_authorizer, require_operation -from ..bootstrap.out_of_box_controls import ( - default_out_of_box_namespace_key, - seed_out_of_box_controls, -) +from ..bootstrap.out_of_box_controls import seed_out_of_box_controls from ..db import AsyncSessionLocal, get_async_db from ..errors import ( APIError, @@ -58,7 +55,7 @@ NotFoundError, ) from ..logging_utils import get_logger -from ..models import Agent, AgentData, Control +from ..models import Agent, AgentData from ..services.condition_traversal import iter_condition_leaves_with_paths from ..services.control_bindings import ControlBindingsService from ..services.control_definitions import parse_control_definition_or_api_error @@ -262,23 +259,17 @@ def _validate_attachment_filters( async def _seed_out_of_box_controls_for_namespace( - db: AsyncSession, *, namespace_key: str, ) -> None: - """Best-effort namespace seeding for browse/list surfaces.""" + """Best-effort idempotent namespace seeding for browse/list surfaces.""" try: - if namespace_key == default_out_of_box_namespace_key(): - return - if await _namespace_has_active_controls(db, namespace_key=namespace_key): - return await seed_out_of_box_controls( session_factory=AsyncSessionLocal, namespace_key=namespace_key, available_evaluators=set(list_evaluators().keys()), ) except Exception: - await db.rollback() _logger.warning( "Out-of-box control seed failed for namespace '%s'; continuing request", namespace_key, @@ -317,22 +308,6 @@ def _should_seed_out_of_box_controls_on_list( ) -async def _namespace_has_active_controls( - db: AsyncSession, - *, - namespace_key: str, -) -> bool: - result = await db.execute( - select(Control.id) - .where( - Control.namespace_key == namespace_key, - Control.deleted_at.is_(None), - ) - .limit(1) - ) - return result.first() is not None - - def _serialize_control_data( control_data: ControlDefinition | UnrenderedTemplateControl, ) -> dict[str, object]: @@ -1324,7 +1299,7 @@ async def list_controls( attachment_target_type=attachment_target_type, attachment_target_id=attachment_target_id, ): - await _seed_out_of_box_controls_for_namespace(db, namespace_key=namespace_key) + 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 25092703..2975dbd2 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -42,7 +42,6 @@ "oob-dangerous-shell-command-match", "oob-high-value-action-requires-approval", "oob-outbound-communication-requires-approval", - "oob-sensitive-tool-requires-approved-role", "oob-only-approved-tools-may-run", ) _AVAILABLE_PHASE_2_EVALUATORS = {"regex", "json", "list"} @@ -118,6 +117,14 @@ def test_out_of_box_catalog_contains_phase_2_templates() -> None: 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: @@ -376,12 +383,19 @@ async def test_regex_out_of_box_controls_match_representative_payloads() -> None shell_spec = _oob_evaluator_spec("oob-dangerous-shell-command-match") shell_evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(shell_spec.config)) - shell_result = await shell_evaluator.evaluate("sudo rm -rf /") - assert shell_result.matched is True + 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_match_missing_approval_only() -> None: +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) @@ -389,13 +403,16 @@ async def test_json_out_of_box_controls_match_missing_approval_only() -> None: high_value_result = await high_value_evaluator.evaluate({"amount": 25000}) low_value_result = await high_value_evaluator.evaluate({"amount": 250}) - approved_result = await high_value_evaluator.evaluate( - {"amount": 25000, "approval": {"approved": True}} - ) + 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 approved_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)) @@ -404,26 +421,26 @@ async def test_json_out_of_box_controls_match_missing_approval_only() -> None: {"to": "customer@example.com", "message": "Hello"} ) internal_result = await outbound_evaluator.evaluate({"query": "customer history"}) - approved_outbound_result = await outbound_evaluator.evaluate( - {"to": "customer@example.com", "message": "Hello", "approved": True} - ) + 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 approved_outbound_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_controls_match_unapproved_values() -> None: - role_spec = _oob_evaluator_spec("oob-sensitive-tool-requires-approved-role") - role_evaluator = ListEvaluator(ListEvaluatorConfig.model_validate(role_spec.config)) - - viewer_result = await role_evaluator.evaluate("viewer") - admin_result = await role_evaluator.evaluate("admin") - - assert viewer_result.matched is True - assert admin_result.matched is False - +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)) diff --git a/server/tests/test_principal_namespace_flow.py b/server/tests/test_principal_namespace_flow.py index af8b0dfd..f02cf61b 100644 --- a/server/tests/test_principal_namespace_flow.py +++ b/server/tests/test_principal_namespace_flow.py @@ -6,15 +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 agent_control_server.bootstrap.out_of_box_controls import OUT_OF_BOX_CONTROL_TEMPLATES -from fastapi import FastAPI, Request -from fastapi.testclient import TestClient +from agent_control_server.models import Control +from .conftest import engine from .utils import VALID_CONTROL_PAYLOAD @@ -40,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, @@ -77,19 +99,95 @@ def _evaluation_payload(agent_name: str) -> dict[str, Any]: def test_controls_list_seeds_out_of_box_controls_for_principal_namespace( app: FastAPI, ) -> None: - set_authorizer(HeaderNamespaceAuthorizer()) + 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, "cloned": "false"}) + 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: From 7f3c6bef84e71f889742be06ef1c1e4efc0ef176 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 30 Jul 2026 11:33:32 -0700 Subject: [PATCH 10/18] fix ts sdk --- .../src/generated/models/control-selector.ts | 2 +- sdks/typescript/src/generated/models/step.ts | 18 +++++++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) 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 { From dfb129e6d0510a1c858bac6823953f018035756d Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 6 Aug 2026 14:16:45 -0700 Subject: [PATCH 11/18] address comments --- .../bootstrap/out_of_box_controls.py | 95 ++++++++++++++++++- .../test_out_of_box_controls_bootstrap.py | 84 +++++++++++++++- server/tests/test_principal_namespace_flow.py | 12 ++- 3 files changed, 184 insertions(+), 7 deletions(-) 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 1e28b2a1..f93ba7b7 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 @@ -33,6 +33,10 @@ _INITIAL_VERSION_NOTE = "Out-of-box control seed" _SLUG_NAME_ADAPTER = TypeAdapter(SlugName) _OUT_OF_BOX_TAGS = ["out-of-box"] +_SQL_TOOL_NAME_PATTERN = ( + r"(?i)(?:^|[._-])(?:sql|execute[_-]?sql|run[_-]?sql|sql[_-]?query|" + r"query[_-]?database|execute[_-]?query)(?:$|[._-])" +) @dataclass(frozen=True, slots=True) @@ -117,16 +121,21 @@ def _leaf_control_payload( decision: str, tags: list[str], steering_message: str | None = None, + step_name_regex: str | None = None, ) -> dict[str, object]: action: dict[str, object] = {"decision": decision} if steering_message is not None: action["steering_context"] = {"message": steering_message} + scope: dict[str, object] = {"step_types": step_types, "stages": stages} + if step_name_regex is not None: + scope["step_name_regex"] = step_name_regex + return { "description": description, "enabled": True, "execution": "server", - "scope": {"step_types": step_types, "stages": stages}, + "scope": scope, "condition": { "selector": {"path": selector_path}, "evaluator": { @@ -308,6 +317,90 @@ def _leaf_control_payload( tags=["tool", "allowlist", "list"], ), ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-owasp-llm05-read-only-sql", + name="oob-owasp-llm05-read-only-sql", + data=_leaf_control_payload( + description=("Block SQL tool calls that are not a single read-only SELECT statement."), + selector_path="input.query", + evaluator_name="sql", + evaluator_config={ + "allowed_operations": ["SELECT"], + "allow_multi_statements": False, + "block_ddl": True, + "block_dcl": True, + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["owasp", "owasp-llm05", "owasp-asi02", "tool", "sql"], + step_name_regex=_SQL_TOOL_NAME_PATTERN, + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-owasp-llm10-bounded-sql-query", + name="oob-owasp-llm10-bounded-sql-query", + data=_leaf_control_payload( + description=("Block SQL queries without bounded results or with excessive complexity."), + selector_path="input.query", + evaluator_name="sql", + evaluator_config={ + "require_limit": True, + "max_limit": 1000, + "max_result_window": 1000, + "max_subquery_depth": 3, + "max_joins": 5, + "max_union_count": 2, + }, + step_types=["tool"], + stages=["pre"], + decision="deny", + tags=["owasp", "owasp-llm10", "tool", "sql", "resource-limit"], + step_name_regex=_SQL_TOOL_NAME_PATTERN, + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-owasp-llm02-common-credential-output-match", + name="oob-owasp-llm02-common-credential-output-match", + data=_leaf_control_payload( + description=("Block LLM output containing common private-key or API-token formats."), + selector_path="output", + evaluator_name="regex", + evaluator_config={ + "pattern": ( + r"(?:-----BEGIN (?:RSA |EC |DSA |OPENSSH )?PRIVATE KEY-----|" + r"\b(?:AKIA|ASIA)[A-Z0-9]{16}\b|" + r"\bgh[pousr]_[A-Za-z0-9]{36,255}\b|" + r"\bAIza[0-9A-Za-z_-]{35}\b|" + r"\bxox[baprs]-[A-Za-z0-9-]{10,}\b)" + ) + }, + step_types=["llm"], + stages=["post"], + decision="deny", + tags=["owasp", "owasp-llm02", "credential", "secret", "regex"], + ), + ), + OutOfBoxControlTemplate.from_payload( + source_id="oob-owasp-llm05-dangerous-uri-output-match", + name="oob-owasp-llm05-dangerous-uri-output-match", + data=_leaf_control_payload( + description=("Block LLM output containing executable or active-content URI schemes."), + selector_path="output", + evaluator_name="regex", + evaluator_config={ + "pattern": ( + r"(?:\b(?:javascript|vbscript)\s*:|" + r"\bdata\s*:\s*(?:text/html|application/xhtml\+xml|image/svg\+xml))" + ), + "flags": ["IGNORECASE"], + }, + step_types=["llm"], + stages=["post"], + decision="deny", + tags=["owasp", "owasp-llm05", "output-handling", "uri", "regex"], + ), + ), ) diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 2975dbd2..bb7b7919 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -12,6 +12,7 @@ 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_evaluators.sql import SQLEvaluator, SQLEvaluatorConfig from agent_control_models import EvaluatorSpec from agent_control_server.bootstrap.out_of_box_controls import ( OUT_OF_BOX_CONTROL_TEMPLATES, @@ -43,8 +44,12 @@ "oob-high-value-action-requires-approval", "oob-outbound-communication-requires-approval", "oob-only-approved-tools-may-run", + "oob-owasp-llm05-read-only-sql", + "oob-owasp-llm10-bounded-sql-query", + "oob-owasp-llm02-common-credential-output-match", + "oob-owasp-llm05-dangerous-uri-output-match", ) -_AVAILABLE_PHASE_2_EVALUATORS = {"regex", "json", "list"} +_AVAILABLE_PHASE_2_EVALUATORS = {"regex", "json", "list", "sql"} def _control_payload(*, evaluator_name: str = "regex") -> dict[str, object]: @@ -125,6 +130,13 @@ def test_out_of_box_catalog_contains_phase_2_templates() -> None: approved_tools_leaf = approved_tools.control.primary_leaf() assert approved_tools_leaf is not None assert approved_tools_leaf.selector.path == "canonical_name" + sql_controls = [ + template + for template in OUT_OF_BOX_CONTROL_TEMPLATES + if "sql" in template.control.tags + ] + assert len(sql_controls) == 2 + assert all(template.control.scope.step_name_regex for template in sql_controls) def test_missing_required_evaluators_returns_sorted_names() -> None: @@ -449,3 +461,73 @@ async def test_list_out_of_box_control_matches_unapproved_tools() -> None: assert delete_result.matched is True assert search_result.matched is False + + +@pytest.mark.asyncio +async def test_owasp_credential_control_matches_common_secret_formats() -> None: + # Given: the OWASP-aligned common credential output control + spec = _oob_evaluator_spec("oob-owasp-llm02-common-credential-output-match") + evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(spec.config)) + + # When: evaluating representative secret and non-secret output + private_key_result = await evaluator.evaluate("-----BEGIN OPENSSH PRIVATE KEY-----\nredacted") + aws_key_result = await evaluator.evaluate("Credential: AKIAIOSFODNN7EXAMPLE") + safe_result = await evaluator.evaluate("The operation completed successfully.") + + # Then: recognizable credentials are blocked while ordinary output passes + assert private_key_result.matched is True + assert aws_key_result.matched is True + assert safe_result.matched is False + + +@pytest.mark.asyncio +async def test_owasp_dangerous_uri_control_matches_active_content_schemes() -> None: + # Given: the OWASP-aligned dangerous URI output control + spec = _oob_evaluator_spec("oob-owasp-llm05-dangerous-uri-output-match") + evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(spec.config)) + + # When: evaluating executable, active-content, and ordinary HTTPS links + javascript_result = await evaluator.evaluate( + 'click' + ) + data_uri_result = await evaluator.evaluate("data:image/svg+xml,") + safe_result = await evaluator.evaluate("https://docs.example.com/safety") + + # Then: active-content schemes are blocked while HTTPS passes + assert javascript_result.matched is True + assert data_uri_result.matched is True + assert safe_result.matched is False + + +@pytest.mark.asyncio +async def test_owasp_read_only_sql_control_blocks_mutation_and_multiple_statements() -> None: + # Given: the OWASP-aligned read-only SQL control + spec = _oob_evaluator_spec("oob-owasp-llm05-read-only-sql") + evaluator = SQLEvaluator(SQLEvaluatorConfig.model_validate(spec.config)) + + # When: evaluating read-only, mutating, and multi-statement SQL + select_result = await evaluator.evaluate("SELECT id FROM users") + delete_result = await evaluator.evaluate("DELETE FROM users") + multiple_result = await evaluator.evaluate("SELECT id FROM users; DROP TABLE users") + + # Then: only the single read-only query passes + assert select_result.matched is False + assert delete_result.matched is True + assert multiple_result.matched is True + + +@pytest.mark.asyncio +async def test_owasp_bounded_sql_control_enforces_result_and_complexity_limits() -> None: + # Given: the OWASP-aligned bounded SQL query control + spec = _oob_evaluator_spec("oob-owasp-llm10-bounded-sql-query") + evaluator = SQLEvaluator(SQLEvaluatorConfig.model_validate(spec.config)) + + # When: evaluating bounded, unbounded, and oversized result windows + bounded_result = await evaluator.evaluate("SELECT id FROM users LIMIT 100") + missing_limit_result = await evaluator.evaluate("SELECT id FROM users") + oversized_window_result = await evaluator.evaluate("SELECT id FROM users LIMIT 1000 OFFSET 1") + + # Then: only the bounded query within the configured result window passes + assert bounded_result.matched is False + assert missing_limit_result.matched is True + assert oversized_window_result.matched is True diff --git a/server/tests/test_principal_namespace_flow.py b/server/tests/test_principal_namespace_flow.py index f02cf61b..295c6573 100644 --- a/server/tests/test_principal_namespace_flow.py +++ b/server/tests/test_principal_namespace_flow.py @@ -6,10 +6,6 @@ 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, @@ -17,6 +13,9 @@ ) from agent_control_server.bootstrap.out_of_box_controls import OUT_OF_BOX_CONTROL_TEMPLATES from agent_control_server.models import Control +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from sqlalchemy.orm import Session from .conftest import engine from .utils import VALID_CONTROL_PAYLOAD @@ -107,7 +106,10 @@ def test_controls_list_seeds_out_of_box_controls_for_principal_namespace( assert filtered.status_code == 200, filtered.text assert filtered.json()["controls"] == [] - resp = namespace_client.get("/api/v1/controls", params={"limit": 10}) + resp = namespace_client.get( + "/api/v1/controls", + params={"limit": len(OUT_OF_BOX_CONTROL_TEMPLATES)}, + ) assert resp.status_code == 200, resp.text expected_names = {template.name for template in OUT_OF_BOX_CONTROL_TEMPLATES} From a64020bfbc45199e968dc3d70425fc466f98a0e5 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 6 Aug 2026 14:23:27 -0700 Subject: [PATCH 12/18] address P2 comments --- ...d7e2b4_out_of_box_control_seed_identity.py | 18 +++++++------ .../bootstrap/out_of_box_controls.py | 16 +++++------ .../test_out_of_box_controls_bootstrap.py | 27 +++++++++++++++++++ 3 files changed, 45 insertions(+), 16 deletions(-) diff --git a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py index 63924de5..f976a2d9 100644 --- a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py +++ b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py @@ -24,16 +24,18 @@ def upgrade() -> None: "controls", sa.Column("seed_opted_out_at", sa.DateTime(timezone=True), nullable=True), ) - op.create_index( - "idx_controls_namespace_seed_source", - "controls", - ["namespace_key", "seed_source_id"], - unique=True, - postgresql_where=sa.text("seed_source_id IS NOT NULL"), - ) + with op.get_context().autocommit_block(): + op.execute( + """ + CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_controls_namespace_seed_source + ON controls (namespace_key, seed_source_id) + WHERE seed_source_id IS NOT NULL + """ + ) def downgrade() -> None: - op.drop_index("idx_controls_namespace_seed_source", table_name="controls") + with op.get_context().autocommit_block(): + op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_controls_namespace_seed_source") op.drop_column("controls", "seed_opted_out_at") op.drop_column("controls", "seed_source_id") 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..2368565d 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 @@ -50,14 +50,14 @@ class OutOfBoxControlTemplate: def __post_init__(self) -> None: object.__setattr__(self, "source_id", _SLUG_NAME_ADAPTER.validate_python(self.source_id)) object.__setattr__(self, "name", _SLUG_NAME_ADAPTER.validate_python(self.name)) - if not self.required_evaluators: - required_evaluators = { - evaluator.name for _, evaluator in self.control.iter_condition_leaf_parts() - } - object.__setattr__(self, "required_evaluators", frozenset(required_evaluators)) - return - - object.__setattr__(self, "required_evaluators", frozenset(self.required_evaluators)) + condition_evaluators = { + evaluator.name for _, evaluator in self.control.iter_condition_leaf_parts() + } + object.__setattr__( + self, + "required_evaluators", + frozenset(self.required_evaluators).union(condition_evaluators), + ) @classmethod def from_payload( diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 59135451..46386635 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -123,6 +123,33 @@ async def test_seed_skips_template_when_required_evaluator_is_missing() -> None: assert _fetch_controls() == [] +@pytest.mark.asyncio +async def test_seed_unions_explicit_requirements_with_condition_evaluators() -> None: + # Given: a regex control with an additional explicit Luna requirement + template = OutOfBoxControlTemplate.from_payload( + source_id="oob-mixed-requirements", + name="oob-mixed-requirements", + data=_control_payload(evaluator_name="regex"), + required_evaluators={"galileo.luna"}, + ) + + # When: seeding on a pod that only has Luna + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"galileo.luna"}, + templates=(template,), + ) + + # Then: the condition's missing regex evaluator prevents seeding + assert template.required_evaluators == frozenset({"galileo.luna", "regex"}) + assert result.created == () + assert len(result.skipped_missing_evaluator) == 1 + assert result.skipped_missing_evaluator[0].name == "oob-mixed-requirements" + assert result.skipped_missing_evaluator[0].missing_evaluators == ("regex",) + assert _fetch_controls() == [] + + @pytest.mark.asyncio async def test_seed_creates_control_version_in_namespace_without_bindings() -> None: template = _template(name="oob-create-control") From f94d14ccd1fcf57e4ff1d229906c607d4803e59c Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 6 Aug 2026 14:44:03 -0700 Subject: [PATCH 13/18] address comments --- ...d7e2b4_out_of_box_control_seed_identity.py | 32 ++++++ .../bootstrap/out_of_box_controls.py | 28 +++-- .../endpoints/controls.py | 56 ++++++++-- .../endpoints/evaluation.py | 16 +++ .../agent_control_server/services/controls.py | 31 ++++-- server/tests/test_controls_additional.py | 64 +++++++++++ .../test_data_model_v1_alembic_migration.py | 101 +++++++++++++++++- .../test_evaluation_legacy_canonical_name.py | 58 ++++++++++ .../test_out_of_box_controls_bootstrap.py | 86 ++++++++++++--- 9 files changed, 427 insertions(+), 45 deletions(-) create mode 100644 server/tests/test_evaluation_legacy_canonical_name.py diff --git a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py index f976a2d9..d9fb6d31 100644 --- a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py +++ b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py @@ -17,6 +17,8 @@ branch_labels = None depends_on = None +_CANONICAL_NAME_SEED_SOURCE_ID = "oob-only-approved-tools-may-run" + def upgrade() -> None: op.add_column("controls", sa.Column("seed_source_id", sa.String(length=255), nullable=True)) @@ -35,6 +37,36 @@ def upgrade() -> None: def downgrade() -> None: + # Older servers reject ``canonical_name`` selectors. Retire only the seeded + # control that still uses that selector, and make both its current payload + # and historical snapshots parseable before rolling the application back. + op.execute( + f""" + UPDATE control_versions AS version + SET snapshot = jsonb_set( + version.snapshot, + '{{data,condition,selector,path}}', + '"name"'::jsonb + ) + FROM controls AS control + WHERE version.control_id = control.id + AND control.seed_source_id = '{_CANONICAL_NAME_SEED_SOURCE_ID}' + AND version.snapshot #>> '{{data,condition,selector,path}}' = 'canonical_name' + """ + ) + op.execute( + f""" + UPDATE controls + SET data = jsonb_set( + data, + '{{condition,selector,path}}', + '"name"'::jsonb + ), + deleted_at = COALESCE(deleted_at, CURRENT_TIMESTAMP) + WHERE seed_source_id = '{_CANONICAL_NAME_SEED_SOURCE_ID}' + AND data #>> '{{condition,selector,path}}' = 'canonical_name' + """ + ) with op.get_context().autocommit_block(): op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_controls_namespace_seed_source") op.drop_column("controls", "seed_opted_out_at") 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 554840c1..d57ba3df 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 @@ -205,8 +205,10 @@ def _leaf_control_payload( evaluator_name="regex", evaluator_config={ "pattern": ( - r"(?:\brm\s+-rf\s+(?:/|~|\$HOME)(?:\s|[|;&]|$)|" - r"\bsudo\s+rm\s+-rf(?:\s|[|;&]|$)|" + r"(?:\brm\s+(?:-(?:rf|fr)|-r\s+-f|-f\s+-r)\s+" + r"(?:\"(?:/|~/?|\$HOME/?)\"|'(?:/|~/?|\$HOME/?)'|" + r"(?:/|~/?|\$HOME/?))(?:\s|[|;&]|$)|" + r"\bsudo\s+rm\s+(?:-(?:rf|fr)|-r\s+-f|-f\s+-r)(?:\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|[|;&]|$)|" @@ -442,6 +444,7 @@ async def seed_out_of_box_controls( available_evaluator_names = set(available_evaluators) async with session_factory() as session: + eligible_templates: list[OutOfBoxControlTemplate] = [] for template in templates: missing = missing_required_evaluators( template.required_evaluators, @@ -456,6 +459,19 @@ async def seed_out_of_box_controls( ) continue + eligible_templates.append(template) + + control_service = ControlService(session) + existing_source_ids, active_names = await control_service.find_existing_seed_controls( + namespace_key=namespace_key, + source_ids={template.source_id for template in eligible_templates}, + names={template.name for template in eligible_templates}, + ) + for template in eligible_templates: + if template.source_id in existing_source_ids or template.name in active_names: + skipped_existing.append(template.name) + continue + outcome = await _seed_one_control( session, namespace_key=namespace_key, @@ -483,14 +499,6 @@ async def _seed_one_control( template: OutOfBoxControlTemplate, ) -> str: control_service = ControlService(session) - if await control_service.seed_source_exists( - template.source_id, - namespace_key=namespace_key, - ): - return "existing" - if await control_service.active_control_name_exists(template.name, namespace_key=namespace_key): - return "existing" - control = control_service.create_control( namespace_key=namespace_key, name=template.name, diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index a2ff9f1a..dce876ad 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -1,6 +1,8 @@ +import asyncio import datetime as dt import uuid from copy import deepcopy +from functools import partial from typing import Any from agent_control_engine import list_evaluators @@ -95,6 +97,9 @@ _GENERATED_CLONE_NAME_ATTEMPTS = 5 _TRUE_QUERY_VALUES = {"1", "true", "t", "yes", "y", "on"} _SLUG_NAME_ADAPTER = TypeAdapter(SlugName) +_OUT_OF_BOX_RECONCILIATION_TIMEOUT_SECONDS = 3.0 +_MAX_PENDING_OUT_OF_BOX_RECONCILIATIONS = 128 +_out_of_box_reconciliation_tasks: dict[str, asyncio.Task[None]] = {} def _is_target_context_value(value: object) -> bool: @@ -258,16 +263,23 @@ def _validate_attachment_filters( ) -async def _seed_out_of_box_controls_for_namespace( +async def _run_out_of_box_controls_reconciliation( *, namespace_key: str, ) -> None: - """Best-effort idempotent namespace seeding for browse/list surfaces.""" + """Run one bounded, best-effort namespace reconciliation.""" try: - await seed_out_of_box_controls( - session_factory=AsyncSessionLocal, - namespace_key=namespace_key, - available_evaluators=set(list_evaluators().keys()), + async with asyncio.timeout(_OUT_OF_BOX_RECONCILIATION_TIMEOUT_SECONDS): + await seed_out_of_box_controls( + session_factory=AsyncSessionLocal, + namespace_key=namespace_key, + available_evaluators=set(list_evaluators().keys()), + ) + except TimeoutError: + _logger.warning( + "Out-of-box control reconciliation timed out for namespace '%s'; " + "continuing request", + namespace_key, ) except Exception: _logger.warning( @@ -277,6 +289,38 @@ async def _seed_out_of_box_controls_for_namespace( ) +def _remove_out_of_box_reconciliation_task( + namespace_key: str, + task: asyncio.Future[None], +) -> None: + if _out_of_box_reconciliation_tasks.get(namespace_key) is task: + _out_of_box_reconciliation_tasks.pop(namespace_key, None) + + +async def _seed_out_of_box_controls_for_namespace( + *, + namespace_key: str, +) -> None: + """Join or start the bounded reconciliation for a namespace.""" + task = _out_of_box_reconciliation_tasks.get(namespace_key) + if task is None: + if len(_out_of_box_reconciliation_tasks) >= _MAX_PENDING_OUT_OF_BOX_RECONCILIATIONS: + _logger.warning( + "Out-of-box control reconciliation queue is full; skipping namespace '%s'", + namespace_key, + ) + return + task = asyncio.create_task( + _run_out_of_box_controls_reconciliation(namespace_key=namespace_key) + ) + _out_of_box_reconciliation_tasks[namespace_key] = task + task.add_done_callback( + partial(_remove_out_of_box_reconciliation_task, namespace_key) + ) + + await asyncio.shield(task) + + def _should_seed_out_of_box_controls_on_list( *, cursor: int | None, diff --git a/server/src/agent_control_server/endpoints/evaluation.py b/server/src/agent_control_server/endpoints/evaluation.py index a31d757d..adc465a9 100644 --- a/server/src/agent_control_server/endpoints/evaluation.py +++ b/server/src/agent_control_server/endpoints/evaluation.py @@ -117,6 +117,21 @@ def _sanitize_evaluation_response(response: EvaluationResponse) -> EvaluationRes ) +def _normalize_legacy_qualified_step_name(request: EvaluationRequest) -> EvaluationRequest: + """Derive a canonical tool name for clients that predate ``canonical_name``.""" + step = request.step + if step.type != "tool" or step.canonical_name is not None: + return request + + _, separator, canonical_name = step.name.rpartition(".") + if not separator or not canonical_name: + return request + + return request.model_copy( + update={"step": step.model_copy(update={"canonical_name": canonical_name})} + ) + + async def _evaluation_context(request: Request) -> dict[str, object]: """Surface target identifiers to the runtime authorizer.""" try: @@ -196,6 +211,7 @@ async def evaluate( on the server; SDKs reconstruct and emit those events separately through the observability ingestion endpoint. """ + request = _normalize_legacy_qualified_step_name(request) engine_controls = await _load_engine_controls(request, principal) engine = ControlEngine(engine_controls) try: diff --git a/server/src/agent_control_server/services/controls.py b/server/src/agent_control_server/services/controls.py index 619cbed1..6af39721 100644 --- a/server/src/agent_control_server/services/controls.py +++ b/server/src/agent_control_server/services/controls.py @@ -1,7 +1,7 @@ from __future__ import annotations import datetime as dt -from collections.abc import Sequence +from collections.abc import Collection, Sequence from dataclasses import dataclass from typing import Any, Literal, cast @@ -163,19 +163,32 @@ def mark_control_deleted(control: Control, *, deleted_at: dt.datetime) -> None: if control.seed_source_id is not None: control.seed_opted_out_at = deleted_at - async def seed_source_exists( + async def find_existing_seed_controls( self, - seed_source_id: str, *, namespace_key: str, - ) -> bool: - """Return whether a control has claimed an immutable seed identity.""" - stmt = select(Control.id).where( + source_ids: Collection[str], + names: Collection[str], + ) -> tuple[frozenset[str], frozenset[str]]: + """Bulk-load claimed seed identities and conflicting active names.""" + if not source_ids and not names: + return frozenset(), frozenset() + + stmt = select(Control.seed_source_id, Control.name, Control.deleted_at).where( Control.namespace_key == namespace_key, - Control.seed_source_id == seed_source_id, + or_( + Control.seed_source_id.in_(source_ids), + Control.name.in_(names), + ), ) - result = await self._db.execute(stmt) - return result.first() is not None + rows = (await self._db.execute(stmt)).all() + existing_source_ids = frozenset( + source_id for source_id, _, _ in rows if source_id is not None + ) + active_names = frozenset( + name for _, name, deleted_at in rows if deleted_at is None + ) + return existing_source_ids, active_names async def get_control_or_404( self, diff --git a/server/tests/test_controls_additional.py b/server/tests/test_controls_additional.py index 59884f55..8cefc8f7 100644 --- a/server/tests/test_controls_additional.py +++ b/server/tests/test_controls_additional.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import uuid from collections.abc import AsyncGenerator @@ -542,6 +543,69 @@ async def test_seed_out_of_box_controls_failure_is_best_effort( seed.assert_awaited_once() +@pytest.mark.asyncio +async def test_seed_out_of_box_controls_deduplicates_concurrent_namespace_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given: one namespace reconciliation that remains in flight + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def seed(**_kwargs: object) -> None: + nonlocal calls + calls += 1 + started.set() + await release.wait() + + monkeypatch.setattr(controls_module, "seed_out_of_box_controls", seed) + + # When: two list requests reconcile the same namespace concurrently + first = asyncio.create_task( + controls_module._seed_out_of_box_controls_for_namespace( + namespace_key="shared-reconciliation-namespace" + ) + ) + await started.wait() + second = asyncio.create_task( + controls_module._seed_out_of_box_controls_for_namespace( + namespace_key="shared-reconciliation-namespace" + ) + ) + await asyncio.sleep(0) + release.set() + await asyncio.gather(first, second) + + # Then: both requests joined one reconciliation + assert calls == 1 + + +@pytest.mark.asyncio +async def test_seed_out_of_box_controls_bounds_reconciliation_wait( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given: a reconciliation that cannot complete within the read-path deadline + cancelled = asyncio.Event() + + async def seed(**_kwargs: object) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + monkeypatch.setattr(controls_module, "seed_out_of_box_controls", seed) + monkeypatch.setattr(controls_module, "_OUT_OF_BOX_RECONCILIATION_TIMEOUT_SECONDS", 0.01) + + # When: a list request starts reconciliation + await controls_module._seed_out_of_box_controls_for_namespace( + namespace_key="bounded-reconciliation-namespace" + ) + + # Then: the deadline cancels the database work and returns control to the request + assert cancelled.is_set() + + @pytest.mark.asyncio async def test_resolve_clone_name_reports_generated_name_exhaustion() -> None: control_service = MagicMock() diff --git a/server/tests/test_data_model_v1_alembic_migration.py b/server/tests/test_data_model_v1_alembic_migration.py index 53334732..00c27e56 100644 --- a/server/tests/test_data_model_v1_alembic_migration.py +++ b/server/tests/test_data_model_v1_alembic_migration.py @@ -2,22 +2,23 @@ from __future__ import annotations +import json import uuid from pathlib import Path import pytest +from agent_control_server.config import db_config +from alembic import command from alembic.config import Config from sqlalchemy import create_engine, inspect, text from sqlalchemy.engine import Engine, make_url -from agent_control_server.config import db_config -from alembic import command - SERVER_DIR = Path(__file__).resolve().parents[1] PRE_MIGRATION_REVISION = "c1e9f9c4a1d2" MIGRATION_REVISION = "a7f3b1e0d9c5" OBSERVABILITY_NAMESPACE_REVISION = "b6f4c2d8e9a1" CLONE_LINEAGE_REVISION = "e2b7f4a9c6d1" +SEED_IDENTITY_REVISION = "f3a1c8d7e2b4" _BASE_DB_URL = make_url(db_config.get_url()) pytestmark = pytest.mark.skipif( @@ -396,6 +397,100 @@ def test_control_clone_lineage_migration_adds_composite_fk_and_partial_index( assert "ix_events_agent_time" not in indexes +def test_seed_identity_downgrade_retires_canonical_name_control( + alembic_config: Config, + temp_engine: Engine, +) -> None: + # Given: the seeded control and its initial version use the new selector + command.upgrade(alembic_config, SEED_IDENTITY_REVISION) + control_data = { + "condition": { + "selector": {"path": "canonical_name"}, + "evaluator": {"name": "list", "config": {"values": ["web_search"]}}, + }, + "action": {"decision": "deny"}, + } + with temp_engine.begin() as conn: + control_id = conn.execute( + text( + """ + INSERT INTO controls (namespace_key, name, data, seed_source_id) + VALUES ( + 'default', + 'oob-only-approved-tools-may-run', + CAST(:data AS jsonb), + 'oob-only-approved-tools-may-run' + ) + RETURNING id + """ + ), + {"data": json.dumps(control_data)}, + ).scalar_one() + conn.execute( + text( + """ + INSERT INTO control_versions ( + control_id, version_num, event_type, snapshot, note + ) + VALUES ( + :control_id, + 1, + 'created', + CAST(:snapshot AS jsonb), + 'Out-of-box control seed' + ) + """ + ), + { + "control_id": control_id, + "snapshot": json.dumps( + { + "name": "oob-only-approved-tools-may-run", + "data": control_data, + } + ), + }, + ) + + # When: rolling back to the last revision that rejects canonical_name + command.downgrade(alembic_config, CLONE_LINEAGE_REVISION) + + # Then: the seed is inactive and all persisted definitions use a supported selector + with temp_engine.begin() as conn: + control = conn.execute( + text( + """ + SELECT + deleted_at, + data #>> '{condition,selector,path}' AS selector_path + FROM controls + WHERE id = :control_id + """ + ), + {"control_id": control_id}, + ).mappings().one() + snapshot_selector_path = conn.execute( + text( + """ + SELECT snapshot #>> '{data,condition,selector,path}' + FROM control_versions + WHERE control_id = :control_id + """ + ), + {"control_id": control_id}, + ).scalar_one() + + assert control["deleted_at"] is not None + assert control["selector_path"] == "name" + assert snapshot_selector_path == "name" + assert "seed_source_id" not in _column_names(temp_engine, "controls") + assert "seed_opted_out_at" not in _column_names(temp_engine, "controls") + assert "idx_controls_namespace_seed_source" not in _index_names( + temp_engine, + "controls", + ) + + def test_downgrade_rejects_cross_namespace_agents_duplicates( alembic_config: Config, temp_engine: Engine ) -> None: diff --git a/server/tests/test_evaluation_legacy_canonical_name.py b/server/tests/test_evaluation_legacy_canonical_name.py new file mode 100644 index 00000000..e6a7d6f4 --- /dev/null +++ b/server/tests/test_evaluation_legacy_canonical_name.py @@ -0,0 +1,58 @@ +"""Compatibility coverage for canonical tool names in evaluation requests.""" + +from agent_control_models import EvaluationRequest, Step +from fastapi.testclient import TestClient + +from .utils import create_and_assign_policy + + +def test_canonical_name_allowlist_supports_mixed_sdk_versions(client: TestClient) -> None: + # Given: an allowlist using the canonical tool identity sent by current SDKs + control_data = { + "description": "Allow web search", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["tool"], "stages": ["pre"]}, + "selector": {"path": "canonical_name"}, + "evaluator": { + "name": "list", + "config": { + "values": ["web_search"], + "logic": "any", + "match_on": "no_match", + "match_mode": "exact", + "case_sensitive": False, + }, + }, + "action": {"decision": "deny"}, + } + agent_name, _ = create_and_assign_policy( + client, + control_data, + agent_name="MixedVersionAgent", + ) + legacy_request = EvaluationRequest( + agent_name=agent_name, + step=Step(type="tool", name="writer.web_search", input={}), + stage="pre", + ) + current_request = EvaluationRequest( + agent_name=agent_name, + step=Step( + type="tool", + name="writer.web_search", + canonical_name="web_search", + input={}, + ), + stage="pre", + ) + + # When: legacy and current SDK payloads are evaluated by the same server + responses = [ + client.post("/api/v1/evaluation", json=request.model_dump(mode="json")) + for request in (legacy_request, current_request) + ] + + # Then: the approved tool is allowed for both client versions + assert [response.status_code for response in responses] == [200, 200] + assert [response.json()["is_safe"] for response in responses] == [True, True] diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 941b161c..da44201e 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -31,10 +31,10 @@ ) from agent_control_server.services.controls import ControlService from pydantic import ValidationError -from sqlalchemy import Table, func, select +from sqlalchemy import Table, event, func, select from sqlalchemy.orm import Session -from .conftest import AsyncSessionTest, engine +from .conftest import AsyncSessionTest, async_engine, engine _EXPECTED_OOB_CONTROL_NAMES = ( "oob-ssn-match", @@ -287,6 +287,43 @@ async def test_seed_default_catalog_is_idempotent() -> None: assert len(_fetch_versions()) == len(_EXPECTED_OOB_CONTROL_NAMES) +@pytest.mark.asyncio +async def test_seed_existing_catalog_uses_one_bulk_lookup() -> None: + # Given: a namespace whose complete catalog is already seeded + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators=_AVAILABLE_PHASE_2_EVALUATORS, + ) + statements: list[str] = [] + + def record_statement( + _conn: object, + _cursor: object, + statement: str, + _parameters: object, + _context: object, + _executemany: bool, + ) -> None: + statements.append(statement) + + event.listen(async_engine.sync_engine, "before_cursor_execute", record_statement) + try: + # When: reconciling the already-seeded namespace + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators=_AVAILABLE_PHASE_2_EVALUATORS, + ) + finally: + event.remove(async_engine.sync_engine, "before_cursor_execute", record_statement) + + # Then: all seed identities are resolved by one database statement + assert result.skipped_existing == _EXPECTED_OOB_CONTROL_NAMES + assert len(statements) == 1 + assert statements[0].lstrip().startswith("SELECT") + + @pytest.mark.asyncio async def test_seed_is_idempotent_for_existing_active_control_names() -> None: template = _template(name="oob-idempotent-control") @@ -379,25 +416,16 @@ async def test_seed_treats_duplicate_insert_integrity_error_as_skip( templates=(template,), ) - async def active_control_name_exists( - self: ControlService, - name: str, - *, - namespace_key: str, - exclude_control_id: int | None = None, - ) -> bool: - return False - - async def seed_source_exists( + async def find_existing_seed_controls( self: ControlService, - seed_source_id: str, *, namespace_key: str, - ) -> bool: - return False + source_ids: object, + names: object, + ) -> tuple[frozenset[str], frozenset[str]]: + return frozenset(), frozenset() - monkeypatch.setattr(ControlService, "active_control_name_exists", active_control_name_exists) - monkeypatch.setattr(ControlService, "seed_source_exists", seed_source_exists) + monkeypatch.setattr(ControlService, "find_existing_seed_controls", find_existing_seed_controls) result = await seed_out_of_box_controls( session_factory=AsyncSessionTest, @@ -433,6 +461,30 @@ async def test_regex_out_of_box_controls_match_representative_payloads() -> None assert shell_result.matched is True, command +@pytest.mark.asyncio +async def test_dangerous_shell_control_matches_equivalent_recursive_rm_forms() -> None: + # Given: the destructive shell command control + shell_spec = _oob_evaluator_spec("oob-dangerous-shell-command-match") + shell_evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(shell_spec.config)) + + # When: evaluating equivalent recursive deletion spellings and a scoped deletion + destructive_results = [ + await shell_evaluator.evaluate(command) + for command in ( + 'rm -rf "$HOME"', + "rm -rf ~/", + "rm -fr /", + "rm -r -f /", + "rm -f -r '$HOME/'", + ) + ] + scoped_result = await shell_evaluator.evaluate("rm -rf /tmp/build-output") + + # Then: equivalent root/home deletions are blocked without blocking scoped deletion + assert all(result.matched is True for result in destructive_results) + assert scoped_result.matched is False + + @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") From ebd68dd3e6dfc146cc46ca54012558bd5a92eb24 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 10 Aug 2026 14:58:39 -0700 Subject: [PATCH 14/18] address comments --- ...4a9b2c6f1_out_of_box_control_seed_index.py | 53 +++++++++ ...d7e2b4_out_of_box_control_seed_identity.py | 11 +- .../test_data_model_v1_alembic_migration.py | 106 ++++++++++++++++++ 3 files changed, 160 insertions(+), 10 deletions(-) create mode 100644 server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py diff --git a/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py b/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py new file mode 100644 index 00000000..c4b61f23 --- /dev/null +++ b/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py @@ -0,0 +1,53 @@ +"""add out-of-box control seed index + +Revision ID: d7e4a9b2c6f1 +Revises: f3a1c8d7e2b4 +Create Date: 2026-08-10 12:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "d7e4a9b2c6f1" +down_revision = "f3a1c8d7e2b4" +branch_labels = None +depends_on = None + +_INDEX_NAME = "idx_controls_namespace_seed_source" + + +def _index_is_invalid() -> bool: + result = op.get_bind().execute( + sa.text( + """ + SELECT NOT pg_index.indisvalid + FROM pg_index + WHERE pg_index.indexrelid = to_regclass(:index_name) + """ + ), + {"index_name": _INDEX_NAME}, + ) + return bool(result.scalar_one_or_none()) + + +def upgrade() -> None: + with op.get_context().autocommit_block(): + if _index_is_invalid(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") + op.execute( + f""" + CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} + ON controls (namespace_key, seed_source_id) + WHERE seed_source_id IS NOT NULL + """ + ) + + +def downgrade() -> None: + with op.get_context().autocommit_block(): + op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") diff --git a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py index f976a2d9..8e31e14c 100644 --- a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py +++ b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py @@ -9,6 +9,7 @@ from __future__ import annotations import sqlalchemy as sa + from alembic import op # revision identifiers, used by Alembic. @@ -24,18 +25,8 @@ def upgrade() -> None: "controls", sa.Column("seed_opted_out_at", sa.DateTime(timezone=True), nullable=True), ) - with op.get_context().autocommit_block(): - op.execute( - """ - CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_controls_namespace_seed_source - ON controls (namespace_key, seed_source_id) - WHERE seed_source_id IS NOT NULL - """ - ) def downgrade() -> None: - with op.get_context().autocommit_block(): - op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_controls_namespace_seed_source") op.drop_column("controls", "seed_opted_out_at") op.drop_column("controls", "seed_source_id") diff --git a/server/tests/test_data_model_v1_alembic_migration.py b/server/tests/test_data_model_v1_alembic_migration.py index 53334732..78588b41 100644 --- a/server/tests/test_data_model_v1_alembic_migration.py +++ b/server/tests/test_data_model_v1_alembic_migration.py @@ -18,6 +18,8 @@ MIGRATION_REVISION = "a7f3b1e0d9c5" OBSERVABILITY_NAMESPACE_REVISION = "b6f4c2d8e9a1" CLONE_LINEAGE_REVISION = "e2b7f4a9c6d1" +SEED_IDENTITY_REVISION = "f3a1c8d7e2b4" +SEED_INDEX_REVISION = "d7e4a9b2c6f1" _BASE_DB_URL = make_url(db_config.get_url()) pytestmark = pytest.mark.skipif( @@ -120,6 +122,22 @@ def _pg_index_definition(engine: Engine, index_name: str) -> str: ) +def _pg_index_is_valid(engine: Engine, index_name: str) -> bool: + with engine.begin() as conn: + return bool( + conn.execute( + text( + """ + SELECT pg_index.indisvalid + FROM pg_index + WHERE pg_index.indexrelid = to_regclass(:index_name) + """ + ), + {"index_name": index_name}, + ).scalar_one() + ) + + def _pg_constraint_definition(engine: Engine, constraint_name: str) -> tuple[str, str]: with engine.begin() as conn: row = conn.execute( @@ -396,6 +414,94 @@ def test_control_clone_lineage_migration_adds_composite_fk_and_partial_index( assert "ix_events_agent_time" not in indexes +def test_control_seed_index_migration_is_split_from_column_additions( + alembic_config: Config, temp_engine: Engine +) -> None: + # Given: the identity revision has added columns without creating the index + command.upgrade(alembic_config, SEED_IDENTITY_REVISION) + assert "seed_source_id" in _column_names(temp_engine, "controls") + assert "seed_opted_out_at" in _column_names(temp_engine, "controls") + assert "idx_controls_namespace_seed_source" not in _index_names( + temp_engine, "controls" + ) + + # When: the separate concurrent index revision is applied + command.upgrade(alembic_config, SEED_INDEX_REVISION) + + # Then: the partial unique index is valid + assert "idx_controls_namespace_seed_source" in _index_names(temp_engine, "controls") + assert _pg_index_is_valid(temp_engine, "idx_controls_namespace_seed_source") + index_def = _pg_index_definition(temp_engine, "idx_controls_namespace_seed_source") + assert "CREATE UNIQUE INDEX idx_controls_namespace_seed_source" in index_def + assert "ON public.controls USING btree (namespace_key, seed_source_id)" in index_def + assert "WHERE (seed_source_id IS NOT NULL)" in index_def + + +def test_control_seed_index_downgrade_preserves_identity_columns( + alembic_config: Config, temp_engine: Engine +) -> None: + # Given: both seed identity revisions have been applied + command.upgrade(alembic_config, SEED_INDEX_REVISION) + + # When: downgrading only the concurrent index revision + command.downgrade(alembic_config, SEED_IDENTITY_REVISION) + + # Then: the index is removed while the identity columns remain + assert "idx_controls_namespace_seed_source" not in _index_names( + temp_engine, "controls" + ) + assert "seed_source_id" in _column_names(temp_engine, "controls") + assert "seed_opted_out_at" in _column_names(temp_engine, "controls") + + +def test_control_seed_index_migration_rebuilds_an_invalid_existing_index( + alembic_config: Config, temp_engine: Engine +) -> None: + # Given: a failed concurrent build left an invalid same-name index + command.upgrade(alembic_config, SEED_IDENTITY_REVISION) + with temp_engine.begin() as conn: + conn.execute( + text( + """ + INSERT INTO controls (namespace_key, name, data, seed_source_id) + VALUES + ('default', 'seed-one', '{}'::jsonb, 'duplicate-seed'), + ('default', 'seed-two', '{}'::jsonb, 'duplicate-seed') + """ + ) + ) + + with pytest.raises(Exception): + with temp_engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: + conn.execute( + text( + """ + CREATE UNIQUE INDEX CONCURRENTLY idx_controls_namespace_seed_source + ON controls (namespace_key, seed_source_id) + WHERE seed_source_id IS NOT NULL + """ + ) + ) + + assert not _pg_index_is_valid(temp_engine, "idx_controls_namespace_seed_source") + with temp_engine.begin() as conn: + conn.execute( + text( + """ + UPDATE controls + SET seed_source_id = 'distinct-seed' + WHERE name = 'seed-two' + """ + ) + ) + + # When: the index revision is retried + command.upgrade(alembic_config, SEED_INDEX_REVISION) + + # Then: it replaces the invalid index with a valid unique index + assert _pg_index_is_valid(temp_engine, "idx_controls_namespace_seed_source") + + def test_downgrade_rejects_cross_namespace_agents_duplicates( alembic_config: Config, temp_engine: Engine ) -> None: From 46362d7a8cc31fb97eaf08f8255de73aa8bee8c4 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 10 Aug 2026 17:08:59 -0700 Subject: [PATCH 15/18] coverage --- .../test_out_of_box_controls_bootstrap.py | 133 ++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 46386635..f96b6fd7 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -3,11 +3,14 @@ import datetime as dt import uuid from copy import deepcopy +from types import SimpleNamespace from typing import cast import pytest +from agent_control_server.bootstrap import out_of_box_controls as bootstrap_module from agent_control_server.bootstrap.out_of_box_controls import ( OutOfBoxControlTemplate, + OutOfBoxSeedResult, default_out_of_box_namespace_key, missing_required_evaluators, seed_out_of_box_controls, @@ -23,6 +26,7 @@ from agent_control_server.services.controls import ControlService from pydantic import ValidationError from sqlalchemy import Table, func, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from .conftest import AsyncSessionTest, engine @@ -88,6 +92,25 @@ def test_missing_required_evaluators_returns_sorted_names() -> None: assert missing == ("galileo.luna", "regex") +def test_seed_result_reports_created_and_skipped_counts() -> None: + # Given: a result containing every skip category + result = OutOfBoxSeedResult( + created=("created-one", "created-two"), + skipped_existing=("existing",), + skipped_missing_evaluator=( + bootstrap_module.SkippedOutOfBoxControl( + name="missing", + missing_evaluators=("regex",), + ), + ), + skipped_conflict=("conflict",), + ) + + # When/Then: its summary counts include the corresponding entries + assert result.created_count == 2 + assert result.skipped_count == 3 + + def test_template_from_payload_validates_control_definition() -> None: payload = deepcopy(_control_payload()) payload["condition"] = { @@ -103,6 +126,24 @@ def test_template_from_payload_validates_control_definition() -> None: ) +@pytest.mark.asyncio +async def test_seed_empty_catalog_returns_without_opening_session() -> None: + # Given: an empty catalog and a session factory that must not be used + def unexpected_session_factory() -> None: + raise AssertionError("empty catalog should not open a database session") + + # When: bootstrap runs with no templates + result = await seed_out_of_box_controls( + session_factory=unexpected_session_factory, # type: ignore[arg-type] + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(), + ) + + # Then: it returns an empty result without touching the database + assert result == OutOfBoxSeedResult() + + @pytest.mark.asyncio async def test_seed_skips_template_when_required_evaluator_is_missing() -> None: template = _template(name="oob-missing-evaluator") @@ -209,6 +250,55 @@ async def test_seed_is_idempotent_for_existing_active_control_names() -> None: assert len(_fetch_versions()) == 1 +@pytest.mark.asyncio +async def test_seed_skips_active_name_claimed_by_another_source() -> None: + # Given: an existing seeded control and a new template reusing its active name + original = _template(name="oob-shared-name", source_id="original-source") + replacement = _template(name="oob-shared-name", source_id="replacement-source") + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(original,), + ) + + # When: bootstrap evaluates the replacement template + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(replacement,), + ) + + # Then: the active name prevents a duplicate with a different seed identity + assert result.skipped_existing == ("oob-shared-name",) + assert len(_fetch_controls()) == 1 + + +@pytest.mark.asyncio +async def test_seed_serializes_default_enabled_value() -> None: + # Given: a template payload that relies on the model's enabled default + payload = _control_payload() + payload.pop("enabled") + template = OutOfBoxControlTemplate.from_payload( + source_id="oob-default-enabled", + name="oob-default-enabled", + data=payload, + ) + + # When: the template is seeded + result = await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + # Then: the stored payload explicitly contains the effective default + assert result.created == ("oob-default-enabled",) + assert _fetch_controls()[0].data["enabled"] is True + + @pytest.mark.asyncio async def test_seed_does_not_duplicate_a_renamed_seeded_control() -> None: template = _template(name="oob-original-name", source_id="stable-seed-id") @@ -309,3 +399,46 @@ 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_seed_reraises_unrelated_integrity_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given: version creation fails for a reason unrelated to seed uniqueness + template = _template(name="oob-unrelated-integrity-error") + error = IntegrityError("statement", {}, RuntimeError("unrelated constraint")) + + async def raise_integrity_error( + self: ControlService, + control: Control, + *, + event_type: str, + note: str | None = None, + ) -> None: + raise error + + monkeypatch.setattr(ControlService, "create_version", raise_integrity_error) + + # When/Then: bootstrap rolls back and propagates the unexpected failure + with pytest.raises(IntegrityError) as exc_info: + await seed_out_of_box_controls( + session_factory=AsyncSessionTest, + namespace_key=DEFAULT_NAMESPACE_KEY, + available_evaluators={"regex"}, + templates=(template,), + ) + + assert exc_info.value is error + assert _fetch_controls() == [] + + +def test_seed_conflict_recognizes_seed_constraint_diagnostic() -> None: + # Given: PostgreSQL reports the immutable seed index through its diagnostic + original = SimpleNamespace( + diag=SimpleNamespace(constraint_name="idx_controls_namespace_seed_source") + ) + error = IntegrityError("statement", {}, original) + + # When/Then: the pure classifier recognizes the race as a seed conflict + assert bootstrap_module._is_control_seed_conflict(error) is True From 1b3dfe7b90e98722fd915598f8aabda99e7b033e Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Tue, 11 Aug 2026 09:49:01 -0700 Subject: [PATCH 16/18] more comments --- ...4a9b2c6f1_out_of_box_control_seed_index.py | 8 +++--- .../test_data_model_v1_alembic_migration.py | 27 ++++++++++++++++--- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py b/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py index c4b61f23..8c7c3c8c 100644 --- a/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py +++ b/server/alembic/versions/d7e4a9b2c6f1_out_of_box_control_seed_index.py @@ -9,8 +9,7 @@ from __future__ import annotations import sqlalchemy as sa - -from alembic import op +from alembic import context, op # revision identifiers, used by Alembic. revision = "d7e4a9b2c6f1" @@ -37,7 +36,10 @@ def _index_is_invalid() -> bool: def upgrade() -> None: with op.get_context().autocommit_block(): - if _index_is_invalid(): + # Offline generation has no live PostgreSQL catalog to inspect. Emitting + # CREATE IF NOT EXISTS is safe there; invalid-index recovery remains an + # online-only retry path. + if not context.is_offline_mode() and _index_is_invalid(): op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}") op.execute( f""" diff --git a/server/tests/test_data_model_v1_alembic_migration.py b/server/tests/test_data_model_v1_alembic_migration.py index 78588b41..5d8175c4 100644 --- a/server/tests/test_data_model_v1_alembic_migration.py +++ b/server/tests/test_data_model_v1_alembic_migration.py @@ -2,17 +2,17 @@ from __future__ import annotations +import io import uuid from pathlib import Path import pytest +from agent_control_server.config import db_config +from alembic import command from alembic.config import Config from sqlalchemy import create_engine, inspect, text from sqlalchemy.engine import Engine, make_url -from agent_control_server.config import db_config -from alembic import command - SERVER_DIR = Path(__file__).resolve().parents[1] PRE_MIGRATION_REVISION = "c1e9f9c4a1d2" MIGRATION_REVISION = "a7f3b1e0d9c5" @@ -437,6 +437,27 @@ def test_control_seed_index_migration_is_split_from_column_additions( assert "WHERE (seed_source_id IS NOT NULL)" in index_def +def test_control_seed_index_migration_supports_offline_sql_generation() -> None: + # Given: an offline Alembic configuration with no live database connection + output = io.StringIO() + config = Config(str(SERVER_DIR / "alembic.ini"), output_buffer=output) + config.set_main_option("script_location", str(SERVER_DIR / "alembic")) + config.set_main_option("sqlalchemy.url", _BASE_DB_URL.render_as_string(False)) + + # When: generating SQL for only the concurrent index revision + command.upgrade( + config, + f"{SEED_IDENTITY_REVISION}:{SEED_INDEX_REVISION}", + sql=True, + ) + + # Then: Alembic emits the index DDL without trying to query pg_index + generated_sql = output.getvalue() + assert "CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS" in generated_sql + assert "idx_controls_namespace_seed_source" in generated_sql + assert "SELECT NOT pg_index.indisvalid" not in generated_sql + + def test_control_seed_index_downgrade_preserves_identity_columns( alembic_config: Config, temp_engine: Engine ) -> None: From 6d4da85148457979ef2603b2b79afc38dc0ccfce Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 13 Aug 2026 13:21:10 -0700 Subject: [PATCH 17/18] claude review --- ...d7e2b4_out_of_box_control_seed_identity.py | 2 +- .../bootstrap/out_of_box_controls.py | 34 ++++++++++++--- .../test_data_model_v1_alembic_migration.py | 6 +-- .../test_out_of_box_controls_bootstrap.py | 43 +++++++++++++++++-- 4 files changed, 72 insertions(+), 13 deletions(-) diff --git a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py index dca2ac9b..5ee5a913 100644 --- a/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py +++ b/server/alembic/versions/f3a1c8d7e2b4_out_of_box_control_seed_identity.py @@ -18,7 +18,7 @@ branch_labels = None depends_on = None -_CANONICAL_NAME_SEED_SOURCE_ID = "oob-only-approved-tools-may-run" +_CANONICAL_NAME_SEED_SOURCE_ID = "oob-example-tool-allowlist" def upgrade() -> None: 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 e72752e3..1717722d 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 @@ -37,6 +37,20 @@ r"(?i)(?:^|[._-])(?:sql|execute[_-]?sql|run[_-]?sql|sql[_-]?query|" r"query[_-]?database|execute[_-]?query)(?:$|[._-])" ) +# Issuer prefix + exact digit length per network, allowing the conventional +# grouping separators (space or dash) so both "4111111111111111" and +# "4111 1111 1111 1111" match. A bare digit-count check (e.g. `\d{13,19}`) +# would also match order numbers, tracking numbers, and other unrelated +# identifiers, so each branch is anchored to a real issuer prefix and length. +_CREDIT_CARD_PATTERN = ( + r"\b(?:" + r"4\d{3}(?:[ -]?\d{4}){3}" # Visa (16 digits) + r"|5[1-5]\d{2}(?:[ -]?\d{4}){3}" # Mastercard 51-55 (16 digits) + r"|(?:2221|222[2-9]|22[3-9]\d|2[3-6]\d{2}|27[01]\d|2720)(?:[ -]?\d{4}){3}" # Mastercard 2221-2720 + r"|3[47]\d{2}(?:[ -]?\d{6})(?:[ -]?\d{5})" # American Express (15 digits, 4-6-5) + r"|6(?:011|5\d{2})(?:[ -]?\d{4}){3}" # Discover (16 digits) + r")\b" +) @dataclass(frozen=True, slots=True) @@ -167,10 +181,13 @@ def _leaf_control_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.", + description=( + "Block LLM output containing Visa, Mastercard, American Express, or " + "Discover card numbers (issuer prefix and length, formatted or unformatted)." + ), selector_path="output", evaluator_name="regex", - evaluator_config={"pattern": r"\b(?:\d[ -]?){13,19}\b"}, + evaluator_config={"pattern": _CREDIT_CARD_PATTERN}, step_types=["llm"], stages=["post"], decision="deny", @@ -299,10 +316,15 @@ def _leaf_control_payload( ), ), OutOfBoxControlTemplate.from_payload( - source_id="oob-only-approved-tools-may-run", - name="oob-only-approved-tools-may-run", + source_id="oob-example-tool-allowlist", + name="oob-example-tool-allowlist", data=_leaf_control_payload( - description="Deny tool calls whose step name is not in the approved tool list.", + description=( + "Example static allowlist: deny tool calls whose canonical name is not " + "in this list. The values below are placeholders, not a live MCP tool " + "catalog — replace them with your own approved tool names before " + "enabling this control." + ), selector_path="canonical_name", evaluator_name="list", evaluator_config={ @@ -315,7 +337,7 @@ def _leaf_control_payload( step_types=["tool"], stages=["pre"], decision="deny", - tags=["tool", "allowlist", "list"], + tags=["tool", "allowlist", "example", "list"], ), ), OutOfBoxControlTemplate.from_payload( diff --git a/server/tests/test_data_model_v1_alembic_migration.py b/server/tests/test_data_model_v1_alembic_migration.py index 1c041180..14d3a5ca 100644 --- a/server/tests/test_data_model_v1_alembic_migration.py +++ b/server/tests/test_data_model_v1_alembic_migration.py @@ -435,9 +435,9 @@ def test_seed_identity_downgrade_retires_canonical_name_control( INSERT INTO controls (namespace_key, name, data, seed_source_id) VALUES ( 'default', - 'oob-only-approved-tools-may-run', + 'oob-example-tool-allowlist', CAST(:data AS jsonb), - 'oob-only-approved-tools-may-run' + 'oob-example-tool-allowlist' ) RETURNING id """ @@ -463,7 +463,7 @@ def test_seed_identity_downgrade_retires_canonical_name_control( "control_id": control_id, "snapshot": json.dumps( { - "name": "oob-only-approved-tools-may-run", + "name": "oob-example-tool-allowlist", "data": control_data, } ), diff --git a/server/tests/test_out_of_box_controls_bootstrap.py b/server/tests/test_out_of_box_controls_bootstrap.py index 25d31c26..74d2a368 100644 --- a/server/tests/test_out_of_box_controls_bootstrap.py +++ b/server/tests/test_out_of_box_controls_bootstrap.py @@ -47,7 +47,7 @@ "oob-dangerous-shell-command-match", "oob-high-value-action-requires-approval", "oob-outbound-communication-requires-approval", - "oob-only-approved-tools-may-run", + "oob-example-tool-allowlist", "oob-owasp-llm05-select-only-sql", "oob-owasp-llm10-bounded-sql-query", "oob-owasp-llm02-common-credential-output-match", @@ -129,7 +129,7 @@ def test_out_of_box_catalog_contains_phase_2_templates() -> None: approved_tools = next( template for template in OUT_OF_BOX_CONTROL_TEMPLATES - if template.name == "oob-only-approved-tools-may-run" + if template.name == "oob-example-tool-allowlist" ) approved_tools_leaf = approved_tools.control.primary_leaf() assert approved_tools_leaf is not None @@ -600,6 +600,43 @@ async def test_regex_out_of_box_controls_match_representative_payloads() -> None assert shell_result.matched is True, command +@pytest.mark.asyncio +async def test_credit_card_control_matches_known_networks_and_ignores_generic_digit_runs() -> None: + # Given: the credit-card-number-match control + spec = _oob_evaluator_spec("oob-credit-card-number-match") + evaluator = RegexEvaluator(RegexEvaluatorConfig.model_validate(spec.config)) + + # When: evaluating real card numbers from each supported network, formatted + # and unformatted, alongside unrelated numbers with a similar digit count + network_results = [ + await evaluator.evaluate(number) + for number in ( + "4111 1111 1111 1111", # Visa + "4111111111111111", # Visa, unformatted + "4111-1111-1111-1111", # Visa, dash-separated + "5500 0000 0000 0004", # Mastercard (51-55 range) + "2223 0000 4841 0010", # Mastercard (2221-2720 range) + "3782 822463 10005", # American Express + "378282246310005", # American Express, unformatted + "6011 0000 0000 0004", # Discover + ) + ] + generic_digit_results = [ + await evaluator.evaluate(text) + for text in ( + "Your order 1234567890123456 has shipped", + "Invoice #: 987654321098765", + "Tracking: 19999999999999999", + "Account number: 12345678901234", + ) + ] + + # Then: only genuine card-shaped numbers are blocked; other long digit + # runs (order/invoice/tracking/account numbers) are not false positives + assert all(result.matched is True for result in network_results) + assert all(result.matched is False for result in generic_digit_results) + + @pytest.mark.asyncio async def test_dangerous_shell_control_matches_equivalent_recursive_rm_forms() -> None: # Given: the destructive shell command control @@ -674,7 +711,7 @@ async def test_json_out_of_box_controls_ignore_caller_controlled_approval_flags( @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_spec = _oob_evaluator_spec("oob-example-tool-allowlist") tool_evaluator = ListEvaluator(ListEvaluatorConfig.model_validate(tool_spec.config)) delete_result = await tool_evaluator.evaluate("delete_user") From b30910e38ab8d619bdcb375fd2334843ae2b3f12 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 13 Aug 2026 13:32:39 -0700 Subject: [PATCH 18/18] ruff --- .../src/agent_control_server/bootstrap/out_of_box_controls.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 1717722d..e63d4431 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 @@ -46,7 +46,8 @@ r"\b(?:" r"4\d{3}(?:[ -]?\d{4}){3}" # Visa (16 digits) r"|5[1-5]\d{2}(?:[ -]?\d{4}){3}" # Mastercard 51-55 (16 digits) - r"|(?:2221|222[2-9]|22[3-9]\d|2[3-6]\d{2}|27[01]\d|2720)(?:[ -]?\d{4}){3}" # Mastercard 2221-2720 + r"|(?:2221|222[2-9]|22[3-9]\d|2[3-6]\d{2}|27[01]\d|2720)" # Mastercard 2221-2720 + r"(?:[ -]?\d{4}){3}" r"|3[47]\d{2}(?:[ -]?\d{6})(?:[ -]?\d{5})" # American Express (15 digits, 4-6-5) r"|6(?:011|5\d{2})(?:[ -]?\d{4}){3}" # Discover (16 digits) r")\b"