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..5223c1d1 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 @@ -10,6 +10,7 @@ from galileo_core.schemas.logging.span import ( AgentSpan, ControlSpan, + HttpSpan, LlmSpan, RetrieverSpan, ToolSpan, @@ -21,7 +22,9 @@ type GalileoRecord = LlmSpan | ToolSpan | RetrieverSpan | Trace | Session type GalileoSpan = LlmSpan | ToolSpan | RetrieverSpan -type GalileoCoreSpan = AgentSpan | WorkflowSpan | LlmSpan | RetrieverSpan | ToolSpan | ControlSpan +type GalileoCoreSpan = ( + AgentSpan | WorkflowSpan | LlmSpan | RetrieverSpan | ToolSpan | ControlSpan | HttpSpan +) _MISSING = object() _SPAN_TYPES = {"llm", "tool", "retriever"} _RECORD_TYPES = _SPAN_TYPES | {"trace", "session"} @@ -129,8 +132,10 @@ def record_from_step( selected_value = selected_data if selected_data is _MISSING: selected_value = selected_input if selected_input is not _MISSING else _MISSING - if selected_value is not _MISSING and selected_value is not None and not isinstance( - selected_value, Mapping + if ( + selected_value is not _MISSING + and selected_value is not None + and not isinstance(selected_value, Mapping) ): raise RecordFactoryError( f"Galileo {record_type} records require structured selector values; " diff --git a/server/DIAGNOSTIC_LOOKUP.md b/server/DIAGNOSTIC_LOOKUP.md new file mode 100644 index 00000000..be446390 --- /dev/null +++ b/server/DIAGNOSTIC_LOOKUP.md @@ -0,0 +1,21 @@ +# Controlled control-lookup diagnostic + +This default-disabled feature is for a dedicated owned load-test app. It is not production caching or production write-event invalidation. Do not enable it on a shared or customer app. + +The Orbit launcher injects its existing RedisConfig async client. The standalone server does not import Orbit or cache authorization. Without GALILEO_REDIS_HOST, Orbit does not construct a Redis client and database lookup remains available. Authentication and runtime-use authorization run on every evaluation, including Redis hits. The Redis connection is closed by the launcher lifespan. + +Enable AGENT_CONTROL_LOOKUP_DIAGNOSTIC_ENABLED only with an explicit AGENT_CONTROL_LOOKUP_DIAGNOSTIC_FIXTURES list and timezone-aware AGENT_CONTROL_LOOKUP_DIAGNOSTIC_EXPIRES_AT. Each fixture contains namespace, agent, target_type, target_id, generation, and controls_sha256. Compute the digest over ordered parsed runtime controls with definition_digest and verify it against fresh read-only fixture evidence before deployment. Duplicate target identities can declare distinct reviewed pair generations. Never put identifiers or contents in metric labels. + +Both modes use the same app image. Send X-LT-Control-Lookup-Cache as off or on and X-LT-Control-Lookup-Generation as the exact allowlisted pair generation. Without both headers the request is rejected; without either the original lookup path remains. Unknown/foreign/expired fixtures are rejected. Use a new reviewed generation for each comparison pair to prevent prior tests warming later pairs. Do not add excluded warmup traffic. + +OFF queries the database for the agent and effective controls, then parses the definitions. ON reads serialized runtime definitions from a versioned, hashed Redis key. It validates their digest and returns independently parsed models. A miss uses the same database loader and writes a bounded JSON value with a 300-second TTL. No ORM objects, results, authorization decisions or model outputs are cached. No pickle, global scan or global Redis cleanup is used. + +Every Redis command is bounded by the default 50ms deadline. A failed read, invalid cached value or timeout uses the DB path. A failed SET keeps the fresh DB result. The metric fallback counters must expose these cases; a passing semantic request does not prove a cache hit. Cache hits do not queue behind local locks. Local locks only coalesce misses within one process. They do not promise a single DB load across all replicas. + +While this diagnostic is enabled, its middleware rejects POST/PUT/PATCH/DELETE configuration operations under agents, controls, control-bindings and policies. This freezes the dedicated application, including after cache expiry, until the diagnostic is disabled or the original app restored. Reads, evaluation, token exchange and observation remain available. No external/direct database writer may change the fixtures. A rejected write or changed fixture blocks further load. Capture the original app digest before deployment and restore it after the test. + +Counter labels publish explicit startup zeros so missing series cannot be mistaken for no fallbacks. Numeric metrics are agent_control_lookup_diagnostic_seconds (mode and stage), agent_control_lookup_diagnostic_events_total (mode and outcome), and agent_control_lookup_diagnostic_sql_total (mode and query). Stages include authorization, pool_acquire, agent_query, effective_query, parse, control_load, redis_get, miss_wait, redis_set, cache_encode, cache_decode, engine and handler. Handler timing excludes authorization and response transport; use load-generator and existing HTTP middleware metrics for complete request latency. Query counters count attempted control-loader execute calls; they do not count authorization, pre-ping or every database operation. + +Accept the live treatment only with verified Redis hits, fewer loader SQL calls, unchanged control matches and zero failures, original SLOs, source/image/resource proof and required connected traces. Compare cold and steady windows and retain misses. Histograms are independent distributions and cannot be subtracted to reconstruct request timing. A single OFF/ON pair is exploratory and does not prove a causal bottleneck or capacity boundary. + +The snapshot-age histogram agent_control_lookup_diagnostic_snapshot_age_seconds measures seconds since successful database loading into the Redis envelope, only on valid hits. It is not the age of the underlying configuration. Cross-pod wall clocks are used; timestamps more than one second in the future or older than the maximum 600-second TTL fall back to the database. The miss_wait timer measures only waiting for the diagnostic process-local miss lock. It is not evaluator concurrency, model scheduling or database pool waiting. Cancellation always releases an acquired lock. diff --git a/server/src/agent_control_server/diagnostic_freeze.py b/server/src/agent_control_server/diagnostic_freeze.py new file mode 100644 index 00000000..77129d50 --- /dev/null +++ b/server/src/agent_control_server/diagnostic_freeze.py @@ -0,0 +1,42 @@ +"""Reject configuration writes during the dedicated-stack diagnostic window.""" + +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +from . import diagnostic_lookup +from .diagnostic_metrics import _EVENTS + + +class DiagnosticFixtureFreeze: + """Keep AC configuration fixed while the dedicated LT image is enabled. + + Reads, token exchange, evaluation and observation remain available. Freeze + lasts until disabling the diagnostic or restoring the original image, even + after cache expiry. Never enable this on a shared application. + """ + + def __init__(self, app: ASGIApp, *, api_base: str) -> None: + self.app = app + self._prefixes = tuple( + api_base + "/" + name + for name in ("agents", "controls", "control-bindings", "policies") + ) + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if ( + diagnostic_lookup.settings.enabled + and scope["type"] == "http" + and scope["method"] in {"POST", "PUT", "PATCH", "DELETE"} + and any( + scope["path"] == p or scope["path"].startswith(p + "/") + for p in self._prefixes + ) + ): + _EVENTS.labels("on", "fixture_write_rejected").inc() + response = JSONResponse( + status_code=409, + content={"detail": "Configuration is frozen for the controlled lookup diagnostic"}, + ) + await response(scope, receive, send) + return + await self.app(scope, receive, send) diff --git a/server/src/agent_control_server/diagnostic_lookup.py b/server/src/agent_control_server/diagnostic_lookup.py new file mode 100644 index 00000000..dd92f3ae --- /dev/null +++ b/server/src/agent_control_server/diagnostic_lookup.py @@ -0,0 +1,227 @@ +"""Opt-in lookup-cache experiment for explicitly frozen, expiring LT fixtures. + +Not production invalidation: the controller must prohibit all fixture writes until +restoration. Authorization always precedes this module. No ORM object is retained. +""" + +import asyncio +import hashlib +import json +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Literal, Protocol + +from agent_control_models import ControlDefinitionRuntime +from pydantic import BaseModel, Field +from pydantic_settings import BaseSettings, SettingsConfigDict + +from .diagnostic_metrics import _EVENTS, _SNAPSHOT_AGE, measured +from .services.controls import RuntimeControl + + +class FrozenFixture(BaseModel): + """Exact frozen fixture and expected digest of its ordered runtime definitions.""" + + namespace: str = Field(min_length=1) + agent: str = Field(min_length=1) + target_type: str = Field(min_length=1) + target_id: str = Field(min_length=1) + generation: str = Field(pattern=r"^[a-z0-9][a-z0-9-]{0,79}$") + controls_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + + +class DiagnosticSettings(BaseSettings): + model_config = SettingsConfigDict(env_prefix="AGENT_CONTROL_LOOKUP_DIAGNOSTIC_", extra="ignore") + enabled: bool = False + fixtures: list[FrozenFixture] = Field(default_factory=list, max_length=32) + expires_at: datetime | None = None + ttl_seconds: float = Field(default=300, gt=0, le=600) + redis_timeout_seconds: float = Field(default=0.05, gt=0, le=0.25) + max_value_bytes: int = Field(default=1048576, ge=1024, le=1048576) + + def fixture( + self, + namespace: str, + agent: str, + target_type: str | None, + target_id: str | None, + generation: str | None = None, + ) -> FrozenFixture | None: + expiry = self.expires_at + if ( + not self.enabled + or expiry is None + or expiry.tzinfo is None + or expiry <= datetime.now(UTC) + ): + return None + return next( + ( + f + for f in self.fixtures + if (f.namespace, f.agent, f.target_type, f.target_id) + == (namespace, agent, target_type, target_id) + and (generation is None or f.generation == generation) + ), + None, + ) + + +settings = DiagnosticSettings() + + +def definition_digest(controls: list[RuntimeControl]) -> str: + """Stable ordered digest; output is a hash, never persisted control contents.""" + data = [(c.id, c.name, c.control.model_dump(mode="json")) for c in controls] + return hashlib.sha256( + json.dumps(data, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +@dataclass(frozen=True) +class CacheKey: + namespace: str + agent: str + target_type: str + target_id: str + generation: str + digest: str + + +class RedisClient(Protocol): + """Small async transport seam; Orbit supplies RedisConfig's async client.""" + + async def get(self, name: str) -> bytes | str | None: ... + async def set(self, name: str, value: str, *, ex: int) -> object: ... + async def aclose(self) -> None: ... + + +class CachedControl(BaseModel): + id: int + name: str + control: ControlDefinitionRuntime + + +class CachedDefinitions(BaseModel): + format_version: Literal[1] = 1 + controls: list[CachedControl] + loaded_at_unix: float = Field(gt=0, allow_inf_nan=False) + + +def redis_key(key: CacheKey) -> str: + """Hash identifiers so keys expose no fixture identity and remain bounded.""" + identity = json.dumps( + [key.namespace, key.agent, key.target_type, key.target_id, key.generation, key.digest], + separators=(",", ":"), + ) + return "galileo:ac:lt-lookup:v1:" + hashlib.sha256(identity.encode()).hexdigest() + + +class LookupCache: + """Redis shared cache restricted to a bounded, explicitly frozen fixture set. + + No local results cache. Every ON request measures a Redis GET. Parse cached JSON + into independent runtime models. Failed, timed-out or corrupt Redis data falls + back to the authorized DB loader, never to stale local definitions. Generation + changes address a new key and TTL bounds old keys; fixture writes are forbidden + throughout this diagnostic. This is not production write-event invalidation. + """ + + def __init__(self, factory: Callable[[], RedisClient] | None = None) -> None: + self._factory = factory + self._client: RedisClient | None = None + self._locks = [asyncio.Lock() for _ in range(64)] + + def _redis(self) -> RedisClient: + if self._client is None: + if self._factory is None: + raise RuntimeError("Redis diagnostic client not configured") + self._client = self._factory() + return self._client + + async def aclose(self) -> None: + if self._client is not None: + await self._client.aclose() + + async def _read(self, key: CacheKey) -> list[RuntimeControl] | None: + with measured("redis_get"): + async with asyncio.timeout(settings.redis_timeout_seconds): + value = await self._redis().get(redis_key(key)) + if value is None: + return None + if len(value) > settings.max_value_bytes: + raise ValueError("Redis diagnostic value exceeds bound") + with measured("cache_decode"): + parsed = CachedDefinitions.model_validate_json(value) + controls = [RuntimeControl(c.id, c.name, c.control) for c in parsed.controls] + if definition_digest(controls) != key.digest: + raise ValueError("Redis diagnostic digest mismatch") + age = time.time() - parsed.loaded_at_unix + if age < -1 or age > 600: + raise ValueError("Redis diagnostic snapshot timestamp outside bound") + _SNAPSHOT_AGE.labels("on").observe(max(0, age)) + return controls + + async def load( + self, + key: CacheKey, + loader: Callable[[], Awaitable[list[RuntimeControl]]], + *, + ttl: float, + ) -> list[RuntimeControl]: + # Hits must not queue behind a process-local lock. Only coalesce misses. + try: + controls = await self._read(key) + except Exception: + _EVENTS.labels("on", "redis_read_fallback").inc() + return await loader() + if controls is not None: + _EVENTS.labels("on", "hit").inc() + return controls + lock = self._locks[hash(key) % len(self._locks)] + with measured("miss_wait"): + await lock.acquire() + try: + try: + controls = await self._read(key) + except Exception: + _EVENTS.labels("on", "redis_read_fallback").inc() + return await loader() + if controls is not None: + _EVENTS.labels("on", "hit").inc() + return controls + _EVENTS.labels("on", "miss").inc() + controls = await loader() + if definition_digest(controls) != key.digest: + _EVENTS.labels("on", "fixture_mismatch").inc() + raise ValueError("Frozen diagnostic fixture digest mismatch") + with measured("cache_encode"): + value = CachedDefinitions( + loaded_at_unix=time.time(), + controls=[ + CachedControl(id=c.id, name=c.name, control=c.control) for c in controls + ], + ).model_dump_json() + if len(value.encode()) > settings.max_value_bytes: + _EVENTS.labels("on", "oversize_fallback").inc() + return controls + try: + with measured("redis_set"): + async with asyncio.timeout(settings.redis_timeout_seconds): + await self._redis().set(redis_key(key), value, ex=max(1, int(ttl))) + except Exception: + _EVENTS.labels("on", "redis_write_fallback").inc() + return controls + finally: + lock.release() + + +cache = LookupCache() + + +def configure_redis_factory(factory: Callable[[], RedisClient]) -> None: + """Wire the deployed service's existing RedisConfig without importing Orbit.""" + global cache + cache = LookupCache(factory) diff --git a/server/src/agent_control_server/diagnostic_metrics.py b/server/src/agent_control_server/diagnostic_metrics.py new file mode 100644 index 00000000..10bb1b72 --- /dev/null +++ b/server/src/agent_control_server/diagnostic_metrics.py @@ -0,0 +1,84 @@ +"""Bounded numeric-only timers for the opt-in lookup diagnostic.""" + +import time +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Literal + +from prometheus_client import Counter, Histogram + +Mode = Literal["off", "on"] +Stage = Literal[ + "authorization", + "pool_acquire", + "agent_query", + "effective_query", + "parse", + "redis_get", + "miss_wait", + "redis_set", + "cache_encode", + "cache_decode", + "control_load", + "engine", + "handler", +] +_active_mode: ContextVar[Mode | None] = ContextVar("lookup_diagnostic_mode", default=None) +_DURATION = Histogram( + "agent_control_lookup_diagnostic_seconds", + "Diagnostic stage elapsed seconds", + ["mode", "stage"], + buckets=(0.0001, 0.0005, 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30, 60), +) +_SNAPSHOT_AGE = Histogram( + "agent_control_lookup_diagnostic_snapshot_age_seconds", + "Age since cached definitions were loaded, not configuration creation age", + ["mode"], + buckets=(0, 1, 5, 10, 30, 60, 120, 180, 240, 300, 600), +) +_EVENTS = Counter( + "agent_control_lookup_diagnostic_events", "Diagnostic cache outcomes", ["mode", "outcome"] +) +_SQL = Counter( + "agent_control_lookup_diagnostic_sql", "Control-loader SQL execute calls", ["mode", "query"] +) + + +@contextmanager +def measured(stage: Stage) -> Iterator[None]: + """Record a bounded-label duration only inside an opted-in diagnostic request.""" + mode = _active_mode.get() + if mode is None: + yield + return + started = time.perf_counter() + try: + yield + finally: + if mode is not None: + _DURATION.labels(mode, stage).observe(time.perf_counter() - started) + + +def count_query(query: Literal["agent", "effective"]) -> None: + mode = _active_mode.get() + if mode is not None: + _SQL.labels(mode, query).inc() + + +# Publish explicit zeros for bounded counter labels. An absent fallback series +# must not be reported as zero merely because no error was observed. +for _mode in ("off", "on"): + for _outcome in ( + "bypass", + "hit", + "miss", + "redis_read_fallback", + "redis_write_fallback", + "oversize_fallback", + "fixture_mismatch", + "fixture_write_rejected", + ): + _EVENTS.labels(_mode, _outcome) + for _query in ("agent", "effective"): + _SQL.labels(_mode, _query) diff --git a/server/src/agent_control_server/endpoints/evaluation.py b/server/src/agent_control_server/endpoints/evaluation.py index 7a46829e..4a45a093 100644 --- a/server/src/agent_control_server/endpoints/evaluation.py +++ b/server/src/agent_control_server/endpoints/evaluation.py @@ -1,7 +1,10 @@ """Evaluation analysis endpoints.""" import json +import time from dataclasses import dataclass +from datetime import UTC, datetime +from typing import cast from agent_control_engine.core import ControlEngine from agent_control_models import ( @@ -12,9 +15,11 @@ JSONObject, ) from agent_control_models.errors import ErrorCode, ValidationErrorItem -from fastapi import APIRouter, Depends, Request +from fastapi import APIRouter, Depends, HTTPException, Request from sqlalchemy import select +from .. import diagnostic_lookup as diagnostic +from .. import diagnostic_metrics as timings from ..auth_framework import Operation, Principal, require_operation from ..db import AsyncSessionLocal from ..errors import APIValidationError, NotFoundError @@ -58,10 +63,7 @@ def _sanitize_condition_trace(trace: object) -> object: if not isinstance(trace, dict): return trace - sanitized = { - key: _sanitize_condition_trace(value) - for key, value in trace.items() - } + sanitized = {key: _sanitize_condition_trace(value) for key, value in trace.items()} raw_error = sanitized.get("error") if isinstance(raw_error, str) and raw_error: @@ -160,12 +162,17 @@ async def _load_engine_controls( namespace_key = principal.namespace_key async with AsyncSessionLocal() as db: - agent_result = await db.execute( - select(Agent).where( - Agent.name == request.agent_name, - Agent.namespace_key == namespace_key, + if timings._active_mode.get() is not None: + with timings.measured("pool_acquire"): + await db.connection() + timings.count_query("agent") + with timings.measured("agent_query"): + agent_result = await db.execute( + select(Agent).where( + Agent.name == request.agent_name, + Agent.namespace_key == namespace_key, + ) ) - ) agent = agent_result.scalar_one_or_none() if agent is None: raise NotFoundError( @@ -187,6 +194,66 @@ async def _load_engine_controls( return [ControlAdapter(c.id, c.name, c.control) for c in runtime_controls] +_authorize_runtime = require_operation(Operation.RUNTIME_USE, context_builder=_evaluation_context) + + +async def _diagnostic_authorize(request: Request) -> Principal: + """Keep the original runtime authorization on every request, including hits.""" + # Authorization duration is captured without labeling caller or target identity. + started = time.perf_counter() + principal = await _authorize_runtime(request) + request.state.lookup_authorization_seconds = time.perf_counter() - started + return principal + + +async def _load_with_diagnostic( + request: EvaluationRequest, + principal: Principal, + mode: timings.Mode, + generation: str | None = None, +) -> list[ControlAdapter]: + fixture = diagnostic.settings.fixture( + principal.namespace_key, + request.agent_name, + request.target_type, + request.target_id, + generation, + ) + if fixture is None: + raise HTTPException(status_code=400, detail="Diagnostic fixture unavailable or expired") + + async def loader() -> list[diagnostic.RuntimeControl]: + controls = await _load_engine_controls(request, principal) + return [diagnostic.RuntimeControl(c.id, c.name, c.control) for c in controls] + + if mode == "off": + timings._EVENTS.labels("off", "bypass").inc() + controls = await loader() + if diagnostic.definition_digest(controls) != fixture.controls_sha256: + raise HTTPException(status_code=409, detail="Frozen diagnostic fixture changed") + else: + key = diagnostic.CacheKey( + fixture.namespace, + fixture.agent, + fixture.target_type, + fixture.target_id, + fixture.generation, + fixture.controls_sha256, + ) + expiry = diagnostic.settings.expires_at + assert expiry is not None + ttl = min(diagnostic.settings.ttl_seconds, (expiry - datetime.now(UTC)).total_seconds()) + try: + controls = await diagnostic.cache.load(key, loader, ttl=ttl) + except ValueError as exc: + raise HTTPException( + status_code=409, detail="Frozen diagnostic fixture changed" + ) from exc + if diagnostic.definition_digest(controls) != fixture.controls_sha256: + raise HTTPException(status_code=409, detail="Frozen diagnostic fixture changed") + return [ControlAdapter(c.id, c.name, c.control) for c in controls] + + @router.post( "", response_model=EvaluationResponse, @@ -195,9 +262,8 @@ async def _load_engine_controls( ) async def evaluate( request: EvaluationRequest, - principal: Principal = Depends( - require_operation(Operation.RUNTIME_USE, context_builder=_evaluation_context) - ), + http_request: Request, + principal: Principal = Depends(_diagnostic_authorize), ) -> EvaluationResponse: """Analyze content for safety and control violations. @@ -213,10 +279,45 @@ async def evaluate( on the server; SDKs reconstruct and emit those events separately through the observability ingestion endpoint. """ - engine_controls = await _load_engine_controls(request, principal) + raw_mode = http_request.headers.get("X-LT-Control-Lookup-Cache") + if raw_mode not in (None, "off", "on"): + raise HTTPException(status_code=400, detail="Invalid diagnostic cache mode") + generation = http_request.headers.get("X-LT-Control-Lookup-Generation") + if (raw_mode is None) != (generation is None): + raise HTTPException( + status_code=400, detail="Diagnostic mode and generation required together" + ) + mode = cast(timings.Mode | None, raw_mode) + token = timings._active_mode.set(mode) + try: + if mode is not None: + timings._DURATION.labels(mode, "authorization").observe( + http_request.state.lookup_authorization_seconds + ) + with timings.measured("handler"): + return await _execute_evaluation(request, principal, mode, generation) + finally: + timings._active_mode.reset(token) + + +async def _execute_evaluation( + request: EvaluationRequest, + principal: Principal, + mode: timings.Mode | None, + generation: str | None = None, +) -> EvaluationResponse: + with timings.measured("control_load"): + engine_controls = ( + await _load_engine_controls(request, principal) + if mode is None + else await _load_with_diagnostic(request, principal, mode, generation) + ) engine = ControlEngine(engine_controls) try: - raw_response = await engine.process(request, extensions=_principal_extensions(principal)) + with timings.measured("engine"): + raw_response = await engine.process( + request, extensions=_principal_extensions(principal) + ) except ValueError: _logger.exception("Evaluation failed due to invalid configuration or input") raise APIValidationError( diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index effa3e1a..b8c2dc2c 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -25,6 +25,7 @@ ) from .config import observability_settings, settings from .db import AsyncSessionLocal, async_engine +from .diagnostic_freeze import DiagnosticFixtureFreeze from .endpoints.agents import router as agent_router from .endpoints.auth import router as auth_router from .endpoints.control_bindings import router as control_binding_router @@ -268,6 +269,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: ) add_prometheus_metrics(app, settings.prometheus_metrics_prefix) + +# The diagnostic image is allowed only on the dedicated owned LT stack. +app.add_middleware( + DiagnosticFixtureFreeze, api_base=f"{settings.api_prefix}/{settings.api_version}" +) app.add_middleware(RequestCorrelationMiddleware) # Configure CORS diff --git a/server/src/agent_control_server/services/controls.py b/server/src/agent_control_server/services/controls.py index afe38bf6..33f7efe1 100644 --- a/server/src/agent_control_server/services/controls.py +++ b/server/src/agent_control_server/services/controls.py @@ -18,6 +18,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import Select +from ..diagnostic_metrics import count_query, measured from ..errors import APIValidationError, NotFoundError from ..models import ( Control, @@ -454,16 +455,19 @@ async def list_runtime_controls_for_agent( method applies the same selection logic and parses each row into the runtime form used by the evaluation engine. """ - db_controls = await self._list_db_controls_for_agent( - agent_name, - namespace_key=namespace_key, - target_type=target_type, - target_id=target_id, - ) - return parse_runtime_controls( - db_controls, - allow_invalid_step_name_regex=allow_invalid_step_name_regex, - ) + with measured("effective_query"): + count_query("effective") + db_controls = await self._list_db_controls_for_agent( + agent_name, + namespace_key=namespace_key, + target_type=target_type, + target_id=target_id, + ) + with measured("parse"): + return parse_runtime_controls( + db_controls, + allow_invalid_step_name_regex=allow_invalid_step_name_regex, + ) async def list_controls_page( self, diff --git a/server/tests/test_diagnostic_lookup.py b/server/tests/test_diagnostic_lookup.py new file mode 100644 index 00000000..645db3d7 --- /dev/null +++ b/server/tests/test_diagnostic_lookup.py @@ -0,0 +1,429 @@ +"""Controlled lookup diagnostic: isolation, coalescing, expiry and authorization.""" + +import asyncio +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from agent_control_models import ControlDefinitionRuntime, EvaluationRequest, Step +from fastapi.testclient import TestClient + +from agent_control_server import diagnostic_lookup as diagnostic +from agent_control_server.auth_framework import Principal +from agent_control_server.endpoints import evaluation +from agent_control_server.models import DEFAULT_NAMESPACE_KEY +from agent_control_server.services.controls import RuntimeControl + +from .utils import VALID_CONTROL_PAYLOAD, create_and_assign_policy + + +def sample() -> list[RuntimeControl]: + return [ + RuntimeControl( + 1, "synthetic", ControlDefinitionRuntime.model_validate(VALID_CONTROL_PAYLOAD) + ) + ] + + +def key() -> diagnostic.CacheKey: + return diagnostic.CacheKey( + "namespace", "agent", "log_stream", "target", "v1", diagnostic.definition_digest(sample()) + ) + + +class FakeRedis: + """Deterministic transport boundary; integration must still verify real Redis.""" + + def __init__(self) -> None: + self.now = 0.0 + self.values: dict[str, tuple[str, float]] = {} + self.reads = 0 + self.writes = 0 + self.read_error: Exception | None = None + self.write_error: Exception | None = None + + async def aclose(self) -> None: + pass + + async def get(self, name: str) -> str | None: + self.reads += 1 + if self.read_error: + raise self.read_error + item = self.values.get(name) + return item[0] if item is not None and item[1] > self.now else None + + async def set(self, name: str, value: str, *, ex: int) -> None: + self.writes += 1 + if self.write_error: + raise self.write_error + self.values[name] = (value, self.now + ex) + + +@pytest.mark.asyncio +async def test_redis_shared_across_instances_and_results_are_isolated() -> None: + # Given: two server cache instances use one Redis transport. + redis = FakeRedis() + first = diagnostic.LookupCache(lambda: redis) + second = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + # When: a second instance evaluates after the first fills Redis. + initial = await first.load(key(), loader, ttl=30) + initial[0].control.condition.evaluator.config["pattern"] = "mutated" + actual = await second.load(key(), loader, ttl=30) + # Then: shared Redis avoids another database load and does not share mutable objects. + assert loader.await_count == 1 + assert redis.reads == 3 and redis.writes == 1 + assert actual[0].control.condition.evaluator.config["pattern"] == "x" + + +@pytest.mark.asyncio +async def test_redis_expiry_namespace_target_and_generation_separation() -> None: + # Given: a deterministic Redis TTL, no local result cache. + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + first = key() + # When: namespace/target/generation changes and an old key expires. + for other in ( + first, + replace(first, namespace="other"), + replace(first, target_id="other"), + replace(first, generation="v2"), + first, + ): + await cache.load(other, loader, ttl=30) + redis.now = 30 + await cache.load(first, loader, ttl=30) + # Then: exact identities share only their own unexpired definition. + assert loader.await_count == 5 + # Five misses recheck under the miss lock; the one warm hit reads once. + assert redis.reads == 11 + + +@pytest.mark.asyncio +async def test_redis_corrupt_read_and_unavailability_fall_back_without_stale_data() -> None: + # Given: a known cache entry followed by corrupt or unavailable Redis. + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + await cache.load(key(), loader, ttl=30) + redis.values[diagnostic.redis_key(key())] = ("invalid-json", 30) + # When: corrupt data then a connection error occurs. + corrupt = await cache.load(key(), loader, ttl=30) + redis.read_error = ConnectionError("synthetic") + unavailable = await cache.load(key(), loader, ttl=30) + # Then: both use the database loader, never stale or corrupt cached controls. + assert loader.await_count == 3 + assert corrupt == unavailable == sample() + + +@pytest.mark.asyncio +async def test_redis_write_failure_does_not_break_evaluation() -> None: + # Given: Redis accepts reads but rejects writes. + redis = FakeRedis() + redis.write_error = ConnectionError("synthetic") + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + # When: two requests encounter a cold cache. + actual = [await cache.load(key(), loader, ttl=30) for _ in range(2)] + # Then: every evaluation uses valid DB data and the cache remains cold. + assert loader.await_count == 2 + assert actual == [sample(), sample()] + + +@pytest.mark.asyncio +async def test_cancelled_failed_and_mismatched_db_loads_never_reach_redis() -> None: + # Given: the internal loader seam yields controlled database failures. + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock( + side_effect=[RuntimeError("synthetic"), asyncio.CancelledError(), [], sample()] + ) + # When: each failed load precedes one valid load. + for error in (RuntimeError, asyncio.CancelledError, ValueError): + with pytest.raises(error): + await cache.load(key(), loader, ttl=30) + result = await cache.load(key(), loader, ttl=30) + # Then: Redis receives only the valid serialized definitions. + assert result == sample() + assert redis.writes == 1 + + +def test_default_off_expiry_and_allowlist() -> None: + # Given: exact frozen fixture plus an explicit expiry. + fixture = diagnostic.FrozenFixture( + namespace="n", + agent="synthetic-agent", + target_type="t", + target_id="id", + generation="v1", + controls_sha256=key().digest, + ) + default = diagnostic.DiagnosticSettings(_env_file=None) + config = diagnostic.DiagnosticSettings( + enabled=True, fixtures=[fixture], expires_at=datetime.now(UTC) + timedelta(hours=1) + ) + # When/Then: disabled, expired, foreign namespace/target and naive time cannot opt in. + assert default.fixture("n", "synthetic-agent", "t", "id") is None + assert config.fixture("n", "synthetic-agent", "t", "id") == fixture + assert config.fixture("foreign", "synthetic-agent", "t", "id") is None + assert config.fixture("n", "synthetic-agent", "t", "foreign") is None + config.expires_at = datetime.now(UTC) - timedelta(seconds=1) + assert config.fixture("n", "synthetic-agent", "t", "id") is None + config.expires_at = datetime(2099, 1, 1) + assert config.fixture("n", "synthetic-agent", "t", "id") is None + + +@pytest.mark.asyncio +async def test_off_always_loads_and_on_reuses_parsed_controls(monkeypatch) -> None: + # Given: explicit internal loader seam avoids an unrelated engine/model call. + controls = sample() + fixture = diagnostic.FrozenFixture( + namespace="n", + agent="synthetic-agent", + target_type="t", + target_id="id", + generation="v1", + controls_sha256=diagnostic.definition_digest(controls), + ) + config = diagnostic.DiagnosticSettings( + enabled=True, fixtures=[fixture], expires_at=datetime.now(UTC) + timedelta(hours=1) + ) + monkeypatch.setattr(diagnostic, "settings", config) + redis = FakeRedis() + monkeypatch.setattr(diagnostic, "cache", diagnostic.LookupCache(lambda: redis)) + loader = AsyncMock( + side_effect=lambda *_: [ + evaluation.ControlAdapter(c.id, c.name, c.control) for c in sample() + ] + ) + monkeypatch.setattr(evaluation, "_load_engine_controls", loader) + request = EvaluationRequest( + agent_name="synthetic-agent", + target_type="t", + target_id="id", + stage="pre", + step=Step(type="llm", name="s", input="x"), + ) + principal = Principal(namespace_key="n") + # When: each mode receives two identical requests. + off = [await evaluation._load_with_diagnostic(request, principal, "off") for _ in range(2)] + on = [await evaluation._load_with_diagnostic(request, principal, "on") for _ in range(2)] + # Then: only ON reuses a load; semantic definitions are identical. + assert loader.await_count == 3 + assert off == on + + +def test_cached_http_evaluation_still_authorizes_every_call( + client: TestClient, monkeypatch +) -> None: + # Given: real HTTP-created fixture, then freeze its exact resolved definition digest. + agent, _ = create_and_assign_policy(client) + request = EvaluationRequest( + agent_name=agent, + target_type="log_stream", + target_id="synthetic", + stage="pre", + step=Step(type="llm", name="s", input="x"), + ) + principal = Principal(namespace_key=DEFAULT_NAMESPACE_KEY) + controls = asyncio.run(evaluation._load_engine_controls(request, principal)) + runtime = [RuntimeControl(c.id, c.name, c.control) for c in controls] + fixture = diagnostic.FrozenFixture( + namespace=DEFAULT_NAMESPACE_KEY, + agent=agent, + target_type="log_stream", + target_id="synthetic", + generation="v1", + controls_sha256=diagnostic.definition_digest(runtime), + ) + monkeypatch.setattr( + diagnostic, + "settings", + diagnostic.DiagnosticSettings( + enabled=True, fixtures=[fixture], expires_at=datetime.now(UTC) + timedelta(hours=1) + ), + ) + redis = FakeRedis() + monkeypatch.setattr(diagnostic, "cache", diagnostic.LookupCache(lambda: redis)) + headers = {"X-LT-Control-Lookup-Cache": "on", "X-LT-Control-Lookup-Generation": "v1"} + body = request.model_dump(mode="json") + # When: warm the cache, then use an invalid credential on that identical request. + first = client.post("/api/v1/evaluation", json=body, headers=headers) + second = client.post("/api/v1/evaluation", json=body, headers=headers) + denied = client.post( + "/api/v1/evaluation", json=body, headers={**headers, "X-API-Key": "invalid-synthetic"} + ) + # Then: valid requests agree but a cache hit never bypasses authorization. + assert first.status_code == second.status_code == 200 + assert first.json()["is_safe"] == second.json()["is_safe"] is False + assert denied.status_code in (401, 403) + + +def test_diagnostic_header_without_allowlist_is_rejected(client: TestClient) -> None: + # Given: the default-OFF configuration and an otherwise valid authenticated request. + body = { + "agent_name": "not-registered", + "stage": "pre", + "step": {"type": "llm", "name": "s", "input": "x"}, + } + # When: attempting to opt in without server-side fixture authorization. + response = client.post( + "/api/v1/evaluation", json=body, headers={"X-LT-Control-Lookup-Cache": "on"} + ) + # Then: reject rather than pretending the request was a cache trial. + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_redis_timeout_uses_database_and_is_counted(monkeypatch) -> None: + # Given: a transport waits forever, bounded by the diagnostic's request deadline. + from agent_control_server import diagnostic_metrics as metrics + + class SlowRedis(FakeRedis): + async def get(self, name: str) -> str | None: + await asyncio.Event().wait() + return None + + slow = SlowRedis() + cache = diagnostic.LookupCache(lambda: slow) + loader = AsyncMock(side_effect=lambda: sample()) + monkeypatch.setattr(diagnostic.settings, "redis_timeout_seconds", 0.001) + before = metrics._EVENTS.labels("on", "redis_read_fallback")._value.get() + # When: the Redis read exceeds its own deadline. + result = await cache.load(key(), loader, ttl=30) + # Then: the evaluation keeps valid DB semantics and exposes the fallback. + assert result == sample() + assert loader.await_count == 1 + assert metrics._EVENTS.labels("on", "redis_read_fallback")._value.get() == before + 1 + + +@pytest.mark.asyncio +async def test_redis_stages_have_numeric_metrics_without_identifiers() -> None: + # Given: an opted-in diagnostic request and an isolated fixed transport. + from agent_control_server import diagnostic_metrics as metrics + + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + mode = metrics._active_mode.set("on") + # When: the request first fills, then reads Redis. + try: + await cache.load(key(), loader, ttl=30) + await cache.load(key(), loader, ttl=30) + finally: + metrics._active_mode.reset(mode) + # Then: all stages use bounded mode/stage labels, with no fixture identifiers. + samples = [s for family in metrics._DURATION.collect() for s in family.samples] + assert {"redis_get", "redis_set", "cache_encode", "cache_decode"} <= { + s.labels["stage"] for s in samples + } + assert all(set(s.labels) <= {"mode", "stage", "le"} for s in samples) + + +@pytest.mark.parametrize( + "path,method", + [ + ("/api/v1/controls", "PUT"), + ("/api/v1/controls/1", "DELETE"), + ("/api/v1/control-bindings/1", "PATCH"), + ("/api/v1/policies", "PUT"), + ("/api/v1/agents/initAgent", "POST"), + ], +) +def test_diagnostic_freezes_configuration_writes_until_disabled( + client: TestClient, + monkeypatch, + path: str, + method: str, +) -> None: + # Given: the owned dedicated app is explicitly frozen, even after cache expiry. + monkeypatch.setattr(diagnostic.settings, "enabled", True) + monkeypatch.setattr(diagnostic.settings, "expires_at", datetime(2000, 1, 1, tzinfo=UTC)) + # When: a configuration mutation is attempted through any management surface. + response = client.request(method, path, json={}) + # Then: reject it without writing, and leave read/health routes available. + assert response.status_code == 409 + assert "frozen" in response.json()["detail"] + assert client.get("/health").status_code == 200 + + +def test_missing_pair_generation_rejected(client: TestClient) -> None: + # Given: a cache-mode header without an explicit cold-start pair identity. + body = { + "agent_name": "not-registered", + "stage": "pre", + "step": {"type": "llm", "name": "s", "input": "x"}, + } + # When: evaluation attempts diagnostic mode without the paired generation. + response = client.post( + "/api/v1/evaluation", json=body, headers={"X-LT-Control-Lookup-Cache": "on"} + ) + # Then: no silently reused generation is accepted. + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_cache_hits_do_not_serialize_behind_fixture_lock() -> None: + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + await cache.load(key(), loader, ttl=30) + original = redis.get + entered = 0 + both_started = asyncio.Event() + + async def overlapping_get(name: str) -> str | None: + nonlocal entered + entered += 1 + if entered == 2: + both_started.set() + await asyncio.wait_for(both_started.wait(), timeout=1) + return await original(name) + + redis.get = overlapping_get # type: ignore[method-assign] + results = await asyncio.gather(*(cache.load(key(), loader, ttl=30) for _ in range(2))) + assert entered == 2 + assert len(results) == 2 + assert loader.await_count == 1 + + +@pytest.mark.asyncio +async def test_snapshot_age_records_hits_and_invalid_timestamp_falls_back(monkeypatch) -> None: + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + loader = AsyncMock(side_effect=lambda: sample()) + monkeypatch.setattr(diagnostic, "time", SimpleNamespace(time=lambda: 1000.0)) + await cache.load(key(), loader, ttl=30) + age_metric = diagnostic._SNAPSHOT_AGE.labels("on") + before = age_metric._sum.get() + monkeypatch.setattr(diagnostic, "time", SimpleNamespace(time=lambda: 1007.0)) + await cache.load(key(), loader, ttl=30) + assert age_metric._sum.get() - before == 7 + assert loader.await_count == 1 + # A future snapshot is not silently treated as a valid old definition. + monkeypatch.setattr(diagnostic, "time", SimpleNamespace(time=lambda: 990.0)) + await cache.load(key(), loader, ttl=30) + assert loader.await_count == 2 + + +@pytest.mark.asyncio +async def test_cancelled_miss_releases_lock() -> None: + redis = FakeRedis() + cache = diagnostic.LookupCache(lambda: redis) + started = asyncio.Event() + + async def slow_loader(): + started.set() + await asyncio.Event().wait() + return sample() + + pending = asyncio.create_task(cache.load(key(), slow_loader, ttl=30)) + await started.wait() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + loader = AsyncMock(side_effect=lambda: sample()) + assert await asyncio.wait_for(cache.load(key(), loader, ttl=30), timeout=1) == sample()