diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py index 75a3e5b6..ce6dcb07 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/records/factory.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Mapping -from typing import Any +from typing import Any, cast from agent_control_models import Step from galileo_core.schemas.logging.session import Session @@ -156,21 +156,21 @@ def record_from_step( **common, input=GalileoRecordNormalizer.tool_input(input_value), output=GalileoRecordNormalizer.tool_output(output_value), - spans=_span_children(child_records, parent_type=record_type), + spans=cast(Any, _span_children(child_records, parent_type=record_type)), ) if record_type == "retriever": return RetrieverSpan( **common, input=GalileoRecordNormalizer.retriever_input(input_value), output=GalileoRecordNormalizer.retriever_output(output_value), - spans=_span_children(child_records, parent_type=record_type), + spans=cast(Any, _span_children(child_records, parent_type=record_type)), ) if record_type == "trace": return Trace( **common, input=GalileoRecordNormalizer.trace_input(input_value), output=GalileoRecordNormalizer.trace_output(output_value), - spans=_span_children(child_records, parent_type=record_type), + spans=cast(Any, _span_children(child_records, parent_type=record_type)), ) traces = [child for child in child_records if isinstance(child, Trace)] if len(traces) != len(child_records): diff --git a/server/pyproject.toml b/server/pyproject.toml index 82b5a2af..9f2b7304 100644 --- a/server/pyproject.toml +++ b/server/pyproject.toml @@ -34,6 +34,7 @@ license = {text = "Apache-2.0"} [project.optional-dependencies] galileo = ["agent-control-evaluator-galileo>=7.5.0"] +redis = ["redis>=5.2.1"] binary = ["psycopg[binary]>=3.1"] c = ["psycopg[c]>=3.1"] diff --git a/server/src/agent_control_server/config.py b/server/src/agent_control_server/config.py index 00335611..31d9418e 100644 --- a/server/src/agent_control_server/config.py +++ b/server/src/agent_control_server/config.py @@ -3,7 +3,7 @@ import logging import secrets from functools import cached_property -from typing import Any +from typing import Any, Literal from agent_control_models import JSONObject from agent_control_telemetry import DEFAULT_CONTROL_EVENT_SINK_NAME @@ -191,6 +191,21 @@ class Settings(BaseSettings): ), ) + control_cache_backend: Literal["postgres", "redis"] = Field( + default="postgres", validation_alias=AliasChoices("AGENT_CONTROL_CONTROL_CACHE_BACKEND") + ) + redis_url: str | None = Field( + default=None, validation_alias=AliasChoices("AGENT_CONTROL_REDIS_URL") + ) + redis_mode: Literal["standalone", "cluster"] = Field( + default="standalone", validation_alias=AliasChoices("AGENT_CONTROL_REDIS_MODE") + ) + control_cache_ttl_seconds: int = Field( + default=60, + ge=1, + validation_alias=AliasChoices("AGENT_CONTROL_CONTROL_CACHE_TTL_SECONDS"), + ) + # Prometheus metrics settings prometheus_metrics_prefix: str = _env_alias_field( "agent_control_server", diff --git a/server/src/agent_control_server/control_cache.py b/server/src/agent_control_server/control_cache.py new file mode 100644 index 00000000..1825b92b --- /dev/null +++ b/server/src/agent_control_server/control_cache.py @@ -0,0 +1,483 @@ +"""Provider-neutral cache for serialized effective-control snapshots.""" + +from __future__ import annotations + +import base64 +import builtins +import json +import logging +import time +from collections.abc import Sequence +from datetime import UTC, datetime +from functools import lru_cache +from typing import Any, Protocol, cast + +from agent_control_models import ControlDefinitionRuntime +from prometheus_client import Counter, Histogram +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from .config import settings +from .models import Agent, ControlBinding, agent_controls, agent_policies, policy_controls +from .services.controls import RuntimeControl + +_logger = logging.getLogger(__name__) +_CACHE_HIT = Counter("agent_control_control_cache_hit_total", "Effective control cache hits") +_CACHE_MISS = Counter("agent_control_control_cache_miss_total", "Effective control cache misses") +_CACHE_FAILURE = Counter("agent_control_control_cache_redis_failure_total", "Redis cache failures") +_CACHE_INVALIDATION = Counter( + "agent_control_control_cache_invalidation_total", "Effective control cache invalidations" +) +_CACHE_INVALIDATION_FAILURE = Counter( + "agent_control_control_cache_invalidation_failure_total", + "Effective control cache invalidation failures", +) +_SNAPSHOT_LOAD = Histogram( + "agent_control_control_cache_snapshot_load_seconds", "Effective control snapshot load time" +) +_SNAPSHOT_SIZE = Histogram( + "agent_control_control_cache_snapshot_size_bytes", "Serialized snapshot size" +) + + +class ControlCache(Protocol): + """Storage interface for effective-control snapshots.""" + + async def get( + self, namespace: str, agent: str, target_type: str | None, target_id: str | None + ) -> list[RuntimeControl] | None: ... + + async def set( + self, + namespace: str, + agent: str, + target_type: str | None, + target_id: str | None, + controls: Sequence[RuntimeControl], + generation: str | None, + ) -> None: ... + + async def get_generation(self, namespace: str, agent: str) -> str | None: ... + + async def invalidate_agent(self, namespace: str, agent: str) -> None: ... + + +class NoOpControlCache: + """Cache implementation that always uses PostgreSQL.""" + + async def get( + self, namespace: str, agent: str, target_type: str | None, target_id: str | None + ) -> None: + del namespace, agent, target_type, target_id + _CACHE_MISS.inc() + return None + + async def set( + self, + namespace: str, + agent: str, + target_type: str | None, + target_id: str | None, + controls: Sequence[RuntimeControl], + generation: str | None = None, + ) -> None: + del namespace, agent, target_type, target_id, controls, generation + + async def get_generation(self, namespace: str, agent: str) -> None: + del namespace, agent + return None + + async def invalidate_agent(self, namespace: str, agent: str) -> None: + del namespace, agent + + +class RedisCommands(Protocol): + """Minimal asynchronous Redis command surface used by this cache.""" + + async def get(self, name: str) -> str | bytes | None: ... + + async def set(self, name: str, value: str, ex: int) -> object: ... + + async def sadd(self, name: str, *values: str) -> object: ... + + async def expire(self, name: str, seconds: int) -> object: ... + + async def smembers(self, name: str) -> builtins.set[str | bytes]: ... + + async def delete(self, *names: str) -> object: ... + + async def eval( + self, script: str, numkeys: int, *keys_and_args: str | int + ) -> str | bytes | int | None: ... + + async def aclose(self) -> None: ... + + +def build_cache_key( + namespace: str, agent: str, target_type: str | None, target_id: str | None +) -> str: + """Build the versioned snapshot key, preserving all effective-control scope dimensions.""" + return ( + f"agent-control:effective-controls:v1:{{{_agent_hash_tag(namespace, agent)}}}:" + f"{_encode_key_component(target_type)}:{_encode_key_component(target_id)}" + ) + + +def _encode_key_component(value: str | None) -> str: + if value is None: + return "!" + return "=" + base64.urlsafe_b64encode(value.encode()).decode().rstrip("=") + + +def _agent_index_key(namespace: str, agent: str) -> str: + return f"agent-control:effective-controls:index:v1:{{{_agent_hash_tag(namespace, agent)}}}" + + +def _agent_generation_key(namespace: str, agent: str) -> str: + return f"agent-control:effective-controls:generation:v1:{{{_agent_hash_tag(namespace, agent)}}}" + + +def _agent_hash_tag(namespace: str, agent: str) -> str: + return f"{_encode_key_component(namespace)}:{_encode_key_component(agent)}" + + +def _stale_key(snapshot_key: str) -> str: + suffix = snapshot_key.removeprefix("agent-control:effective-controls:v1:") + return f"agent-control:effective-controls:stale:v1:{suffix}" + + +class RedisControlCache: + """Redis implementation storing JSON snapshots and a per-agent key index.""" + + def __init__(self, client: RedisCommands, ttl_seconds: int = 60) -> None: + self._client = client + self._ttl_seconds = ttl_seconds + + async def get_generation(self, namespace: str, agent: str) -> str | None: + """Read the version used to guard a subsequent snapshot write.""" + try: + raw = await self._client.get(_agent_generation_key(namespace, agent)) + return raw.decode() if isinstance(raw, bytes) else raw or "0" + except Exception: + _CACHE_FAILURE.inc() + _logger.warning("Effective-control cache generation read failed", exc_info=True) + return None + + async def get( + self, namespace: str, agent: str, target_type: str | None, target_id: str | None + ) -> list[RuntimeControl] | None: + try: + key = build_cache_key(namespace, agent, target_type, target_id) + script = """ + if redis.call('EXISTS', KEYS[1]) == 1 then return nil end + return redis.call('GET', KEYS[2]) + """ + raw = await self._client.eval(script, 2, _stale_key(key), key) + if raw is None: + _CACHE_MISS.inc() + return None + if not isinstance(raw, (str, bytes)): + raise TypeError("Redis returned a non-string effective-control snapshot") + payload = json.loads(raw) + if payload.get("schema_version") != 1 or not isinstance(payload.get("controls"), list): + _CACHE_MISS.inc() + return None + controls = [ + RuntimeControl( + id=item["id"], + name=item["name"], + control=ControlDefinitionRuntime.model_validate(item["control"]), + ) + for item in payload["controls"] + ] + _CACHE_HIT.inc() + return controls + except Exception: # Redis and corrupt cache data must not block evaluation. + _CACHE_FAILURE.inc() + _CACHE_MISS.inc() + _logger.warning( + "Effective-control cache read failed; loading from PostgreSQL", exc_info=True + ) + return None + + async def set( + self, + namespace: str, + agent: str, + target_type: str | None, + target_id: str | None, + controls: Sequence[RuntimeControl], + generation: str | None = None, + ) -> None: + key = build_cache_key(namespace, agent, target_type, target_id) + payload = { + "schema_version": 1, + "controls": [ + {"id": item.id, "name": item.name, "control": item.control.model_dump(mode="json")} + for item in controls + ], + "loaded_at": datetime.now(UTC).isoformat(), + } + serialized = json.dumps(payload, separators=(",", ":")) + _SNAPSHOT_SIZE.observe(len(serialized.encode("utf-8"))) + try: + if generation is None: + return + index_key = _agent_index_key(namespace, agent) + generation_key = _agent_generation_key(namespace, agent) + stale_key = _stale_key(key) + script = """ + local current = redis.call('GET', KEYS[1]) or '0' + if current ~= ARGV[1] then return 0 end + redis.call('SET', KEYS[2], ARGV[2], 'EX', ARGV[3]) + redis.call('SADD', KEYS[3], KEYS[2]) + redis.call('EXPIRE', KEYS[3], ARGV[3] * 2) + redis.call('DEL', KEYS[4]) + return 1 + """ + await self._client.eval( + script, + 4, + generation_key, + key, + index_key, + stale_key, + generation, + serialized, + self._ttl_seconds, + ) + except Exception: + _CACHE_FAILURE.inc() + _logger.warning("Effective-control cache write failed", exc_info=True) + + async def invalidate_agent(self, namespace: str, agent: str) -> None: + index_key = _agent_index_key(namespace, agent) + generation_key = _agent_generation_key(namespace, agent) + try: + script = """ + redis.call('INCR', KEYS[1]) + local keys = redis.call('SMEMBERS', KEYS[2]) + local prefix = 'agent-control:effective-controls:v1:' + for _, key in ipairs(keys) do + local suffix = string.sub(key, string.len(prefix) + 1) + local stale = 'agent-control:effective-controls:stale:v1:' .. suffix + redis.call('SET', stale, '1', 'EX', ARGV[1]) + end + return #keys + """ + await self._client.eval( + script, 2, generation_key, index_key, self._ttl_seconds * 2 + ) + _CACHE_INVALIDATION.inc() + except Exception: + _CACHE_FAILURE.inc() + _logger.warning("Effective-control cache invalidation failed", exc_info=True) + + async def close(self) -> None: + """Close the shared Redis connection pool.""" + try: + await self._client.aclose() + except Exception: + _CACHE_FAILURE.inc() + _logger.warning("Effective-control cache shutdown failed", exc_info=True) + + +def get_control_cache() -> ControlCache: + """Construct the configured cache; PostgreSQL-only remains the default.""" + return _configured_control_cache() + + +@lru_cache(maxsize=1) +def _configured_control_cache() -> ControlCache: + if settings.control_cache_backend.lower() != "redis" or not settings.redis_url: + return NoOpControlCache() + try: + client = _create_redis_client(settings.redis_url, settings.redis_mode) + return RedisControlCache(client, settings.control_cache_ttl_seconds) + except Exception: + _CACHE_FAILURE.inc() + _logger.warning("Redis cache initialization failed; using PostgreSQL", exc_info=True) + return NoOpControlCache() + + +def _create_redis_client(redis_url: str, redis_mode: str) -> RedisCommands: + """Build a redis-py client using a standard redis/rediss URL.""" + from redis.asyncio import Redis # type: ignore[import-not-found] + + options = { + "decode_responses": True, + "socket_connect_timeout": 0.5, + "socket_timeout": 0.5, + } + if redis_mode == "cluster": + from redis.asyncio.cluster import RedisCluster # type: ignore[import-not-found] + + client_factory = cast(Any, RedisCluster) + else: + client_factory = cast(Any, Redis) + return cast(RedisCommands, client_factory.from_url(redis_url, **options)) + + +def observe_snapshot_load(started_at: float) -> None: + """Record database snapshot load duration.""" + _SNAPSHOT_LOAD.observe(time.perf_counter() - started_at) + + +async def close_control_cache() -> None: + """Close the shared Redis client during application shutdown.""" + cache = get_control_cache() + if isinstance(cache, RedisControlCache): + await cache.close() + + +async def invalidate_namespace_agents(db: AsyncSession, namespace: str) -> None: + """Invalidate every agent scope in a namespace after a namespace-wide mutation.""" + cache = get_control_cache() + if isinstance(cache, NoOpControlCache): + return + try: + result = await db.execute(select(Agent.name).where(Agent.namespace_key == namespace)) + for (agent_name,) in result.all(): + await cache.invalidate_agent(namespace, agent_name) + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to discover namespace cache scopes after commit; relying on TTL", + exc_info=True, + ) + + +async def invalidate_agent_scope(namespace: str, agent: str) -> None: + """Invalidate all target variants for one agent scope.""" + try: + await get_control_cache().invalidate_agent(namespace, agent) + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to invalidate agent cache scope after commit; relying on TTL", + exc_info=True, + ) + + +async def invalidate_control_agents( + db: AsyncSession, namespace: str, control_id: int +) -> None: + """Invalidate agents related to a changed control through direct or policy bindings.""" + cache = get_control_cache() + if isinstance(cache, NoOpControlCache): + return + direct_agents = select(agent_controls.c.agent_name).where( + agent_controls.c.namespace_key == namespace, + agent_controls.c.control_id == control_id, + ) + policy_agents = ( + select(agent_policies.c.agent_name) + .join( + policy_controls, + (policy_controls.c.namespace_key == agent_policies.c.namespace_key) + & (policy_controls.c.policy_id == agent_policies.c.policy_id), + ) + .where( + agent_policies.c.namespace_key == namespace, + policy_controls.c.control_id == control_id, + ) + ) + try: + result = await db.execute(direct_agents.union(policy_agents)) + for (agent_name,) in result.all(): + await cache.invalidate_agent(namespace, agent_name) + + target_binding = await db.execute( + select(ControlBinding.id).where( + ControlBinding.namespace_key == namespace, + ControlBinding.control_id == control_id, + ).limit(1) + ) + if target_binding.first() is not None: + await invalidate_namespace_agents(db, namespace) + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to discover control cache scopes after commit; relying on TTL", + exc_info=True, + ) + + +async def get_control_cache_scopes( + db: AsyncSession, namespace: str, control_id: int +) -> tuple[list[str], bool]: + """Capture related agent scopes before a mutation removes association rows.""" + try: + if isinstance(get_control_cache(), NoOpControlCache): + return [], False + direct_agents = select(agent_controls.c.agent_name).where( + agent_controls.c.namespace_key == namespace, + agent_controls.c.control_id == control_id, + ) + policy_agents = ( + select(agent_policies.c.agent_name) + .join( + policy_controls, + (policy_controls.c.namespace_key == agent_policies.c.namespace_key) + & (policy_controls.c.policy_id == agent_policies.c.policy_id), + ) + .where( + agent_policies.c.namespace_key == namespace, + policy_controls.c.control_id == control_id, + ) + ) + agents_result = await db.execute(direct_agents.union(policy_agents)) + target_result = await db.execute( + select(ControlBinding.id).where( + ControlBinding.namespace_key == namespace, + ControlBinding.control_id == control_id, + ).limit(1) + ) + return [name for (name,) in agents_result.all()], target_result.first() is not None + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to capture control cache scopes before mutation; relying on TTL", + exc_info=True, + ) + return [], False + + +async def invalidate_control_cache_scopes( + db: AsyncSession, namespace: str, agents: Sequence[str], has_target_bindings: bool +) -> None: + """Invalidate captured control scopes after commit.""" + try: + cache = get_control_cache() + for agent_name in agents: + await cache.invalidate_agent(namespace, agent_name) + if has_target_bindings: + await invalidate_namespace_agents(db, namespace) + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to invalidate control cache scopes after commit; relying on TTL", + exc_info=True, + ) + + +async def invalidate_policy_agents(db: AsyncSession, namespace: str, policy_id: int) -> None: + """Invalidate agents assigned to a changed policy.""" + cache = get_control_cache() + if isinstance(cache, NoOpControlCache): + return + try: + result = await db.execute( + select(agent_policies.c.agent_name).where( + agent_policies.c.namespace_key == namespace, + agent_policies.c.policy_id == policy_id, + ) + ) + for (agent_name,) in result.all(): + await cache.invalidate_agent(namespace, agent_name) + except Exception: + _CACHE_INVALIDATION_FAILURE.inc() + _logger.warning( + "Unable to discover policy cache scopes after commit; relying on TTL", + exc_info=True, + ) diff --git a/server/src/agent_control_server/endpoints/agents.py b/server/src/agent_control_server/endpoints/agents.py index e397134b..8466add1 100644 --- a/server/src/agent_control_server/endpoints/agents.py +++ b/server/src/agent_control_server/endpoints/agents.py @@ -37,6 +37,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from ..auth_framework import Operation, Principal, get_authorizer, require_operation +from ..control_cache import invalidate_agent_scope from ..db import get_async_db from ..errors import ( APIValidationError, @@ -939,6 +940,7 @@ async def init_agent( resource="Agent", operation="update", ) + await invalidate_agent_scope(namespace_key, existing.name) controls = await ControlService(db).list_controls_for_agent( existing.name, @@ -1127,6 +1129,7 @@ async def add_agent_policy( ) await db.execute(stmt) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1220,6 +1223,7 @@ async def set_agent_policy( .on_conflict_do_nothing() ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1342,6 +1346,7 @@ async def remove_agent_policy( ) ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1382,6 +1387,7 @@ async def remove_all_agent_policies( ) ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1439,6 +1445,7 @@ async def delete_agent_policy( ) ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1501,6 +1508,7 @@ async def add_agent_control( namespace_key=namespace_key, ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -1544,6 +1552,7 @@ async def remove_agent_control( ) await db.commit() + await invalidate_agent_scope(namespace_key, agent.name) except Exception: await db.rollback() _logger.error( @@ -2000,6 +2009,7 @@ async def patch_agent( resource="Agent", operation="patch", ) + await invalidate_agent_scope(namespace_key, agent.name) return PatchAgentResponse( steps_removed=steps_removed, diff --git a/server/src/agent_control_server/endpoints/control_bindings.py b/server/src/agent_control_server/endpoints/control_bindings.py index 93fc2d82..7454203e 100644 --- a/server/src/agent_control_server/endpoints/control_bindings.py +++ b/server/src/agent_control_server/endpoints/control_bindings.py @@ -26,6 +26,7 @@ from ..auth_framework import Operation, Principal, get_authorizer, require_operation from ..auth_framework.core import IdentityResolver +from ..control_cache import invalidate_namespace_agents from ..db import AsyncSessionLocal, get_async_db from ..errors import BadRequestError, ForbiddenError, NotFoundError from ..models import ControlBinding @@ -158,6 +159,7 @@ async def create_control_binding( enabled=request.enabled, ) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) await db.refresh(binding) return CreateControlBindingResponse(binding_id=binding.id) @@ -286,6 +288,7 @@ async def patch_control_binding_by_key( enabled=request.enabled, ) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) return PatchControlBindingResponse(success=True, enabled=binding.enabled) @@ -314,6 +317,7 @@ async def patch_control_binding( enabled=request.enabled, ) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) return PatchControlBindingResponse(success=True, enabled=binding.enabled) @@ -337,6 +341,7 @@ async def delete_control_binding( service = ControlBindingsService(db) await service.delete_binding(namespace_key=principal.namespace_key, binding_id=binding_id) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) return DeleteControlBindingResponse(success=True) @@ -369,6 +374,7 @@ async def upsert_control_binding_by_key( enabled=request.enabled, ) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) await db.refresh(binding) return UpsertControlBindingResponse( binding_id=binding.id, @@ -404,4 +410,5 @@ async def delete_control_binding_by_key( control_id=request.control_id, ) await db.commit() + await invalidate_namespace_agents(db, principal.namespace_key) return DeleteControlBindingByKeyResponse(deleted=deleted) diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index 5d8b5b9a..e85030b8 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -43,6 +43,12 @@ from sqlalchemy.ext.asyncio import AsyncSession from ..auth_framework import Operation, Principal, get_authorizer, require_operation +from ..control_cache import ( + get_control_cache_scopes, + invalidate_control_agents, + invalidate_control_cache_scopes, + invalidate_namespace_agents, +) from ..db import get_async_db from ..errors import ( APIError, @@ -838,6 +844,7 @@ async def clone_and_bind_control( enabled=request.target_binding.enabled, ) await db.commit() + await invalidate_namespace_agents(db, namespace_key) except APIError: await db.rollback() raise @@ -1098,6 +1105,7 @@ async def set_control_data( note="Edited", ) await db.commit() + await invalidate_control_agents(db, principal.namespace_key, control_id) except Exception: await db.rollback() _logger.error( @@ -1407,6 +1415,9 @@ async def delete_control( target_binding_ids = await bindings_service.list_binding_ids_for_control( namespace_key=namespace_key, control_id=control_id ) + cache_agents, has_target_bindings = await get_control_cache_scopes( + db, namespace_key, control_id + ) if ( associated_policy_ids @@ -1493,6 +1504,9 @@ async def delete_control( note="Deleted", ) await db.commit() + await invalidate_control_cache_scopes( + db, namespace_key, cache_agents, has_target_bindings + ) _logger.info("Soft-deleted control '%s' (%s)", control.name, control_id) except Exception: await db.rollback() @@ -1629,6 +1643,7 @@ async def patch_control( note="Edited", ) await db.commit() + await invalidate_control_agents(db, namespace_key, control_id) _logger.info(f"Updated control '{control.name}' ({control_id})") except IntegrityError as exc: await db.rollback() diff --git a/server/src/agent_control_server/endpoints/evaluation.py b/server/src/agent_control_server/endpoints/evaluation.py index 7a46829e..e89fc540 100644 --- a/server/src/agent_control_server/endpoints/evaluation.py +++ b/server/src/agent_control_server/endpoints/evaluation.py @@ -1,6 +1,7 @@ """Evaluation analysis endpoints.""" import json +import time from dataclasses import dataclass from agent_control_engine.core import ControlEngine @@ -16,6 +17,7 @@ from sqlalchemy import select from ..auth_framework import Operation, Principal, require_operation +from ..control_cache import get_control_cache, observe_snapshot_load from ..db import AsyncSessionLocal from ..errors import APIValidationError, NotFoundError from ..logging_utils import get_logger @@ -158,6 +160,13 @@ async def _load_engine_controls( ) -> list[ControlAdapter]: """Load and materialize controls before evaluator execution starts.""" namespace_key = principal.namespace_key + cache = get_control_cache() + cached = await cache.get( + namespace_key, request.agent_name, request.target_type, request.target_id + ) + if cached is not None: + return [ControlAdapter(c.id, c.name, c.control) for c in cached] + cache_generation = await cache.get_generation(namespace_key, request.agent_name) async with AsyncSessionLocal() as db: agent_result = await db.execute( @@ -176,6 +185,7 @@ async def _load_engine_controls( hint="Register the agent via initAgent before evaluating.", ) + load_started = time.perf_counter() runtime_controls = await ControlService(db).list_runtime_controls_for_agent( request.agent_name, namespace_key=namespace_key, @@ -183,6 +193,16 @@ async def _load_engine_controls( target_id=request.target_id, allow_invalid_step_name_regex=True, ) + observe_snapshot_load(load_started) + + await cache.set( + namespace_key, + request.agent_name, + request.target_type, + request.target_id, + runtime_controls, + cache_generation, + ) return [ControlAdapter(c.id, c.name, c.control) for c in runtime_controls] diff --git a/server/src/agent_control_server/endpoints/policies.py b/server/src/agent_control_server/endpoints/policies.py index ddda7127..f0475f9f 100644 --- a/server/src/agent_control_server/endpoints/policies.py +++ b/server/src/agent_control_server/endpoints/policies.py @@ -10,6 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from ..auth_framework import Operation, Principal, require_operation +from ..control_cache import invalidate_policy_agents from ..db import get_async_db from ..errors import ConflictError, DatabaseError, NotFoundError from ..logging_utils import get_logger @@ -146,6 +147,7 @@ async def add_control_to_policy( namespace_key=namespace_key, ) await db.commit() + await invalidate_policy_agents(db, namespace_key, policy_id) except Exception: await db.rollback() _logger.error( @@ -228,6 +230,7 @@ async def remove_control_from_policy( namespace_key=namespace_key, ) await db.commit() + await invalidate_policy_agents(db, namespace_key, policy_id) except Exception: await db.rollback() _logger.error( diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index bd0b9efd..cb426332 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -24,6 +24,7 @@ seed_out_of_box_controls, ) from .config import observability_settings, settings +from .control_cache import close_control_cache from .db import AsyncSessionLocal, async_engine from .endpoints.agents import router as agent_router from .endpoints.auth import router as auth_router @@ -232,6 +233,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: await app.state.event_store.close() logger.info("EventStore closed") + await close_control_cache() + await async_engine.dispose() logger.info("Database engine disposed") diff --git a/server/tests/test_agent_cache_invalidation.py b/server/tests/test_agent_cache_invalidation.py new file mode 100644 index 00000000..9e6e3ba1 --- /dev/null +++ b/server/tests/test_agent_cache_invalidation.py @@ -0,0 +1,80 @@ +"""Regression tests for invalidating Agent snapshots after data mutations.""" + +from __future__ import annotations + +import agent_control_server.endpoints.agents as agents_endpoint +import pytest +from agent_control_server.models import DEFAULT_NAMESPACE_KEY +from fastapi.testclient import TestClient + + +def _init_payload(agent_name: str, description: str, step_name: str) -> dict[str, object]: + return { + "agent": { + "agent_name": agent_name, + "agent_description": description, + "agent_version": "1.0", + "agent_metadata": {}, + }, + "steps": [ + { + "type": "tool", + "name": step_name, + "input_schema": {}, + "output_schema": {}, + } + ], + } + + +def test_init_agent_invalidates_existing_agent_after_update( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + invalidations: list[tuple[str, str]] = [] + + async def record_invalidation(namespace: str, agent: str) -> None: + invalidations.append((namespace, agent)) + + monkeypatch.setattr(agents_endpoint, "invalidate_agent_scope", record_invalidation) + agent_name = "cache-init-update-agent" + + created = client.post( + "/api/v1/agents/initAgent", json=_init_payload(agent_name, "before", "step-a") + ) + assert created.status_code == 200 + assert created.json()["created"] is True + assert invalidations == [] + + updated = client.post( + "/api/v1/agents/initAgent", json=_init_payload(agent_name, "after", "step-a") + ) + + assert updated.status_code == 200 + assert updated.json()["created"] is False + assert invalidations == [(DEFAULT_NAMESPACE_KEY, agent_name)] + + +def test_patch_agent_invalidates_after_agent_data_commit( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + invalidations: list[tuple[str, str]] = [] + + async def record_invalidation(namespace: str, agent: str) -> None: + invalidations.append((namespace, agent)) + + monkeypatch.setattr(agents_endpoint, "invalidate_agent_scope", record_invalidation) + agent_name = "cache-patch-update-agent" + initialized = client.post( + "/api/v1/agents/initAgent", json=_init_payload(agent_name, "before", "remove-me") + ) + assert initialized.status_code == 200 + assert invalidations == [] + + patched = client.patch( + f"/api/v1/agents/{agent_name}", + json={"remove_steps": [{"type": "tool", "name": "remove-me"}]}, + ) + + assert patched.status_code == 200 + assert patched.json()["steps_removed"] == [{"type": "tool", "name": "remove-me"}] + assert invalidations == [(DEFAULT_NAMESPACE_KEY, agent_name)] diff --git a/server/tests/test_control_cache.py b/server/tests/test_control_cache.py new file mode 100644 index 00000000..32821668 --- /dev/null +++ b/server/tests/test_control_cache.py @@ -0,0 +1,608 @@ +"""Focused unit tests for the effective-control cache implementations.""" + +from __future__ import annotations + +import asyncio +import sys +import time +import types + +import pytest +from agent_control_models import ControlDefinitionRuntime +from agent_control_server import control_cache as cache_module +from agent_control_server.config import Settings +from agent_control_server.control_cache import ( + NoOpControlCache, + RedisControlCache, + _agent_generation_key, + _agent_index_key, + _stale_key, + build_cache_key, + get_control_cache_scopes, + invalidate_control_agents, + invalidate_control_cache_scopes, + invalidate_namespace_agents, + invalidate_policy_agents, +) +from agent_control_server.services.controls import RuntimeControl + + +class FakeRedis: + def __init__(self) -> None: + self.values: dict[str, tuple[str, float | None]] = {} + self.sets: dict[str, set[str]] = {} + + async def get(self, name: str) -> str | None: + value = self.values.get(name) + if value is None: + return None + payload, expires_at = value + if expires_at is not None and expires_at <= time.monotonic(): + self.values.pop(name, None) + return None + return payload + + async def set(self, name: str, value: str, ex: int) -> None: + self.values[name] = (value, time.monotonic() + ex) + + async def sadd(self, name: str, *values: str) -> None: + self.sets.setdefault(name, set()).update(values) + + async def expire(self, name: str, seconds: int) -> None: + del name, seconds + + async def smembers(self, name: str) -> set[str]: + return self.sets.get(name, set()).copy() + + async def delete(self, *names: str) -> None: + for name in names: + self.values.pop(name, None) + self.sets.pop(name, None) + + async def eval( + self, script: str, numkeys: int, *keys_and_args: str | int + ) -> str | int | None: + keys = [str(value) for value in keys_and_args[:numkeys]] + args = keys_and_args[numkeys:] + if "INCR" in script: + generation_key, index_key = keys + current = self.values.get(generation_key, ("0", None))[0] + self.values[generation_key] = (str(int(current) + 1), None) + indexed = self.sets.get(index_key, set()) + ttl = int(args[0]) + for key in indexed: + suffix = key.removeprefix("agent-control:effective-controls:v1:") + stale_key = f"agent-control:effective-controls:stale:v1:{suffix}" + await self.set(stale_key, "1", ttl) + return len(indexed) + if "EXISTS" in script: + stale_key, snapshot_key = keys + if await self.get(stale_key) is not None: + return None + return await self.get(snapshot_key) + generation_key, snapshot_key, index_key, stale_key = keys + expected_generation, payload, ttl = args + current = self.values.get(generation_key, ("0", None))[0] + if current != str(expected_generation): + return 0 + await self.set(snapshot_key, str(payload), int(ttl)) + await self.sadd(index_key, snapshot_key) + await self.delete(stale_key) + return 1 + + async def aclose(self) -> None: + return None + + def expire_now(self, name: str) -> None: + payload, _ = self.values[name] + self.values[name] = (payload, time.monotonic() - 1) + + +class BrokenRedis(FakeRedis): + async def get(self, name: str) -> str | None: + del name + raise ConnectionError("redis unavailable") + + +class FailingRedisCommands(FakeRedis): + async def eval( + self, script: str, numkeys: int, *keys_and_args: str | int + ) -> str | int | None: + del script, numkeys, keys_and_args + raise ConnectionError("redis command failed") + + async def aclose(self) -> None: + raise ConnectionError("redis close failed") + + +class NonStringSnapshotRedis(FakeRedis): + async def eval( + self, script: str, numkeys: int, *keys_and_args: str | int + ) -> str | int | None: + del script, numkeys, keys_and_args + return 1 + + +class FakeResult: + def __init__(self, rows: list[tuple[str]]) -> None: + self.rows = rows + + def all(self) -> list[tuple[str]]: + return self.rows + + def first(self) -> tuple[str] | None: + return self.rows[0] if self.rows else None + + +class FakeDatabase: + def __init__(self, results: list[FakeResult]) -> None: + self.results = iter(results) + + async def execute(self, statement: object) -> FakeResult: + del statement + return next(self.results) + + +class FailingDatabase: + async def execute(self, statement: object) -> FakeResult: + del statement + raise RuntimeError("database scope lookup failed") + + +class MetricSpy: + def __init__(self) -> None: + self.calls = 0 + + def inc(self) -> None: + self.calls += 1 + + +def test_postgres_only_defaults_to_no_op_cache() -> None: + settings = Settings(_env_file=None) + assert settings.control_cache_backend == "postgres" + assert settings.redis_url is None + assert settings.redis_mode == "standalone" + assert asyncio.run(NoOpControlCache().get("ns", "agent", None, None)) is None + + +def test_configured_cache_uses_redis_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None: + client = FakeRedis() + calls: list[tuple[str, str]] = [] + monkeypatch.setattr(cache_module.settings, "control_cache_backend", "redis") + monkeypatch.setattr(cache_module.settings, "redis_url", "redis://cache.example:6379") + monkeypatch.setattr(cache_module.settings, "redis_mode", "cluster") + monkeypatch.setattr(cache_module.settings, "control_cache_ttl_seconds", 17) + + def create_client(url: str, mode: str) -> FakeRedis: + calls.append((url, mode)) + return client + + monkeypatch.setattr(cache_module, "_create_redis_client", create_client) + cache_module._configured_control_cache.cache_clear() + try: + configured = cache_module.get_control_cache() + finally: + cache_module._configured_control_cache.cache_clear() + + assert isinstance(configured, RedisControlCache) + assert configured._ttl_seconds == 17 + assert calls == [("redis://cache.example:6379", "cluster")] + + +def test_configured_cache_falls_back_when_redis_initialization_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(cache_module.settings, "control_cache_backend", "redis") + monkeypatch.setattr(cache_module.settings, "redis_url", "redis://cache.example:6379") + + def fail_to_create_client(url: str, mode: str) -> FakeRedis: + del url, mode + raise ImportError("redis package is unavailable") + + monkeypatch.setattr(cache_module, "_create_redis_client", fail_to_create_client) + cache_module._configured_control_cache.cache_clear() + try: + configured = cache_module.get_control_cache() + finally: + cache_module._configured_control_cache.cache_clear() + + assert isinstance(configured, NoOpControlCache) + + +def test_cache_key_includes_all_scope_dimensions() -> None: + assert ( + build_cache_key("ns", "agent", None, None) + == "agent-control:effective-controls:v1:{=bnM:=YWdlbnQ}:!:!" + ) + base = build_cache_key("ns", "agent:other", "account", "42") + assert base != build_cache_key("other", "agent", "account", "42") + assert base != build_cache_key("ns", "agent", "other:account", "42") + assert base != build_cache_key("ns", "agent:other", "workspace", "42") + assert base != build_cache_key("ns", "agent:other", "account", "43") + assert base != build_cache_key("ns:agent", "other", "account", "42") + + +def test_redis_script_keys_share_hash_tag_and_target_scope() -> None: + key = build_cache_key("ns", "agent", "log_stream", "logstream-1") + keys = [ + key, + _agent_index_key("ns", "agent"), + _agent_generation_key("ns", "agent"), + _stale_key(key), + ] + hash_tags = [item[item.index("{") + 1 : item.index("}")] for item in keys] + assert len(set(hash_tags)) == 1 + assert key != build_cache_key("ns", "agent", "log_stream", "logstream-2") + assert key != build_cache_key("ns", "agent", "account", "logstream-1") + + +@pytest.mark.parametrize( + "redis_mode,expected_client", [("standalone", "redis"), ("cluster", "cluster")] +) +def test_redis_client_factory_uses_selected_topology( + monkeypatch: pytest.MonkeyPatch, redis_mode: str, expected_client: str +) -> None: + calls: list[tuple[str, str, dict[str, object]]] = [] + + class StandaloneRedis: + @classmethod + def from_url(cls, url: str, **options: object) -> object: + calls.append(("redis", url, options)) + return cls() + + class ClusterRedis: + @classmethod + def from_url(cls, url: str, **options: object) -> object: + calls.append(("cluster", url, options)) + return cls() + + redis_package = types.ModuleType("redis") + redis_asyncio = types.ModuleType("redis.asyncio") + redis_asyncio.Redis = StandaloneRedis # type: ignore[attr-defined] + cluster_module = types.ModuleType("redis.asyncio.cluster") + cluster_module.RedisCluster = ClusterRedis # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "redis", redis_package) + monkeypatch.setitem(sys.modules, "redis.asyncio", redis_asyncio) + monkeypatch.setitem(sys.modules, "redis.asyncio.cluster", cluster_module) + + client = cache_module._create_redis_client("rediss://cache.example:6380", redis_mode) + + assert isinstance(client, StandaloneRedis if expected_client == "redis" else ClusterRedis) + assert calls == [ + ( + expected_client, + "rediss://cache.example:6380", + { + "decode_responses": True, + "socket_connect_timeout": 0.5, + "socket_timeout": 0.5, + }, + ) + ] + + +def test_redis_miss_hit_and_ttl_expiry() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=1) + assert await cache.get("ns", "agent", None, None) is None + generation = await cache.get_generation("ns", "agent") + await cache.set("ns", "agent", None, None, [], generation) + assert await cache.get("ns", "agent", None, None) == [] + client.expire_now(build_cache_key("ns", "agent", None, None)) + assert await cache.get("ns", "agent", None, None) is None + + asyncio.run(run()) + + +def test_redis_failure_falls_back_to_database() -> None: + async def run() -> None: + cache = RedisControlCache(BrokenRedis(), ttl_seconds=10) + assert await cache.get("ns", "agent", None, None) is None + assert await cache.get_generation("ns", "agent") is None + await cache.set("ns", "agent", None, None, []) + await cache.invalidate_agent("ns", "agent") + await cache.close() + + asyncio.run(run()) + + +def test_redis_corrupt_snapshot_is_treated_as_a_miss() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=10) + key = build_cache_key("ns", "agent", None, None) + client.values[key] = ('{"schema_version":2,"controls":[]}', None) + assert await cache.get("ns", "agent", None, None) is None + + asyncio.run(run()) + + +def test_redis_non_string_snapshot_is_treated_as_a_miss() -> None: + async def run() -> None: + cache = RedisControlCache(NonStringSnapshotRedis(), ttl_seconds=10) + assert await cache.get("ns", "agent", None, None) is None + + asyncio.run(run()) + + +def test_redis_write_and_invalidation_failures_are_best_effort() -> None: + async def run() -> None: + cache = RedisControlCache(FailingRedisCommands(), ttl_seconds=10) + await cache.set("ns", "agent", None, None, [], "0") + await cache.invalidate_agent("ns", "agent") + await cache.close() + + asyncio.run(run()) + + +def test_invalidation_removes_all_target_snapshots_for_agent() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + await cache.set("ns", "agent", None, None, [], await cache.get_generation("ns", "agent")) + await cache.set( + "ns", "agent", "account", "1", [], await cache.get_generation("ns", "agent") + ) + await cache.set( + "ns", "other-agent", None, None, [], await cache.get_generation("ns", "other-agent") + ) + await cache.invalidate_agent("ns", "agent") + assert await cache.get("ns", "agent", None, None) is None + assert await cache.get("ns", "agent", "account", "1") is None + assert await cache.get("ns", "other-agent", None, None) == [] + + asyncio.run(run()) + + +def test_control_mutation_invalidates_relationship_bound_agents( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + await cache.set( + "ns", "related-agent", None, None, [], await cache.get_generation("ns", "related-agent") + ) + await cache.set( + "ns", + "unrelated-agent", + None, + None, + [], + await cache.get_generation("ns", "unrelated-agent"), + ) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + db = FakeDatabase([FakeResult([("related-agent",)]), FakeResult([])]) + await invalidate_control_agents(db, "ns", 7) # type: ignore[arg-type] + assert await cache.get("ns", "related-agent", None, None) is None + assert await cache.get("ns", "unrelated-agent", None, None) == [] + + asyncio.run(run()) + + +def test_control_target_binding_invalidates_namespace_after_relationship_lookup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + generation = await cache.get_generation("ns", "agent") + await cache.set("ns", "agent", "log_stream", "stream-1", [], generation) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + db = FakeDatabase( + [FakeResult([]), FakeResult([("binding-id",)]), FakeResult([("agent",)])] + ) + + await invalidate_control_agents(db, "ns", 7) # type: ignore[arg-type] + + assert await cache.get("ns", "agent", "log_stream", "stream-1") is None + + asyncio.run(run()) + + +def test_control_cache_scopes_capture_and_invalidate_related_agents( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + generation = await cache.get_generation("ns", "related-agent") + await cache.set("ns", "related-agent", None, None, [], generation) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + db = FakeDatabase( + [ + FakeResult([("related-agent",)]), + FakeResult([("binding-id",)]), + FakeResult([("related-agent",)]), + ] + ) + + agents, has_target_bindings = await get_control_cache_scopes( + db, "ns", 7 # type: ignore[arg-type] + ) + assert agents == ["related-agent"] + assert has_target_bindings + await invalidate_control_cache_scopes( + db, "ns", agents, has_target_bindings # type: ignore[arg-type] + ) + assert await cache.get("ns", "related-agent", None, None) is None + + asyncio.run(run()) + + +def test_policy_mutation_invalidates_assigned_agents(monkeypatch: pytest.MonkeyPatch) -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + assigned_generation = await cache.get_generation("ns", "assigned-agent") + await cache.set("ns", "assigned-agent", None, None, [], assigned_generation) + await cache.set( + "ns", "other-agent", None, None, [], await cache.get_generation("ns", "other-agent") + ) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + db = FakeDatabase([FakeResult([("assigned-agent",)])]) + await invalidate_policy_agents(db, "ns", 4) # type: ignore[arg-type] + assert await cache.get("ns", "assigned-agent", None, None) is None + assert await cache.get("ns", "other-agent", None, None) == [] + + asyncio.run(run()) + + +def test_target_binding_change_invalidates_namespace_agent_scopes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + first_generation = await cache.get_generation("ns", "agent-one") + second_generation = await cache.get_generation("ns", "agent-two") + await cache.set("ns", "agent-one", "account", "a", [], first_generation) + await cache.set("ns", "agent-two", None, None, [], second_generation) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + db = FakeDatabase([FakeResult([("agent-one",), ("agent-two",)])]) + await invalidate_namespace_agents(db, "ns") # type: ignore[arg-type] + assert await cache.get("ns", "agent-one", "account", "a") is None + assert await cache.get("ns", "agent-two", None, None) is None + + asyncio.run(run()) + + +def test_post_commit_scope_lookup_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: + async def run() -> None: + cache = RedisControlCache(FakeRedis(), ttl_seconds=30) + failure_metric = MetricSpy() + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + monkeypatch.setattr(cache_module, "_CACHE_INVALIDATION_FAILURE", failure_metric) + + await invalidate_namespace_agents(FailingDatabase(), "ns") # type: ignore[arg-type] + + assert failure_metric.calls == 1 + + asyncio.run(run()) + + +def test_control_scope_capture_failure_falls_back_to_ttl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + cache = RedisControlCache(FakeRedis(), ttl_seconds=30) + failure_metric = MetricSpy() + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + monkeypatch.setattr(cache_module, "_CACHE_INVALIDATION_FAILURE", failure_metric) + + agents, has_target_bindings = await get_control_cache_scopes( + FailingDatabase(), "ns", 7 # type: ignore[arg-type] + ) + + assert agents == [] + assert has_target_bindings is False + assert failure_metric.calls == 1 + + asyncio.run(run()) + + +def test_control_and_policy_scope_lookup_failures_are_best_effort( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def run() -> None: + cache = RedisControlCache(FakeRedis(), ttl_seconds=30) + failure_metric = MetricSpy() + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + monkeypatch.setattr(cache_module, "_CACHE_INVALIDATION_FAILURE", failure_metric) + + await invalidate_control_agents(FailingDatabase(), "ns", 7) # type: ignore[arg-type] + await invalidate_policy_agents(FailingDatabase(), "ns", 7) # type: ignore[arg-type] + + assert failure_metric.calls == 2 + + asyncio.run(run()) + + +def test_close_control_cache_closes_redis_client(monkeypatch: pytest.MonkeyPatch) -> None: + async def run() -> None: + cache = RedisControlCache(FakeRedis(), ttl_seconds=30) + closed: list[bool] = [] + + async def aclose() -> None: + closed.append(True) + + monkeypatch.setattr(cache._client, "aclose", aclose) + monkeypatch.setattr(cache_module, "get_control_cache", lambda: cache) + await cache_module.close_control_cache() + assert closed == [True] + + asyncio.run(run()) + + +def test_runtime_control_snapshot_round_trip() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + control = ControlDefinitionRuntime.model_validate( + { + "execution": "server", + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "regex", "config": {"pattern": "bad"}}, + }, + "action": {"decision": "deny"}, + } + ) + snapshot = [RuntimeControl(id=21, name="block-bad", control=control)] + generation = await cache.get_generation("ns", "agent") + await cache.set("ns", "agent", None, None, snapshot, generation) + restored = await cache.get("ns", "agent", None, None) + assert restored == snapshot + assert restored is not None + assert restored[0].id == 21 + assert restored[0].name == "block-bad" + assert restored[0].control == control + + asyncio.run(run()) + + +def test_generation_prevents_stale_snapshot_repopulation() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + generation = await cache.get_generation("ns", "agent") + await cache.invalidate_agent("ns", "agent") + await cache.set("ns", "agent", None, None, [], generation) + assert await cache.get("ns", "agent", None, None) is None + + asyncio.run(run()) + + +def test_stale_snapshot_is_bypassed_until_postgres_refresh() -> None: + async def run() -> None: + client = FakeRedis() + cache = RedisControlCache(client, ttl_seconds=30) + old = _runtime_control(1, "old-snapshot") + fresh = _runtime_control(2, "fresh-snapshot") + generation = await cache.get_generation("ns", "agent") + await cache.set("ns", "agent", None, None, [old], generation) + + await cache.invalidate_agent("ns", "agent") + assert await cache.get("ns", "agent", None, None) is None + + fresh_generation = await cache.get_generation("ns", "agent") + await cache.set("ns", "agent", None, None, [fresh], fresh_generation) + restored = await cache.get("ns", "agent", None, None) + assert restored == [fresh] + + asyncio.run(run()) + + +def _runtime_control(control_id: int, name: str) -> RuntimeControl: + control = ControlDefinitionRuntime.model_validate( + { + "execution": "server", + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "regex", "config": {"pattern": "bad"}}, + }, + "action": {"decision": "deny"}, + } + ) + return RuntimeControl(id=control_id, name=name, control=control)