|
| 1 | +"""Public engine contract tests for per-control execution durations.""" |
| 2 | + |
| 3 | +import asyncio |
| 4 | +from dataclasses import dataclass |
| 5 | +from typing import Any, Literal |
| 6 | +from unittest.mock import AsyncMock, Mock |
| 7 | + |
| 8 | +import pytest |
| 9 | +from agent_control_engine.core import ControlEngine |
| 10 | +from agent_control_evaluators import Evaluator, EvaluatorMetadata, register_evaluator |
| 11 | +from agent_control_models import ControlDefinition, EvaluationRequest, EvaluatorResult, Step |
| 12 | +from pydantic import BaseModel |
| 13 | + |
| 14 | + |
| 15 | +class DurationConfig(BaseModel): |
| 16 | + matched: bool = False |
| 17 | + raises_error: bool = False |
| 18 | + blocks: bool = False |
| 19 | + |
| 20 | + |
| 21 | +class DurationEvaluator(Evaluator[DurationConfig]): |
| 22 | + metadata = EvaluatorMetadata( |
| 23 | + name="test-duration", |
| 24 | + version="1.0.0", |
| 25 | + description="Returns a configured outcome or waits for cancellation", |
| 26 | + ) |
| 27 | + config_model = DurationConfig |
| 28 | + |
| 29 | + async def evaluate(self, data: Any) -> EvaluatorResult: |
| 30 | + # Yield once so concurrent controls all start before one completes. |
| 31 | + await asyncio.sleep(0) |
| 32 | + if self.config.blocks: |
| 33 | + await asyncio.Event().wait() |
| 34 | + if self.config.raises_error: |
| 35 | + raise RuntimeError("evaluation failed") |
| 36 | + return EvaluatorResult(matched=self.config.matched, confidence=1.0) |
| 37 | + |
| 38 | + |
| 39 | +@dataclass |
| 40 | +class DurationControl: |
| 41 | + id: int |
| 42 | + name: str |
| 43 | + control: ControlDefinition |
| 44 | + |
| 45 | + |
| 46 | +@pytest.fixture(autouse=True) |
| 47 | +def register_duration_evaluator() -> None: |
| 48 | + register_evaluator(DurationEvaluator) |
| 49 | + |
| 50 | + |
| 51 | +@pytest.fixture |
| 52 | +def evaluation_request() -> EvaluationRequest: |
| 53 | + return EvaluationRequest( |
| 54 | + agent_name="duration-test-agent", |
| 55 | + step=Step(type="llm", name="test-step", input="test", output=None), |
| 56 | + stage="pre", |
| 57 | + ) |
| 58 | + |
| 59 | + |
| 60 | +def make_control( |
| 61 | + control_id: int = 1, |
| 62 | + *, |
| 63 | + context: Literal["server", "sdk"] = "server", |
| 64 | + matched: bool = False, |
| 65 | + raises_error: bool = False, |
| 66 | + blocks: bool = False, |
| 67 | + decision: Literal["observe", "deny"] = "observe", |
| 68 | +) -> DurationControl: |
| 69 | + return DurationControl( |
| 70 | + id=control_id, |
| 71 | + name=f"duration-control-{control_id}", |
| 72 | + control=ControlDefinition( |
| 73 | + description="Duration test control", |
| 74 | + enabled=True, |
| 75 | + execution=context, |
| 76 | + scope={"stages": ["pre"], "step_types": ["llm"]}, |
| 77 | + condition={ |
| 78 | + "selector": {"path": "input"}, |
| 79 | + "evaluator": { |
| 80 | + "name": "test-duration", |
| 81 | + "config": { |
| 82 | + "matched": matched, |
| 83 | + "raises_error": raises_error, |
| 84 | + "blocks": blocks, |
| 85 | + }, |
| 86 | + }, |
| 87 | + }, |
| 88 | + action={"decision": decision}, |
| 89 | + ), |
| 90 | + ) |
| 91 | + |
| 92 | + |
| 93 | +@pytest.mark.asyncio |
| 94 | +@pytest.mark.parametrize("context", ["server", "sdk"]) |
| 95 | +@pytest.mark.parametrize("category", ["matches", "non_matches", "errors"]) |
| 96 | +@pytest.mark.parametrize("elapsed_seconds", [0.0, 0.0125]) |
| 97 | +async def test_completed_controls_report_elapsed_milliseconds( |
| 98 | + monkeypatch: pytest.MonkeyPatch, |
| 99 | + evaluation_request: EvaluationRequest, |
| 100 | + context: Literal["server", "sdk"], |
| 101 | + category: str, |
| 102 | + elapsed_seconds: float, |
| 103 | +) -> None: |
| 104 | + # Given: a completed control and a controlled clock instead of wall-clock assertions |
| 105 | + control = make_control( |
| 106 | + context=context, |
| 107 | + matched=category == "matches", |
| 108 | + raises_error=category == "errors", |
| 109 | + ) |
| 110 | + clock = Mock(side_effect=[10.0, 10.0 + elapsed_seconds]) |
| 111 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 112 | + |
| 113 | + # When: evaluating via the public engine entrypoint |
| 114 | + response = await ControlEngine([control], context=context).process(evaluation_request) |
| 115 | + |
| 116 | + # Then: all completed outcome categories retain their own duration, including zero |
| 117 | + results = getattr(response, category) |
| 118 | + assert results is not None |
| 119 | + assert len(results) == 1 |
| 120 | + assert results[0].execution_duration_ms == pytest.approx(elapsed_seconds * 1000) |
| 121 | + assert results[0].result.matched is (category == "matches") |
| 122 | + assert bool(results[0].result.error) is (category == "errors") |
| 123 | + assert response.is_safe is True |
| 124 | + |
| 125 | + |
| 126 | +@pytest.mark.asyncio |
| 127 | +async def test_parallel_controls_keep_individual_durations( |
| 128 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 129 | +) -> None: |
| 130 | + # Given: parallel controls start together and complete at different clock readings |
| 131 | + clock = Mock(side_effect=[10.0, 10.0, 10.0125, 10.025]) |
| 132 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 133 | + controls = [make_control(1), make_control(2)] |
| 134 | + |
| 135 | + # When: evaluating both controls in the same engine request |
| 136 | + response = await ControlEngine(controls).process(evaluation_request) |
| 137 | + |
| 138 | + # Then: each control reports its own elapsed time rather than the batch duration |
| 139 | + assert response.non_matches is not None |
| 140 | + durations = {result.control_id: result.execution_duration_ms for result in response.non_matches} |
| 141 | + assert durations == pytest.approx({1: 12.5, 2: 25.0}) |
| 142 | + |
| 143 | + |
| 144 | +@pytest.mark.asyncio |
| 145 | +async def test_condition_traversal_error_retains_elapsed_duration( |
| 146 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 147 | +) -> None: |
| 148 | + # Given: an internal traversal failure forced to exercise the outer error handler |
| 149 | + engine = ControlEngine([make_control(decision="deny")]) |
| 150 | + monkeypatch.setattr( |
| 151 | + engine, "_evaluate_condition", AsyncMock(side_effect=ValueError("condition failed")) |
| 152 | + ) |
| 153 | + clock = Mock(side_effect=[10.0, 10.007]) |
| 154 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 155 | + |
| 156 | + # When: evaluating via the public engine entrypoint |
| 157 | + response = await engine.process(evaluation_request) |
| 158 | + |
| 159 | + # Then: the completed error is timed and still fails closed |
| 160 | + assert response.is_safe is False |
| 161 | + assert response.errors is not None |
| 162 | + assert len(response.errors) == 1 |
| 163 | + assert response.errors[0].execution_duration_ms == pytest.approx(7.0) |
| 164 | + assert response.errors[0].result.error == "ValueError: condition failed" |
| 165 | + |
| 166 | + |
| 167 | +@pytest.mark.asyncio |
| 168 | +async def test_timed_out_evaluation_retains_elapsed_duration( |
| 169 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 170 | +) -> None: |
| 171 | + # Given: a blocking evaluator with a bounded timeout and a controlled elapsed clock |
| 172 | + monkeypatch.setattr(DurationEvaluator, "get_timeout_seconds", lambda self: 0.001) |
| 173 | + clock = Mock(side_effect=[10.0, 10.001]) |
| 174 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 175 | + |
| 176 | + # When: the evaluator times out without producing a normal result |
| 177 | + response = await ControlEngine([make_control(blocks=True, decision="deny")]).process( |
| 178 | + evaluation_request |
| 179 | + ) |
| 180 | + |
| 181 | + # Then: the timeout error is a completed evaluation with an elapsed duration |
| 182 | + assert response.is_safe is False |
| 183 | + assert response.errors is not None |
| 184 | + assert len(response.errors) == 1 |
| 185 | + assert response.errors[0].execution_duration_ms == pytest.approx(1.0) |
| 186 | + assert response.errors[0].result.error is not None |
| 187 | + assert response.errors[0].result.error.startswith("TimeoutError:") |
| 188 | + |
| 189 | + |
| 190 | +@pytest.mark.asyncio |
| 191 | +async def test_cancelled_controls_do_not_report_partial_durations( |
| 192 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 193 | +) -> None: |
| 194 | + # Given: a blocking control that is cancelled by a completed deny control |
| 195 | + clock = Mock(side_effect=[10.0, 10.0, 10.020]) |
| 196 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 197 | + controls = [make_control(1, blocks=True), make_control(2, matched=True, decision="deny")] |
| 198 | + |
| 199 | + # When: a deny match cancels the unfinished evaluation |
| 200 | + response = await ControlEngine(controls).process(evaluation_request) |
| 201 | + |
| 202 | + # Then: only the completed deny is returned, with no fabricated cancelled result |
| 203 | + assert response.is_safe is False |
| 204 | + assert response.matches is not None |
| 205 | + assert len(response.matches) == 1 |
| 206 | + assert response.matches[0].control_id == 2 |
| 207 | + assert response.matches[0].execution_duration_ms == pytest.approx(20.0) |
| 208 | + assert response.non_matches is None |
| 209 | + assert response.errors is None |
| 210 | + assert clock.call_count == 3 |
| 211 | + |
| 212 | + |
| 213 | +@pytest.mark.asyncio |
| 214 | +async def test_skipped_control_does_not_start_duration_measurement( |
| 215 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 216 | +) -> None: |
| 217 | + # Given: a disabled control that must not evaluate |
| 218 | + control = make_control() |
| 219 | + control.control.enabled = False |
| 220 | + clock = Mock(side_effect=AssertionError("skipped control must not be timed")) |
| 221 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 222 | + |
| 223 | + # When: processing a request without applicable controls |
| 224 | + response = await ControlEngine([control]).process(evaluation_request) |
| 225 | + |
| 226 | + # Then: there are no results or clock reads for the skipped control |
| 227 | + assert response.matches is None |
| 228 | + assert response.non_matches is None |
| 229 | + assert response.errors is None |
| 230 | + clock.assert_not_called() |
| 231 | + |
| 232 | + |
| 233 | +@pytest.mark.asyncio |
| 234 | +async def test_precheck_error_has_no_execution_duration( |
| 235 | + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest |
| 236 | +) -> None: |
| 237 | + # Given: validation is bypassed to exercise the defensive invalid-regex precheck |
| 238 | + control = make_control() |
| 239 | + payload = control.control.model_dump() |
| 240 | + payload["scope"]["step_name_regex"] = "[invalid(regex" |
| 241 | + control.control = ControlDefinition.model_validate( |
| 242 | + payload, context={"allow_invalid_step_name_regex": True} |
| 243 | + ) |
| 244 | + clock = Mock(side_effect=AssertionError("precheck failure must not be timed")) |
| 245 | + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) |
| 246 | + |
| 247 | + # When: processing a control that fails before evaluation |
| 248 | + response = await ControlEngine([control]).process(evaluation_request) |
| 249 | + |
| 250 | + # Then: the error retains an unknown duration rather than zero |
| 251 | + assert response.errors is not None |
| 252 | + assert len(response.errors) == 1 |
| 253 | + assert response.errors[0].execution_duration_ms is None |
| 254 | + clock.assert_not_called() |
0 commit comments