From c99a3850c39bdb9a564df45cab00f580ae3b0090 Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Mon, 29 Jun 2026 12:18:51 -0700 Subject: [PATCH] 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 +