diff --git a/engine/src/agent_control_engine/core.py b/engine/src/agent_control_engine/core.py index ab170e82..58e270eb 100644 --- a/engine/src/agent_control_engine/core.py +++ b/engine/src/agent_control_engine/core.py @@ -7,6 +7,7 @@ import functools import logging import os +import time from collections.abc import Sequence from dataclasses import dataclass from typing import Any, Literal, Protocol @@ -201,6 +202,7 @@ class _EvalTask: item: ControlWithIdentity task: asyncio.Task[None] | None = None result: EvaluatorResult | None = None + execution_duration_ms: float | None = None @dataclass @@ -673,6 +675,7 @@ async def process( async def evaluate_control(eval_task: _EvalTask) -> None: """Evaluate a single control, respecting cancellation and timeout.""" + started_at = time.perf_counter() try: evaluation = await self._evaluate_condition( eval_task.item, @@ -701,6 +704,9 @@ async def evaluate_control(eval_task: _EvalTask) -> None: error_msg, message_prefix="Condition evaluation failed", ) + finally: + if eval_task.result is not None: + eval_task.execution_duration_ms = (time.perf_counter() - started_at) * 1000 # Create and start all tasks for eval_task in eval_tasks: @@ -753,6 +759,7 @@ async def wait_for_deny() -> None: control_name=eval_task.item.name, action=eval_task.item.control.action.decision, result=eval_task.result, + execution_duration_ms=eval_task.execution_duration_ms, steering_context=eval_task.item.control.action.steering_context, ) ) @@ -776,6 +783,7 @@ async def wait_for_deny() -> None: control_name=eval_task.item.name, action=eval_task.item.control.action.decision, result=eval_task.result, + execution_duration_ms=eval_task.execution_duration_ms, steering_context=steer_ctx, ) ) @@ -792,6 +800,7 @@ async def wait_for_deny() -> None: control_name=eval_task.item.name, action=eval_task.item.control.action.decision, result=eval_task.result, + execution_duration_ms=eval_task.execution_duration_ms, steering_context=eval_task.item.control.action.steering_context, ) ) diff --git a/engine/tests/test_execution_duration.py b/engine/tests/test_execution_duration.py new file mode 100644 index 00000000..b52c154b --- /dev/null +++ b/engine/tests/test_execution_duration.py @@ -0,0 +1,254 @@ +"""Public engine contract tests for per-control execution durations.""" + +import asyncio +from dataclasses import dataclass +from typing import Any, Literal +from unittest.mock import AsyncMock, Mock + +import pytest +from agent_control_engine.core import ControlEngine +from agent_control_evaluators import Evaluator, EvaluatorMetadata, register_evaluator +from agent_control_models import ControlDefinition, EvaluationRequest, EvaluatorResult, Step +from pydantic import BaseModel + + +class DurationConfig(BaseModel): + matched: bool = False + raises_error: bool = False + blocks: bool = False + + +class DurationEvaluator(Evaluator[DurationConfig]): + metadata = EvaluatorMetadata( + name="test-duration", + version="1.0.0", + description="Returns a configured outcome or waits for cancellation", + ) + config_model = DurationConfig + + async def evaluate(self, data: Any) -> EvaluatorResult: + # Yield once so concurrent controls all start before one completes. + await asyncio.sleep(0) + if self.config.blocks: + await asyncio.Event().wait() + if self.config.raises_error: + raise RuntimeError("evaluation failed") + return EvaluatorResult(matched=self.config.matched, confidence=1.0) + + +@dataclass +class DurationControl: + id: int + name: str + control: ControlDefinition + + +@pytest.fixture(autouse=True) +def register_duration_evaluator() -> None: + register_evaluator(DurationEvaluator) + + +@pytest.fixture +def evaluation_request() -> EvaluationRequest: + return EvaluationRequest( + agent_name="duration-test-agent", + step=Step(type="llm", name="test-step", input="test", output=None), + stage="pre", + ) + + +def make_control( + control_id: int = 1, + *, + context: Literal["server", "sdk"] = "server", + matched: bool = False, + raises_error: bool = False, + blocks: bool = False, + decision: Literal["observe", "deny"] = "observe", +) -> DurationControl: + return DurationControl( + id=control_id, + name=f"duration-control-{control_id}", + control=ControlDefinition( + description="Duration test control", + enabled=True, + execution=context, + scope={"stages": ["pre"], "step_types": ["llm"]}, + condition={ + "selector": {"path": "input"}, + "evaluator": { + "name": "test-duration", + "config": { + "matched": matched, + "raises_error": raises_error, + "blocks": blocks, + }, + }, + }, + action={"decision": decision}, + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context", ["server", "sdk"]) +@pytest.mark.parametrize("category", ["matches", "non_matches", "errors"]) +@pytest.mark.parametrize("elapsed_seconds", [0.0, 0.0125]) +async def test_completed_controls_report_elapsed_milliseconds( + monkeypatch: pytest.MonkeyPatch, + evaluation_request: EvaluationRequest, + context: Literal["server", "sdk"], + category: str, + elapsed_seconds: float, +) -> None: + # Given: a completed control and a controlled clock instead of wall-clock assertions + control = make_control( + context=context, + matched=category == "matches", + raises_error=category == "errors", + ) + clock = Mock(side_effect=[10.0, 10.0 + elapsed_seconds]) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + + # When: evaluating via the public engine entrypoint + response = await ControlEngine([control], context=context).process(evaluation_request) + + # Then: all completed outcome categories retain their own duration, including zero + results = getattr(response, category) + assert results is not None + assert len(results) == 1 + assert results[0].execution_duration_ms == pytest.approx(elapsed_seconds * 1000) + assert results[0].result.matched is (category == "matches") + assert bool(results[0].result.error) is (category == "errors") + assert response.is_safe is True + + +@pytest.mark.asyncio +async def test_parallel_controls_keep_individual_durations( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: parallel controls start together and complete at different clock readings + clock = Mock(side_effect=[10.0, 10.0, 10.0125, 10.025]) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + controls = [make_control(1), make_control(2)] + + # When: evaluating both controls in the same engine request + response = await ControlEngine(controls).process(evaluation_request) + + # Then: each control reports its own elapsed time rather than the batch duration + assert response.non_matches is not None + durations = {result.control_id: result.execution_duration_ms for result in response.non_matches} + assert durations == pytest.approx({1: 12.5, 2: 25.0}) + + +@pytest.mark.asyncio +async def test_condition_traversal_error_retains_elapsed_duration( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: an internal traversal failure forced to exercise the outer error handler + engine = ControlEngine([make_control(decision="deny")]) + monkeypatch.setattr( + engine, "_evaluate_condition", AsyncMock(side_effect=ValueError("condition failed")) + ) + clock = Mock(side_effect=[10.0, 10.007]) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + + # When: evaluating via the public engine entrypoint + response = await engine.process(evaluation_request) + + # Then: the completed error is timed and still fails closed + assert response.is_safe is False + assert response.errors is not None + assert len(response.errors) == 1 + assert response.errors[0].execution_duration_ms == pytest.approx(7.0) + assert response.errors[0].result.error == "ValueError: condition failed" + + +@pytest.mark.asyncio +async def test_timed_out_evaluation_retains_elapsed_duration( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: a blocking evaluator with a bounded timeout and a controlled elapsed clock + monkeypatch.setattr(DurationEvaluator, "get_timeout_seconds", lambda self: 0.001) + clock = Mock(side_effect=[10.0, 10.001]) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + + # When: the evaluator times out without producing a normal result + response = await ControlEngine([make_control(blocks=True, decision="deny")]).process( + evaluation_request + ) + + # Then: the timeout error is a completed evaluation with an elapsed duration + assert response.is_safe is False + assert response.errors is not None + assert len(response.errors) == 1 + assert response.errors[0].execution_duration_ms == pytest.approx(1.0) + assert response.errors[0].result.error is not None + assert response.errors[0].result.error.startswith("TimeoutError:") + + +@pytest.mark.asyncio +async def test_cancelled_controls_do_not_report_partial_durations( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: a blocking control that is cancelled by a completed deny control + clock = Mock(side_effect=[10.0, 10.0, 10.020]) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + controls = [make_control(1, blocks=True), make_control(2, matched=True, decision="deny")] + + # When: a deny match cancels the unfinished evaluation + response = await ControlEngine(controls).process(evaluation_request) + + # Then: only the completed deny is returned, with no fabricated cancelled result + assert response.is_safe is False + assert response.matches is not None + assert len(response.matches) == 1 + assert response.matches[0].control_id == 2 + assert response.matches[0].execution_duration_ms == pytest.approx(20.0) + assert response.non_matches is None + assert response.errors is None + assert clock.call_count == 3 + + +@pytest.mark.asyncio +async def test_skipped_control_does_not_start_duration_measurement( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: a disabled control that must not evaluate + control = make_control() + control.control.enabled = False + clock = Mock(side_effect=AssertionError("skipped control must not be timed")) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + + # When: processing a request without applicable controls + response = await ControlEngine([control]).process(evaluation_request) + + # Then: there are no results or clock reads for the skipped control + assert response.matches is None + assert response.non_matches is None + assert response.errors is None + clock.assert_not_called() + + +@pytest.mark.asyncio +async def test_precheck_error_has_no_execution_duration( + monkeypatch: pytest.MonkeyPatch, evaluation_request: EvaluationRequest +) -> None: + # Given: validation is bypassed to exercise the defensive invalid-regex precheck + control = make_control() + payload = control.control.model_dump() + payload["scope"]["step_name_regex"] = "[invalid(regex" + control.control = ControlDefinition.model_validate( + payload, context={"allow_invalid_step_name_regex": True} + ) + clock = Mock(side_effect=AssertionError("precheck failure must not be timed")) + monkeypatch.setattr("agent_control_engine.core.time.perf_counter", clock) + + # When: processing a control that fails before evaluation + response = await ControlEngine([control]).process(evaluation_request) + + # Then: the error retains an unknown duration rather than zero + assert response.errors is not None + assert len(response.errors) == 1 + assert response.errors[0].execution_duration_ms is None + clock.assert_not_called() diff --git a/examples/README.md b/examples/README.md index f780be53..535ce1b1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -14,6 +14,7 @@ This directory contains runnable examples for Agent Control. Each example has it | Customer Support Agent | Enterprise scenario with PII protection, prompt-injection defense, and multiple tools. | https://docs.agentcontrol.dev/examples/customer-support | | DeepEval | Build a custom evaluator using DeepEval GEval metrics. | https://docs.agentcontrol.dev/examples/deepeval | | Galileo Luna Direct | Direct Luna scorer invoke evaluation (`scorer_id` required; `scorer_label`/`scorer_version_id` optional) with a composite Agent Control condition. | `examples/galileo_luna/` | +| SDK-local timing | Compare per-control duration with whole SDK call time using synthetic local controls and no network. | `examples/sdk_local_timing/` | | LangChain SQL Agent | Protect a SQL agent from dangerous queries with server-side controls. | https://docs.agentcontrol.dev/examples/langchain-sql | | Steer Action Demo | Banking transfer agent showcasing observe, deny, and steer actions. | https://docs.agentcontrol.dev/examples/steer-action-demo | | Target Context | Bind controls to opaque external targets (e.g. `env=prod`) and let the SDK pin one target per session. | https://docs.agentcontrol.dev/examples/target-context | diff --git a/examples/sdk_local_timing/README.md b/examples/sdk_local_timing/README.md new file mode 100644 index 00000000..a91b6589 --- /dev/null +++ b/examples/sdk_local_timing/README.md @@ -0,0 +1,37 @@ +# SDK-local control timing example + +This example runs five synthetic controls through the Python SDK's real +`check_evaluation_with_local` path. It covers list allow/block, regex steer, +SQL block, and JSON validation. Each case uses one control, so its result and +duration have a clear owner. + +PR #283 measures the full control condition evaluation. It includes evaluator +work and any wait for an evaluator slot. It excludes control loading and SDK +call overhead. The separate `sdk_call` distribution includes those costs. +These values are not isolated evaluator-call timings. + +From the repository root: + +```bash +make sync +uv run --package agent-control-sdk python examples/sdk_local_timing/demo.py --iterations 100 --warmup 5 +``` + +The default output is a small per-control p50/p95 table. Pass `--json` for +counts, min, p50, p95 and max values in a machine-readable report. The +`control` number comes from `ControlMatch.execution_duration_ms`; `sdk_call` +is measured around the SDK call by this example. + +For a bounded SDK-local throughput smoke, add `--concurrency 8 --iterations 1000`. +The table then includes observed SDK calls per second for each independent +scenario. This uses one Python process and one event loop without a target +request rate; it is not a production capacity number. + +The example disables SDK observability for its run and restores the previous +settings afterward. An HTTP mock rejects all outbound requests. No account, +API key, server, Luna scorer, or GPU is needed. The warmup runs are excluded +from the reported nearest-rank p50 and p95 values. + +This example does not establish SDK capacity or server/Luna3 throughput. Those +require a configured workload, stack, acceptance thresholds and a separate +load test. diff --git a/examples/sdk_local_timing/demo.py b/examples/sdk_local_timing/demo.py new file mode 100644 index 00000000..4cce0214 --- /dev/null +++ b/examples/sdk_local_timing/demo.py @@ -0,0 +1,240 @@ +"""Compare SDK-local per-control time with the full SDK call, without network I/O.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import time +from dataclasses import dataclass +from typing import Any, Literal + +import httpx +from agent_control import AgentControlClient +from agent_control.evaluation import check_evaluation_with_local +from agent_control.settings import configure_settings, get_settings +from agent_control_models import Step + + +@dataclass(frozen=True) +class Scenario: + name: str + evaluator: str + config: dict[str, Any] + selector: str + decision: Literal["deny", "steer", "observe"] + step_input: dict[str, Any] + expected_bucket: Literal["matches", "non_matches"] + + +LIST_CONFIG = { + "values": ["delete_customer"], + "logic": "any", + "match_on": "match", + "match_mode": "exact", + "case_sensitive": False, +} + +SCENARIOS = ( + Scenario( + "list-allow", + "list", + LIST_CONFIG, + "input.action", + "deny", + {"action": "read_customer"}, + "non_matches", + ), + Scenario( + "list-block", + "list", + LIST_CONFIG, + "input.action", + "deny", + {"action": "delete_customer"}, + "matches", + ), + Scenario( + "regex-steer", + "regex", + {"pattern": "transfer"}, + "input.action", + "steer", + {"action": "transfer"}, + "matches", + ), + Scenario( + "sql-block", + "sql", + {"blocked_operations": ["DROP"]}, + "input.sql", + "deny", + {"sql": "DROP TABLE customers"}, + "matches", + ), + Scenario( + "json-valid", + "json", + {"required_fields": ["customer_id"]}, + "input.payload", + "observe", + {"payload": '{"customer_id":"synthetic-1"}'}, + "non_matches", + ), +) + + +def control_payload(scenario: Scenario, control_id: int) -> dict[str, Any]: + return { + "id": control_id, + "name": scenario.name, + "control": { + "enabled": True, + "execution": "sdk", + "scope": {"step_types": ["tool"], "stages": ["pre"], "step_names": ["timing_probe"]}, + "condition": { + "selector": {"path": scenario.selector}, + "evaluator": {"name": scenario.evaluator, "config": scenario.config}, + }, + "action": {"decision": scenario.decision}, + }, + } + + +def distribution(values: list[float]) -> dict[str, float | int]: + """Summarize a nonempty sample using nearest-rank percentiles.""" + ordered = sorted(values) + + def rank(fraction: float) -> float: + return ordered[max(0, math.ceil(fraction * len(ordered)) - 1)] + + return { + "count": len(ordered), + "min_ms": ordered[0], + "p50_ms": rank(0.50), + "p95_ms": rank(0.95), + "max_ms": ordered[-1], + } + + +def reject_network(request: httpx.Request) -> httpx.Response: + raise AssertionError(f"SDK-local demo attempted HTTP to {request.url.host}") + + +def format_table(report: dict[str, Any]) -> str: + """Render each local control's timing beside the enclosing SDK call.""" + rows = [ + f"iterations={report['iterations']} warmup={report['warmup']} " + f"concurrency={report['concurrency']}", + "scenario outcome control p50/p95 (ms) SDK call p50/p95 (ms) calls/s", + ] + for name, result in report["scenarios"].items(): + control = result["control"] + sdk_call = result["sdk_call"] + rows.append( + f"{name:<13} {result['outcome']:<12} " + f"{control['p50_ms']:>8.3f}/{control['p95_ms']:>8.3f} " + f"{sdk_call['p50_ms']:>8.3f}/{sdk_call['p95_ms']:>8.3f} " + f"{result['observed_calls_per_sec']:>8.1f}" + ) + return "\n".join(rows) + + +async def run(*, iterations: int, warmup: int, concurrency: int = 1) -> dict[str, Any]: + """Run each synthetic control locally and validate its measured result.""" + if not (1 <= iterations <= 100_000 and 0 <= warmup <= 10_000 and 1 <= concurrency <= 64): + raise ValueError("iterations must be 1..100000, warmup 0..10000, and concurrency 1..64") + + previous_settings = get_settings().model_dump() + configure_settings(observability_enabled=False) + try: + report: dict[str, Any] = { + "iterations": iterations, + "warmup": warmup, + "concurrency": concurrency, + "scenarios": {}, + } + async with AgentControlClient( + base_url="https://unused.invalid", transport=httpx.MockTransport(reject_network) + ) as client: + for control_id, scenario in enumerate(SCENARIOS, start=1): + control = control_payload(scenario, control_id) + step = Step(type="tool", name="timing_probe", input=scenario.step_input) + + async def evaluate_once() -> tuple[float, float]: + started_at = time.perf_counter() + result = await check_evaluation_with_local( + client, "sdk-local-timing-demo", step, "pre", [control] + ) + sdk_ms = (time.perf_counter() - started_at) * 1000 + bucket = getattr(result, scenario.expected_bucket) or [] + other_buckets = ( + ["non_matches"] if scenario.expected_bucket == "matches" else ["matches"] + ) + ["errors"] + if len(bucket) != 1 or any(getattr(result, name) for name in other_buckets): + raise RuntimeError(f"{scenario.name}: unexpected control outcome") + if bucket[0].control_id != control_id: + raise RuntimeError(f"{scenario.name}: wrong control result") + expected_match = scenario.expected_bucket == "matches" + if ( + bucket[0].action != scenario.decision + or bucket[0].result.matched != expected_match + or result.is_safe != (not expected_match) + ): + raise RuntimeError(f"{scenario.name}: unexpected action or safety result") + duration_ms = bucket[0].execution_duration_ms + if duration_ms is None or not math.isfinite(duration_ms) or duration_ms < 0: + raise RuntimeError(f"{scenario.name}: missing or invalid control duration") + return duration_ms, sdk_ms + + for _ in range(warmup): + await evaluate_once() + + worker_count = min(concurrency, iterations) + + async def worker(count: int) -> tuple[list[float], list[float]]: + control_samples: list[float] = [] + sdk_samples: list[float] = [] + for _ in range(count): + duration_ms, sdk_ms = await evaluate_once() + control_samples.append(duration_ms) + sdk_samples.append(sdk_ms) + return control_samples, sdk_samples + + counts = [ + iterations // worker_count + (worker_index < iterations % worker_count) + for worker_index in range(worker_count) + ] + run_started_at = time.perf_counter() + worker_results = await asyncio.gather(*(worker(count) for count in counts)) + elapsed_seconds = time.perf_counter() - run_started_at + control_times = [sample for controls, _ in worker_results for sample in controls] + sdk_times = [sample for _, calls in worker_results for sample in calls] + report["scenarios"][scenario.name] = { + "evaluator": scenario.evaluator, + "outcome": scenario.expected_bucket, + "control": distribution(control_times), + "sdk_call": distribution(sdk_times), + "observed_calls_per_sec": iterations / elapsed_seconds, + } + return report + finally: + configure_settings(**previous_settings) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--concurrency", type=int, default=1) + parser.add_argument("--json", action="store_true", help="emit machine-readable samples summary") + args = parser.parse_args() + report = asyncio.run( + run(iterations=args.iterations, warmup=args.warmup, concurrency=args.concurrency) + ) + print(json.dumps(report, indent=2) if args.json else format_table(report)) + + +if __name__ == "__main__": + main() diff --git a/examples/sdk_local_timing/test_demo.py b/examples/sdk_local_timing/test_demo.py new file mode 100644 index 00000000..b74631be --- /dev/null +++ b/examples/sdk_local_timing/test_demo.py @@ -0,0 +1,41 @@ +"""Behavior checks for the SDK-local timing example.""" + +import pytest + +from examples.sdk_local_timing.demo import distribution, format_table, run + + +def test_distribution_uses_nearest_rank() -> None: + # Given: a sample with a known upper-tail observation. + values = list(range(1, 21)) + + # When: the collector summarizes the sample. + summary = distribution(values) + + # Then: p95 is the nineteenth ordered value. + assert summary == {"count": 20, "min_ms": 1, "p50_ms": 10, "p95_ms": 19, "max_ms": 20} + + +@pytest.mark.asyncio +async def test_cases_run_locally_and_have_control_time() -> None: + # Given: an HTTP transport that rejects any outbound request. + + # When: each synthetic case runs through the SDK local path after one warmup. + report = await run(iterations=2, warmup=1, concurrency=2) + + # Then: each case returns exactly one measured control result. + assert set(report["scenarios"]) == { + "list-allow", + "list-block", + "regex-steer", + "sql-block", + "json-valid", + } + for case in report["scenarios"].values(): + assert case["control"]["count"] == 2 + assert case["control"]["min_ms"] >= 0 + assert case["sdk_call"]["min_ms"] >= case["control"]["min_ms"] + assert case["observed_calls_per_sec"] > 0 + table = format_table(report) + assert "control p50/p95 (ms)" in table + assert "regex-steer" in table diff --git a/models/README.md b/models/README.md index db98868a..fb09f055 100644 --- a/models/README.md +++ b/models/README.md @@ -17,6 +17,20 @@ agent = Agent(agent_name="support-bot", agent_description="Support agent") step = Step(type="llm", name="chat", input="hello") ``` +## Control execution duration + +Evaluation responses include an optional top-level `execution_duration_ms` on each +`ControlMatch` in `matches`, `non_matches`, and `errors`. It measures that control's +condition evaluation in wall-clock milliseconds, including the visited condition +leaves, evaluator work, and time waiting for an evaluator concurrency slot. +Control loading and request/response overhead are outside this measurement. + +Completed evaluations include a nonnegative value, including evaluator failures. +Precheck failures have no duration, and canceled evaluations produce no result. +Older responses without the field remain valid. The Python SDK copies the value +to `ControlExecutionEvent.execution_duration_ms`; the TypeScript SDK exposes it as +`executionDurationMs`. + ## Tests ```bash diff --git a/models/src/agent_control_models/controls.py b/models/src/agent_control_models/controls.py index 1e2bb9e9..46e8fd8e 100644 --- a/models/src/agent_control_models/controls.py +++ b/models/src/agent_control_models/controls.py @@ -957,6 +957,11 @@ class ControlMatch(BaseModel): result: EvaluatorResult = Field( ..., description="Evaluator result (confidence, message, metadata)" ) + execution_duration_ms: float | None = Field( + default=None, + ge=0, + description="Control evaluation wall-clock duration in milliseconds, if executed", + ) steering_context: SteeringContext | None = Field( None, description="Steering context for steer actions if configured" diff --git a/models/tests/test_control_match_duration.py b/models/tests/test_control_match_duration.py new file mode 100644 index 00000000..6b3e13f7 --- /dev/null +++ b/models/tests/test_control_match_duration.py @@ -0,0 +1,58 @@ +"""Control execution duration response compatibility.""" + +import pytest +from agent_control_models.controls import ControlMatch +from pydantic import ValidationError + + +@pytest.mark.parametrize("duration_ms", [None, 0.0, 12.5]) +def test_control_match_preserves_optional_execution_duration(duration_ms: float | None) -> None: + # Given: a response with a recorded duration or an unevaluated control. + payload = { + "control_id": 1, + "control_name": "Example control", + "action": "deny", + "result": {"matched": False, "confidence": 0.0}, + "execution_duration_ms": duration_ms, + } + + # When: parsing and serializing the response model. + match = ControlMatch.model_validate(payload) + reconstructed = ControlMatch.model_validate_json(match.model_dump_json()) + + # Then: the duration stays top-level and zero is preserved. + assert reconstructed.execution_duration_ms == duration_ms + assert match.model_dump()["execution_duration_ms"] == duration_ms + assert "execution_duration_ms" not in match.result.model_dump() + + +def test_control_match_accepts_legacy_response_without_duration() -> None: + # Given: an older server response with no duration field. + payload = { + "control_id": 1, + "control_name": "Example control", + "action": "deny", + "result": {"matched": False, "confidence": 0.0}, + } + + # When: parsing the response model. + match = ControlMatch.model_validate(payload) + + # Then: no duration is fabricated for the older response. + assert match.execution_duration_ms is None + + +def test_control_match_rejects_negative_execution_duration() -> None: + # Given: a response containing an invalid elapsed time. + payload = { + "control_id": 1, + "control_name": "Example control", + "action": "deny", + "result": {"matched": False, "confidence": 0.0}, + "execution_duration_ms": -1, + } + + # When: parsing the response model. + with pytest.raises(ValidationError, match="greater than or equal to 0"): + ControlMatch.model_validate(payload) + # Then: the response follows the same nonnegative duration contract as events. diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index fee78333..38249e98 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -339,6 +339,7 @@ async def _evaluate( "metadata": m.result.metadata, }, "control_execution_id": m.control_execution_id, + "execution_duration_ms": m.execution_duration_ms, "steering_context": ( {"message": m.steering_context.message} if m.steering_context else None @@ -360,6 +361,7 @@ async def _evaluate( "metadata": e.result.metadata, }, "control_execution_id": e.control_execution_id, + "execution_duration_ms": e.execution_duration_ms, "steering_context": ( {"message": e.steering_context.message} if e.steering_context else None @@ -381,6 +383,7 @@ async def _evaluate( "metadata": nm.result.metadata, }, "control_execution_id": nm.control_execution_id, + "execution_duration_ms": nm.execution_duration_ms, "steering_context": ( {"message": nm.steering_context.message} if nm.steering_context else None @@ -707,7 +710,7 @@ def _log_single_control( matched=matched, action=normalize_action(str(control_data.get("action", "observe"))), confidence=control_data.get("result", {}).get("confidence", 0.0), - duration_ms=control_data.get("result", {}).get("execution_duration_ms"), + duration_ms=control_data.get("execution_duration_ms"), control_execution_id=control_data.get("control_execution_id"), check_stage=check_stage, ) diff --git a/sdks/python/src/agent_control/evaluation_events.py b/sdks/python/src/agent_control/evaluation_events.py index 46c9ac15..ac3f48b2 100644 --- a/sdks/python/src/agent_control/evaluation_events.py +++ b/sdks/python/src/agent_control/evaluation_events.py @@ -141,6 +141,7 @@ def _build_events_for_matches( action=match.action, matched=matched, confidence=match.result.confidence, + execution_duration_ms=match.execution_duration_ms, timestamp=now, evaluator_name=evaluator_name, selector_path=selector_path, diff --git a/sdks/python/tests/test_execution_duration.py b/sdks/python/tests/test_execution_duration.py new file mode 100644 index 00000000..54426d76 --- /dev/null +++ b/sdks/python/tests/test_execution_duration.py @@ -0,0 +1,197 @@ +"""Per-control duration survives SDK response parsing and event reconstruction.""" + +from collections.abc import Iterator, Sequence +from functools import partial +from typing import Any, Literal +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from agent_control_models import ControlExecutionEvent, Step +from agent_control_telemetry import REGISTERED_CONTROL_EVENT_SINK_NAME +from agent_control_telemetry.sinks import BaseControlEventSink, SinkResult + +from agent_control import ( + AgentControlClient, + control, + register_control_event_sink, + unregister_control_event_sink, +) +from agent_control._state import state +from agent_control.evaluation import check_evaluation, check_evaluation_with_local +from agent_control.otel_sink import control_event_to_otel_span +from agent_control.settings import configure_settings, get_settings + +ResultCategory = Literal["matches", "errors", "non_matches"] + + +class RecordingSink(BaseControlEventSink): + """Capture the events delivered through the public sink API.""" + + def __init__(self) -> None: + self.events: list[ControlExecutionEvent] = [] + + def write_events(self, events: Sequence[ControlExecutionEvent]) -> SinkResult: + self.events.extend(events) + return SinkResult(accepted=len(events), dropped=0) + + +@pytest.fixture +def recording_sink(monkeypatch: pytest.MonkeyPatch) -> Iterator[RecordingSink]: + # Given: an isolated SDK session exporting to a registered event sink. + previous_settings = get_settings().model_dump() + for name in ("current_agent", "target_type", "target_id", "server_controls"): + monkeypatch.setattr(state, name, None) + sink = RecordingSink() + register_control_event_sink(sink) + configure_settings( + observability_enabled=True, + observability_sink_name=REGISTERED_CONTROL_EVENT_SINK_NAME, + ) + try: + yield sink + finally: + unregister_control_event_sink(sink) + configure_settings(**previous_settings) + + +def _server_response(category: ResultCategory, duration: float | None) -> dict[str, Any]: + match: dict[str, Any] = { + "control_id": 2, + "control_name": "server-control", + "control_execution_id": "server-execution", + "action": "observe", + "result": { + "matched": category == "matches", + "confidence": 1.0, + "error": "Evaluator failed" if category == "errors" else None, + }, + } + if duration is not None: + match["execution_duration_ms"] = duration + return {"is_safe": category != "errors", "confidence": 1.0, category: [match]} + + +def _control_payload(control_id: int, execution: str) -> dict[str, Any]: + return { + "id": control_id, + "name": f"{execution}-control", + "control": { + "execution": execution, + "scope": {"stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "regex", "config": {"pattern": "hello"}}, + }, + "action": {"decision": "observe"}, + }, + } + + +@pytest.mark.parametrize("category", ["matches", "errors", "non_matches"]) +@pytest.mark.parametrize("duration", [12.5, 0.0, None]) +async def test_server_duration_reaches_events_and_otel( + recording_sink: RecordingSink, + category: ResultCategory, + duration: float | None, +) -> None: + # Given: a real SDK client receiving a timed or legacy server response. + transport = httpx.MockTransport( + lambda request: httpx.Response(200, json=_server_response(category, duration)) + ) + + # When: evaluating through the public server-only SDK helper. + async with AgentControlClient(base_url="https://server.test", transport=transport) as client: + result = await check_evaluation( + client, + agent_name="agent-000000000001", + step=Step(type="llm", name="chat", input="hello"), + stage="pre", + ) + + # Then: the typed result and exported event retain the exact per-control value. + matches = getattr(result, category) + assert matches is not None + assert matches[0].execution_duration_ms == duration + assert len(recording_sink.events) == 1 + event = recording_sink.events[0] + assert event.execution_duration_ms == duration + assert "execution_duration_ms" not in event.metadata + span = control_event_to_otel_span(event) + if duration is None: + assert "agent_control.execution_duration_ms" not in span.attributes + else: + assert span.attributes["agent_control.execution_duration_ms"] == duration + assert span.end_time_unix_nano - span.start_time_unix_nano == int(duration * 1_000_000) + + +async def test_mixed_local_and_server_events_keep_their_own_durations( + recording_sink: RecordingSink, +) -> None: + # Given: a real local regex control and a server control with a distinct duration. + transport = httpx.MockTransport( + lambda request: httpx.Response(200, json=_server_response("non_matches", 27.5)) + ) + controls = [_control_payload(1, "sdk"), _control_payload(2, "server")] + + # When: evaluating both controls through the public local-first SDK helper. + async with AgentControlClient(base_url="https://server.test", transport=transport) as client: + result = await check_evaluation_with_local( + client, + agent_name="agent-000000000001", + step=Step(type="llm", name="chat", input="hello"), + stage="pre", + controls=controls, + trace_id="a" * 32, + span_id="b" * 16, + ) + + # Then: local engine timing and server timing each reach the corresponding event. + assert result.matches is not None + assert result.non_matches is not None + local_duration = result.matches[0].execution_duration_ms + assert local_duration is not None + assert local_duration >= 0 + assert len(recording_sink.events) == 2 + durations = {event.control_id: event.execution_duration_ms for event in recording_sink.events} + assert durations == {1: local_duration, 2: 27.5} + + +@pytest.mark.parametrize("category", ["matches", "errors", "non_matches"]) +@pytest.mark.parametrize("duration", [12.5, 0.0, None]) +async def test_decorator_preserves_duration_in_events_and_debug_logs( + recording_sink: RecordingSink, + category: ResultCategory, + duration: float | None, +) -> None: + # Given: a decorated control site using cached controls and a mock HTTP server. + transport = httpx.MockTransport( + lambda request: httpx.Response(200, json=_server_response(category, duration)) + ) + agent = MagicMock(agent_name="agent-000000000001") + + @control() + async def chat(message: str) -> str: + return message + + # When: calling the public decorator through its real evaluation and reconstruction paths. + with ( + patch("agent_control.current_agent", return_value=agent), + patch("agent_control.get_server_controls", return_value=[_control_payload(2, "server")]), + patch( + "agent_control.control_decorators.AgentControlClient", + partial(AgentControlClient, transport=transport), + ), + patch("agent_control.control_decorators.log_control_evaluation") as log_evaluation, + ): + if category == "errors": + with pytest.raises(RuntimeError, match="Control evaluation failed on server"): + await chat("hello") + else: + assert await chat("hello") == "hello" + + # Then: reconstruction retains duration for every category, including zero and missing values. + assert len(recording_sink.events) == 1 + assert recording_sink.events[0].execution_duration_ms == duration + log_evaluation.assert_called_once() + assert log_evaluation.call_args.kwargs["duration_ms"] == duration diff --git a/sdks/typescript/src/generated/models/control-match.ts b/sdks/typescript/src/generated/models/control-match.ts index 6a7caa1f..edf416ab 100644 --- a/sdks/typescript/src/generated/models/control-match.ts +++ b/sdks/typescript/src/generated/models/control-match.ts @@ -38,6 +38,10 @@ export type ControlMatch = { * Name of the control */ controlName: string; + /** + * Control evaluation wall-clock duration in milliseconds, if executed + */ + executionDurationMs?: number | null | undefined; /** * Result from a control evaluator. * @@ -69,6 +73,7 @@ export const ControlMatch$inboundSchema: z.ZodMiniType = control_execution_id: types.optional(types.string()), control_id: types.number(), control_name: types.string(), + execution_duration_ms: z.optional(z.nullable(types.number())), result: EvaluatorResult$inboundSchema, steering_context: z.optional(z.nullable(SteeringContext$inboundSchema)), }), @@ -77,6 +82,7 @@ export const ControlMatch$inboundSchema: z.ZodMiniType = "control_execution_id": "controlExecutionId", "control_id": "controlId", "control_name": "controlName", + "execution_duration_ms": "executionDurationMs", "steering_context": "steeringContext", }); }), diff --git a/sdks/typescript/tests/client-api.test.ts b/sdks/typescript/tests/client-api.test.ts index b0fa91cb..5d1dc591 100644 --- a/sdks/typescript/tests/client-api.test.ts +++ b/sdks/typescript/tests/client-api.test.ts @@ -21,6 +21,47 @@ describe("AgentControlClient API wiring", () => { vi.restoreAllMocks(); }); + it.each([12.5, 0, null, undefined])( + "preserves per-control execution durations through evaluation responses (%s)", + async (executionDurationMs) => { + const match = { + control_id: 1, + control_name: "Example control", + action: "observe", + result: { matched: true, confidence: 1.0 }, + execution_duration_ms: executionDurationMs, + }; + vi.mocked(globalThis.fetch).mockResolvedValueOnce( + jsonResponse({ + is_safe: true, + confidence: 1.0, + matches: [match], + non_matches: [ + { ...match, result: { matched: false, confidence: 1.0 } }, + ], + errors: [ + { + ...match, + result: { matched: false, confidence: 0.0, error: "Evaluation failed" }, + }, + ], + }), + ); + const client = new AgentControlClient(); + client.init({ agentName: "test-agent", serverUrl: "https://api.example.com" }); + + const response = await client.evaluation.evaluate({ + agentName: "test-agent", + stage: "pre", + step: { type: "llm", name: "answer", input: "hello" }, + }); + + expect(response.matches?.[0]?.executionDurationMs).toBe(executionDurationMs); + expect(response.nonMatches?.[0]?.executionDurationMs).toBe(executionDurationMs); + expect(response.errors?.[0]?.executionDurationMs).toBe(executionDurationMs); + }, + ); + it("throws when endpoint groups are accessed before init", () => { const client = new AgentControlClient(); diff --git a/server/tests/test_execution_duration_response.py b/server/tests/test_execution_duration_response.py new file mode 100644 index 00000000..dec932d6 --- /dev/null +++ b/server/tests/test_execution_duration_response.py @@ -0,0 +1,131 @@ +"""HTTP contract tests for per-control execution durations without database I/O.""" + +from unittest.mock import AsyncMock + +import pytest +from agent_control_models import ControlMatch, EvaluationResponse, EvaluatorResult +from fastapi.testclient import TestClient + +from agent_control_server.endpoints import evaluation +from agent_control_server.endpoints.evaluation import SAFE_EVALUATOR_ERROR + + +@pytest.fixture(scope="session", autouse=True) +def db_schema() -> None: + """Override database setup: control loading is mocked for this contract test.""" + + +@pytest.fixture(autouse=True) +def clean_db() -> None: + """Override database cleanup: these tests neither read nor write database rows.""" + + +@pytest.fixture +def engine_response(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Stub evaluation I/O while retaining the real HTTP response handling.""" + monkeypatch.setattr(evaluation, "_load_engine_controls", AsyncMock(return_value=[])) + process = AsyncMock() + monkeypatch.setattr(evaluation.ControlEngine, "process", process) + return process + + +def test_evaluation_response_preserves_duration_for_each_result_category( + client: TestClient, + engine_response: AsyncMock, +) -> None: + # Given: engine results with distinct per-control timings, including zero. + # Control loading and engine processing are stubbed to avoid database/evaluator I/O. + engine_response.return_value = EvaluationResponse( + is_safe=False, + confidence=0.9, + matches=[ + ControlMatch( + control_id=1, + control_name="matched-control", + action="deny", + execution_duration_ms=12.5, + result=EvaluatorResult(matched=True, confidence=0.9), + ) + ], + errors=[ + ControlMatch( + control_id=2, + control_name="failed-control", + action="observe", + execution_duration_ms=3.25, + result=EvaluatorResult( + matched=False, + confidence=0.0, + error="RuntimeError: private evaluator detail", + ), + ) + ], + non_matches=[ + ControlMatch( + control_id=3, + control_name="non-matched-control", + action="observe", + execution_duration_ms=0.0, + result=EvaluatorResult(matched=False, confidence=0.9), + ) + ], + ) + + # When: a caller evaluates a step through the public HTTP endpoint. + response = client.post( + "/api/v1/evaluation", + json={ + "agent_name": "duration-contract-agent", + "stage": "pre", + "step": {"type": "llm", "name": "answer", "input": "question"}, + }, + ) + + # Then: each category retains its own top-level duration through serialization. + assert response.status_code == 200 + payload = response.json() + assert payload["matches"][0]["execution_duration_ms"] == 12.5 + assert payload["errors"][0]["execution_duration_ms"] == 3.25 + assert payload["non_matches"][0]["execution_duration_ms"] == 0.0 + assert payload["errors"][0]["result"]["error"] == SAFE_EVALUATOR_ERROR + assert payload["errors"][0]["result"]["message"] == SAFE_EVALUATOR_ERROR + assert "private evaluator detail" not in response.text + for category in ("matches", "errors", "non_matches"): + assert "execution_duration_ms" not in payload[category][0]["result"] + + +@pytest.mark.parametrize("has_duration_field", [False, True], ids=["omitted", "null"]) +def test_evaluation_response_preserves_unknown_duration_as_null( + client: TestClient, + engine_response: AsyncMock, + has_duration_field: bool, +) -> None: + # Given: an old result without timing, or an explicitly unknown timing. + # Control loading and engine processing are stubbed to avoid database/evaluator I/O. + match_payload: dict[str, object] = { + "control_id": 1, + "control_name": "untimed-control", + "action": "observe", + "result": {"matched": False, "confidence": 0.0}, + } + if has_duration_field: + match_payload["execution_duration_ms"] = None + engine_response.return_value = EvaluationResponse( + is_safe=True, + confidence=1.0, + non_matches=[ControlMatch.model_validate(match_payload)], + ) + + # When: a caller evaluates a step through the public HTTP endpoint. + response = client.post( + "/api/v1/evaluation", + json={ + "agent_name": "duration-contract-agent", + "stage": "pre", + "step": {"type": "llm", "name": "answer", "input": "question"}, + }, + ) + + # Then: unknown timing remains null rather than becoming a measured zero. + assert response.status_code == 200 + assert response.json()["non_matches"][0]["execution_duration_ms"] is None