diff --git a/models/src/agent_control_models/__init__.py b/models/src/agent_control_models/__init__.py index 6b5562a7..ca50b477 100644 --- a/models/src/agent_control_models/__init__.py +++ b/models/src/agent_control_models/__init__.py @@ -69,6 +69,7 @@ from .observability import ( BatchEventsRequest, BatchEventsResponse, + ControlCheckEvent, ControlExecutionEvent, ControlStats, ControlStatsResponse, @@ -220,6 +221,7 @@ "ValidateControlDataRequest", "ValidateControlDataResponse", # Observability models + "ControlCheckEvent", "ControlExecutionEvent", "BatchEventsRequest", "BatchEventsResponse", diff --git a/models/src/agent_control_models/evaluation.py b/models/src/agent_control_models/evaluation.py index 50b5791b..7e6c4a82 100644 --- a/models/src/agent_control_models/evaluation.py +++ b/models/src/agent_control_models/evaluation.py @@ -170,6 +170,21 @@ class EvaluationResult(EvaluationResponse): This is what SDK users interact with. """ + control_check_id: str | None = Field( + default=None, + description="Unique ID of this SDK control-check invocation, when measured", + ) + control_check_elapsed_ms: float | None = Field( + default=None, + ge=0, + allow_inf_nan=False, + description=( + "SDK caller-visible elapsed milliseconds for one pre/post control check, " + "including local evaluation and conditional server request; excludes the " + "protected tool/LLM function" + ), + ) + def is_confident(self, threshold: float = 0.8) -> bool: """ Check if the result confidence exceeds a threshold. diff --git a/models/src/agent_control_models/observability.py b/models/src/agent_control_models/observability.py index 41cadd27..a0d30ce0 100644 --- a/models/src/agent_control_models/observability.py +++ b/models/src/agent_control_models/observability.py @@ -9,10 +9,10 @@ """ from datetime import UTC, datetime -from typing import Any, Literal +from typing import Any, Literal, Self from uuid import uuid4 -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator from .actions import ( ActionDecision, @@ -582,3 +582,31 @@ class ControlStatsResponse(BaseModel): @classmethod def validate_and_normalize_agent_name(cls, value: str) -> str: return normalize_agent_name(str(value)) + + +class ControlCheckEvent(BaseModel): + """One SDK control-check invocation, separate from individual control executions. + + Start is captured at invocation and end is anchored using monotonic elapsed + nanoseconds. No input/output or exception text is included. A completed check + can be unsafe or contain handled evaluator errors; error/cancelled describe + failures of the SDK check itself. + """ + + control_check_id: str = Field(min_length=1) + trace_id: str = Field(min_length=1) + span_id: str = Field(min_length=1, description="Parent step span ID") + agent_name: str = Field(min_length=1) + check_stage: Literal["pre", "post"] + status: Literal["completed", "error", "cancelled"] + elapsed_ms: float = Field(ge=0, allow_inf_nan=False) + start_time_unix_nano: int = Field(ge=0) + end_time_unix_nano: int = Field(ge=0) + is_safe: bool | None = None + + @model_validator(mode="after") + def validate_time_order(self) -> Self: + """Reject an end before the invocation started.""" + if self.end_time_unix_nano < self.start_time_unix_nano: + raise ValueError("control-check end must not precede its start") + return self diff --git a/sdks/python/README.md b/sdks/python/README.md index 9522cf2e..4477bed2 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -118,3 +118,58 @@ It force-flushes the provider during shutdown but leaves provider shutdown to th application. If no provider is passed, the OTEL sink first checks for a globally registered SDK provider, then falls back to its existing Agent Control-owned OTLP pipeline. + + +## Whole control-check elapsed time + +`evaluate_controls()`, `check_evaluation_with_local()` and the server-only +`evaluation.check_evaluation()` return `EvaluationResult.control_check_elapsed_ms` +and a unique `control_check_id`. Each value measures **one pre or post check** +with a monotonic clock. It includes preparation inside the evaluation helper, +local controls, and the conditional server call (network/auth refresh/retry wait +included). It excludes `init()`, decorator payload preparation before the helper, +the protected tool/LLM function, and export of the whole-check event itself. +High-level `evaluate_controls()` step construction and client setup happen before +the timed helper and are excluded, as is parent trace-context resolution. +The decorator server-only fallback also measures its HTTP evaluation attempt; +its already-built payload and client initialization are outside that boundary. +If local evaluation fails and the decorator falls back, these are separate check +attempts with separate IDs rather than one fabricated combined elapsed value. +Existing per-control `execution_duration_ms` values are unchanged. + +```python +# After agent_control.init(...) has configured the SDK session. +result = await agent_control.evaluate_controls( + agent_name="my-agent", step_name="search", input={"query": "example"}, + stage="pre", step_type="tool" +) +print(result.control_check_id, result.control_check_elapsed_ms) +``` + +Local denial, a completed result containing evaluator errors, and a check with +no applicable controls still have measured elapsed time. The last case measures +the check's preparation, rather than assigning zero as a placeholder. A genuinely +measured zero is valid. Older/manually constructed results have `None` unless +actually measured by these helpers. An SDK check that raises or is cancelled +still propagates that outcome; it does not return a fabricated evaluation result. + +With observability enabled and an `otel` sink selected, the same provider emits +one `agent_control.control_check` span as a child of the correlated step. Its +start is captured when the check begins and its end uses that start plus measured +monotonic elapsed nanoseconds; these are not reconstructed from control children. +Attributes include `agent_control.control_check_elapsed_ms`, the check ID, +`check_stage`, and `control_check_status` (`completed`, `error`, `cancelled`). +Safe IDs/stage/elapsed/status are also supplied in the JSON `metadata` attribute. +No input/output or exception text is recorded. Existing Orbit ingestion treats +this as a workflow span; it is **not** an individual shield/control span. Showing +these spans/metadata in a dedicated session summary requires UI/Orbit follow-up. + +Existing HTTP control-event ingestion accepts only individual control events; +it does not receive whole-check events. Custom registered sinks may opt in by +implementing `agent_control_telemetry.ControlCheckEventSink` in addition to their +existing `write_events` method. Control-only sinks remain compatible. Sink +failures never replace a result, exception or cancellation. + +An elapsed value is neither an average nor the maximum of child controls. +Several checks may overlap in a session, so summing their elapsed values does +not establish session wall time or a counterfactual "agent overhead". diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 04d040f6..7a36de7f 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -59,6 +59,7 @@ async def handle_input(user_message: str) -> str: from agent_control_models import ( Agent, ControlAction, + ControlCheckEvent, ControlDefinition, ControlMatch, ControlScope, @@ -112,6 +113,7 @@ async def handle_input(user_message: str) -> str: sync_shutdown_observability, unregister_control_event_sink, unregister_control_event_sink_factory, + write_control_check_events, write_events, ) from .otel_sink import control_event_to_otel_span @@ -1631,6 +1633,7 @@ async def main(): # Observability "init_observability", "add_event", + "write_control_check_events", "write_events", "shutdown_observability", "is_observability_enabled", @@ -1654,6 +1657,7 @@ async def main(): "StepSchema", "EvaluationRequest", "EvaluationResult", + "ControlCheckEvent", "ControlDefinition", "TemplateControlInput", "TemplateDefinition", diff --git a/sdks/python/src/agent_control/control_check.py b/sdks/python/src/agent_control/control_check.py new file mode 100644 index 00000000..0b499b71 --- /dev/null +++ b/sdks/python/src/agent_control/control_check.py @@ -0,0 +1,74 @@ +"""SDK caller-visible timing for one control-check invocation.""" + +import asyncio +import logging +import time +from collections.abc import Awaitable, Callable +from typing import Any, Literal +from uuid import uuid4 + +from agent_control_models import ControlCheckEvent, EvaluationResult + +from .observability import is_observability_enabled, write_control_check_events + +logger = logging.getLogger(__name__) + +async def measure_control_check[CheckResult: (EvaluationResult, dict[str, Any])]( + evaluate: Callable[[], Awaitable[CheckResult]], + *, + agent_name: str, + stage: Literal["pre", "post"], + trace_id: str | None, + span_id: str | None, +) -> CheckResult: + """Measure the actual local-plus-remote check, retaining exceptions/cancellation. + + The monotonic clock measures evaluation-helper work, including preparation + and existing per-control event handling. Check-event construction/export is + outside the elapsed boundary. Epoch time is captured only for observability; + its span end is the captured start plus measured monotonic nanoseconds. + """ + emit_event = is_observability_enabled() + start_unix_nano = time.time_ns() if emit_event else 0 + started = time.perf_counter_ns() + result: CheckResult | None = None + status: Literal["completed", "error", "cancelled"] = "error" + try: + result = await evaluate() + status = "completed" + return result + except asyncio.CancelledError: + status = "cancelled" + raise + finally: + elapsed_ns = time.perf_counter_ns() - started + elapsed_ms = elapsed_ns / 1_000_000 + check_id = str(uuid4()) if result is not None or emit_event else "" + is_safe: bool | None = None + if isinstance(result, EvaluationResult): + result.control_check_id = check_id + result.control_check_elapsed_ms = elapsed_ms + is_safe = result.is_safe + elif result is not None: + # Preserve the decorator's raw server response, including extensions. + result["control_check_id"] = check_id + result["control_check_elapsed_ms"] = elapsed_ms + if isinstance(result.get("is_safe"), bool): + is_safe = result["is_safe"] + if emit_event: + try: + event = ControlCheckEvent( + control_check_id=check_id, + trace_id=trace_id or "0" * 32, + span_id=span_id or "0" * 16, + agent_name=str(agent_name), + check_stage=stage, + status=status, + elapsed_ms=elapsed_ms, + start_time_unix_nano=start_unix_nano, + end_time_unix_nano=start_unix_nano + elapsed_ns, + is_safe=is_safe, + ) + write_control_check_events([event]) + except Exception: + logger.warning("Control-check telemetry failed", exc_info=True) diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index bfe6c02e..9f2387d3 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -34,13 +34,14 @@ async def chat(message: str) -> str: from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass, field -from typing import Any, TypeVar +from typing import Any, Literal, TypeVar, cast from agent_control_models import JSONObject, Step, normalize_action from agent_control_telemetry import get_trace_context_from_provider from agent_control import AgentControlClient from agent_control._state import state +from agent_control.control_check import measure_control_check from agent_control.evaluation import ( _post_evaluation_request, _resolve_session_target, @@ -337,6 +338,8 @@ async def _evaluate( "is_safe": result.is_safe, "confidence": result.confidence, "reason": result.reason, + "control_check_id": result.control_check_id, + "control_check_elapsed_ms": result.control_check_elapsed_ms, "matches": [ { "control_id": m.control_id, @@ -420,16 +423,25 @@ async def _evaluate( payload["target_type"] = target_type payload["target_id"] = target_id - response = await _post_evaluation_request( - client, - request_payload=payload, - headers=headers, - target_type=target_type, - target_id=target_id, + async def server_check() -> dict[str, Any]: + response = await _post_evaluation_request( + client, + request_payload=payload, + headers=headers, + target_type=target_type, + target_id=target_id, + ) + response.raise_for_status() + result_dict: dict[str, Any] = response.json() + return result_dict + + return await measure_control_check( + server_check, + agent_name=event_agent_name or agent_name, + stage=cast(Literal["pre", "post"], stage), + trace_id=trace_id, + span_id=span_id, ) - response.raise_for_status() - result_dict: dict[str, Any] = response.json() - return result_dict def _unexpected_control_failure_message(stage: str, error: Exception) -> str: diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 800b4bf4..f29cc79e 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -21,6 +21,7 @@ from ._state import state from .client import AgentControlClient +from .control_check import measure_control_check from .evaluation_events import build_control_execution_events, enqueue_observability_events from .observability import is_observability_enabled from .tracing import get_trace_and_span_ids @@ -257,6 +258,44 @@ async def check_evaluation( *, target_type: str | None = None, target_id: str | None = None, +) -> EvaluationResult: + """Evaluate one stage and measure its SDK caller-visible control-check elapsed time. + + Includes preparation inside this helper, SDK-local controls (when supplied), + and the conditional remote call, including network/token-refresh delays. + Excludes SDK initialization and the protected tool/LLM function. Legacy or + manually constructed results retain null timing. Errors/cancellation still + propagate; an enabled compatible sink records the check outcome separately. + """ + resolved_trace_id, resolved_span_id = get_trace_and_span_ids() + return await measure_control_check( + lambda: _check_evaluation( + client=client, + agent_name=agent_name, + step=step, + stage=stage, + target_type=target_type, + target_id=target_id, + trace_id=resolved_trace_id, + span_id=resolved_span_id, + ), + agent_name=agent_name, + stage=stage, + trace_id=resolved_trace_id, + span_id=resolved_span_id, + ) + + +async def _check_evaluation( + client: AgentControlClient, + agent_name: str, + step: Step, + stage: Literal["pre", "post"], + *, + target_type: str | None = None, + target_id: str | None = None, + trace_id: str | None = None, + span_id: str | None = None, ) -> EvaluationResult: """Check if agent interaction is safe through the public SDK helper. @@ -274,7 +313,7 @@ async def check_evaluation( _validate_target_pair(target_type, target_id) normalized_name = ensure_agent_name(agent_name) - resolved_trace_id, resolved_span_id = get_trace_and_span_ids() + resolved_trace_id, resolved_span_id = trace_id, span_id request = EvaluationRequest( agent_name=normalized_name, step=step, @@ -321,6 +360,52 @@ async def check_evaluation_with_local( *, target_type: str | None = None, target_id: str | None = None, +) -> EvaluationResult: + """Evaluate one stage and measure its SDK caller-visible control-check elapsed time. + + Includes preparation inside this helper, SDK-local controls (when supplied), + and the conditional remote call, including network/token-refresh delays. + Excludes SDK initialization and the protected tool/LLM function. Legacy or + manually constructed results retain null timing. Errors/cancellation still + propagate; an enabled compatible sink records the check outcome separately. + """ + resolved_trace_id, resolved_span_id = trace_id, span_id + if trace_id is None or span_id is None: + current_trace_id, current_span_id = get_trace_and_span_ids() + resolved_trace_id = trace_id or current_trace_id + resolved_span_id = span_id or current_span_id + return await measure_control_check( + lambda: _check_evaluation_with_local( + client=client, + agent_name=agent_name, + step=step, + stage=stage, + target_type=target_type, + target_id=target_id, + controls=controls, + trace_id=resolved_trace_id, + span_id=resolved_span_id, + event_agent_name=event_agent_name, + ), + agent_name=event_agent_name or agent_name, + stage=stage, + trace_id=resolved_trace_id, + span_id=resolved_span_id, + ) + + +async def _check_evaluation_with_local( + client: AgentControlClient, + agent_name: str, + step: Step, + stage: Literal["pre", "post"], + controls: list[dict[str, Any]], + trace_id: str | None = None, + span_id: str | None = None, + event_agent_name: str | None = None, + *, + target_type: str | None = None, + target_id: str | None = None, ) -> EvaluationResult: """Evaluate controls with local-first execution and SDK-owned event emission. diff --git a/sdks/python/src/agent_control/observability.py b/sdks/python/src/agent_control/observability.py index 39015bdb..aa3cf326 100644 --- a/sdks/python/src/agent_control/observability.py +++ b/sdks/python/src/agent_control/observability.py @@ -75,6 +75,7 @@ ) from agent_control_telemetry.sinks import ( BaseControlEventSink, + ControlCheckEventSink, ControlEventSink, SinkResult, ) @@ -82,7 +83,7 @@ from agent_control.settings import configure_settings, get_settings if TYPE_CHECKING: - from agent_control_models import ControlExecutionEvent + from agent_control_models import ControlCheckEvent, ControlExecutionEvent from opentelemetry.sdk.trace import TracerProvider from .otel_sink import ( @@ -1281,6 +1282,35 @@ def write_events(events: Sequence[ControlExecutionEvent]) -> SinkResult: return SinkResult(accepted=accepted, dropped=dropped) +def write_control_check_events(events: Sequence[ControlCheckEvent]) -> SinkResult: + """Deliver whole-check events to selected sinks with the optional capability. + + Control-only sinks remain unchanged. No HTTP payload is sent for these + events. Delivery failures are best-effort and never change an evaluation. + """ + check_sinks = tuple( + sink + for sink in _get_active_control_event_sinks() + if isinstance(sink, ControlCheckEventSink) + ) + if not check_sinks: + return SinkResult(accepted=0, dropped=len(events)) + + accepted = len(events) + dropped = 0 + for sink in check_sinks: + try: + result = sink.write_control_check_events(events) + accepted = min(accepted, result.accepted) + dropped = max(dropped, result.dropped) + except Exception: + logger.warning("Control-check sink write failed", exc_info=True) + accepted = 0 + dropped = len(events) + dropped = max(dropped, len(events) - accepted) + return SinkResult(accepted=min(accepted, len(events) - dropped), dropped=dropped) + + def sync_shutdown_observability() -> None: """Synchronously shut down observability and flush remaining events.""" _shutdown_built_in_event_sink() diff --git a/sdks/python/src/agent_control/otel_sink.py b/sdks/python/src/agent_control/otel_sink.py index aca0b4fe..abad1166 100644 --- a/sdks/python/src/agent_control/otel_sink.py +++ b/sdks/python/src/agent_control/otel_sink.py @@ -10,7 +10,7 @@ from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, cast -from agent_control_models import ControlExecutionEvent, JSONObject +from agent_control_models import ControlCheckEvent, ControlExecutionEvent, JSONObject from agent_control_telemetry.sinks import BaseControlEventSink, SinkResult from .settings import get_settings @@ -156,6 +156,37 @@ def control_event_to_otel_span(event: ControlExecutionEvent) -> OTELControlEvent ) +def control_check_event_to_otel_span(event: ControlCheckEvent) -> OTELControlEventSpan: + """Convert a measured whole check without inventing individual control identity.""" + attributes: dict[str, AttributeValue] = { + "agent_control.control_check_id": event.control_check_id, + "agent_control.agent_name": event.agent_name, + "agent_control.check_stage": event.check_stage, + "agent_control.control_check_status": event.status, + "agent_control.control_check_elapsed_ms": event.elapsed_ms, + } + attributes["metadata"] = json.dumps( + { + "control_check_id": event.control_check_id, + "check_stage": event.check_stage, + "control_check_elapsed_ms": event.elapsed_ms, + "control_check_status": event.status, + }, + sort_keys=True, + ) + if event.is_safe is not None: + attributes["agent_control.is_safe"] = event.is_safe + return OTELControlEventSpan( + name="agent_control.control_check", + trace_id=event.trace_id, + parent_span_id=event.span_id, + attributes=attributes, + start_time_unix_nano=event.start_time_unix_nano, + end_time_unix_nano=event.end_time_unix_nano, + error_message="Control check failed" if event.status == "error" else None, + ) + + class _NoOpControlEventSink(BaseControlEventSink): """Sink that accepts events but intentionally emits nothing.""" @@ -196,6 +227,17 @@ def write_events(self, events: Sequence[ControlExecutionEvent]) -> SinkResult: return SinkResult(accepted=accepted, dropped=dropped) + def write_control_check_events(self, events: Sequence[ControlCheckEvent]) -> SinkResult: + """Write genuine SDK check spans using this sink's existing provider.""" + accepted = 0 + for event in events: + try: + self._write_span(control_check_event_to_otel_span(event)) + accepted += 1 + except Exception: + logger.warning("Failed to emit control check to OTEL", exc_info=True) + return SinkResult(accepted=accepted, dropped=len(events) - accepted) + def flush(self) -> None: force_flush = getattr(self._tracer_provider, "force_flush", None) if callable(force_flush): @@ -209,7 +251,9 @@ def close(self) -> None: shutdown() def _write_event(self, event: ControlExecutionEvent) -> None: - span_data = control_event_to_otel_span(event) + self._write_span(control_event_to_otel_span(event)) + + def _write_span(self, span_data: OTELControlEventSpan) -> None: parent_context = self._build_parent_context(span_data) span = self._tracer.start_span( span_data.name, diff --git a/sdks/python/tests/test_control_check_elapsed.py b/sdks/python/tests/test_control_check_elapsed.py new file mode 100644 index 00000000..3ce51d76 --- /dev/null +++ b/sdks/python/tests/test_control_check_elapsed.py @@ -0,0 +1,386 @@ +"""Public SDK contract tests for actual whole-check elapsed timing.""" + +import asyncio +import json +from collections.abc import Iterator, Sequence +from typing import Any +from unittest.mock import Mock + +import httpx +import pytest +from agent_control_models import ControlExecutionEvent, EvaluationResult, Step +from agent_control_telemetry import BaseControlEventSink, SinkResult +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from pydantic import ValidationError + +from agent_control import ( + AgentControlClient, + ControlCheckEvent, + check_evaluation_with_local, + control_check, + evaluation, + init_observability, + register_control_event_sink, + sync_shutdown_observability, + unregister_control_event_sink, +) +from agent_control.settings import configure_settings, get_settings + + +class CheckCollector(BaseControlEventSink): + """Opt-in sink collecting both event types without a server.""" + + def __init__(self) -> None: + self.checks: list[ControlCheckEvent] = [] + + def write_events(self, events: Sequence[ControlExecutionEvent]) -> SinkResult: + return SinkResult(accepted=len(events)) + + def write_control_check_events(self, events: Sequence[ControlCheckEvent]) -> SinkResult: + self.checks.extend(events) + return SinkResult(accepted=len(events)) + + +@pytest.fixture(autouse=True) +def settings() -> Iterator[None]: + previous = get_settings().model_dump() + configure_settings(observability_enabled=False) + yield + sync_shutdown_observability() + configure_settings(**previous) + + +@pytest.fixture +def collector() -> Iterator[CheckCollector]: + sink = CheckCollector() + register_control_event_sink(sink) + configure_settings(observability_enabled=True, observability_sink_name="registered") + yield sink + unregister_control_event_sink(sink) + + +def control_payload(execution: str, *, deny: bool = False) -> dict[str, Any]: + return { + "id": 1 if execution == "sdk" else 2, + "name": f"{execution}-control", + "control": { + "description": "Timing test", + "execution": execution, + "scope": {"step_types": ["llm"], "stages": ["pre", "post"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": {"name": "regex", "config": {"pattern": "hello"}}, + }, + "action": {"decision": "deny" if deny else "observe"}, + }, + } + + +def step() -> Step: + return Step(type="llm", name="chat", input="hello") + + +@pytest.mark.asyncio +async def test_five_remote_controls_use_measured_batch_not_mean_or_max(monkeypatch) -> None: + # Given: five control durations and an independently measured HTTP round trip. + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 100_000_000])) + wall_clock = Mock(side_effect=AssertionError("disabled telemetry must not read epoch clock")) + monkeypatch.setattr(control_check.time, "time_ns", wall_clock) + matches = [ + { + "control_id": number, + "control_name": f"c{number}", + "action": "observe", + "result": {"matched": False, "confidence": 1}, + "execution_duration_ms": number * 10, + } + for number in range(1, 6) + ] + transport = httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"is_safe": True, "confidence": 1, "non_matches": matches} + ) + ) + + # When: using the public SDK helper. + async with AgentControlClient(base_url="https://test.invalid", transport=transport) as client: + result = await check_evaluation_with_local( + client, "test-agent", step(), "pre", [control_payload("server")] + ) + + # Then: the batch is measured directly; individual control durations are preserved. + assert result.control_check_elapsed_ms == 100 + assert result.control_check_id + assert [item.execution_duration_ms for item in result.non_matches] == [10, 20, 30, 40, 50] + assert result.control_check_elapsed_ms != 30 # average + assert result.control_check_elapsed_ms != 50 # maximum + wall_clock.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stage", ["pre", "post"]) +async def test_mixed_local_remote_check_has_one_event_and_parent( + stage, monkeypatch, collector +) -> None: + # Given: local and remote controls, plus a deterministic whole-check clock. + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 140_000_000])) + monkeypatch.setattr(control_check.time, "time_ns", lambda: 1_000_000_000) + transport = httpx.MockTransport( + lambda request: httpx.Response(200, json={"is_safe": True, "confidence": 1}) + ) + + # When: one stage checks both local and conditional remote controls. + async with AgentControlClient(base_url="https://test.invalid", transport=transport) as client: + result = await check_evaluation_with_local( + client, + "test-agent", + step(), + stage, + [control_payload("sdk"), control_payload("server")], + trace_id="a" * 32, + span_id="b" * 16, + ) + + # Then: one measured event identifies this invocation and preserves local outcomes. + assert result.control_check_elapsed_ms == 140 + assert result.matches[0].execution_duration_ms is not None + assert len(collector.checks) == 1 + event = collector.checks[0] + assert event.control_check_id == result.control_check_id + assert event.trace_id == "a" * 32 + assert event.span_id == "b" * 16 + assert event.check_stage == stage + assert event.status == "completed" + assert event.start_time_unix_nano == 1_000_000_000 + assert event.end_time_unix_nano == 1_140_000_000 + + +@pytest.mark.asyncio +async def test_local_deny_short_circuits_server_but_still_measures(monkeypatch, collector) -> None: + # Given: a local deny and a server control that must never be contacted. + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 5_000_000])) + transport = httpx.MockTransport(lambda request: pytest.fail("server called after local deny")) + + # When: evaluating the controls through the public SDK. + async with AgentControlClient(base_url="https://test.invalid", transport=transport) as client: + result = await check_evaluation_with_local( + client, + "test-agent", + step(), + "pre", + [control_payload("sdk", deny=True), control_payload("server")], + ) + + # Then: a completed unsafe batch has its own elapsed time. + assert not result.is_safe + assert result.control_check_elapsed_ms == 5 + assert collector.checks[0].status == "completed" + assert collector.checks[0].is_safe is False + + +@pytest.mark.asyncio +async def test_empty_check_and_measured_zero_are_not_unknown(monkeypatch, collector) -> None: + # Given: no applicable controls and a clock with genuinely equal readings. + monkeypatch.setattr(control_check.time, "perf_counter_ns", lambda: 5) + async with AgentControlClient(base_url="https://test.invalid") as client: + # When: completing an empty check. + result = await check_evaluation_with_local(client, "test-agent", step(), "post", []) + + # Then: measured zero is retained, while legacy result timing remains unknown. + assert result.control_check_elapsed_ms == 0 + assert collector.checks[0].elapsed_ms == 0 + assert EvaluationResult(is_safe=True, confidence=1).control_check_elapsed_ms is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_error_and_cancellation_propagate_and_emit_status(cancelled, monkeypatch, collector): + # Given: an actual HTTP transport failure or cancellation during the check. + error = asyncio.CancelledError() if cancelled else httpx.ConnectError("private detail") + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 8_000_000])) + + async def fail(request: httpx.Request) -> httpx.Response: + raise error + + # When: calling the public helper. + async with AgentControlClient( + base_url="https://test.invalid", transport=httpx.MockTransport(fail) + ) as client: + with pytest.raises(type(error)) as caught: + await check_evaluation_with_local( + client, "test-agent", step(), "pre", [control_payload("server")] + ) + + # Then: the original failure survives and no sensitive exception text is recorded. + assert caught.value is error + assert collector.checks[0].status == ("cancelled" if cancelled else "error") + assert collector.checks[0].is_safe is None + assert "private detail" not in collector.checks[0].model_dump_json() + + +@pytest.mark.asyncio +async def test_telemetry_failure_preserves_completed_result(monkeypatch, collector) -> None: + # Given: an enabled sink whose check-event export fails. + monkeypatch.setattr( + control_check, "write_control_check_events", Mock(side_effect=RuntimeError("sink failed")) + ) + async with AgentControlClient(base_url="https://test.invalid") as client: + # When: completing an empty check. + result = await check_evaluation_with_local(client, "test-agent", step(), "pre", []) + + # Then: the valid business result and measured timing survive. + assert result.is_safe + assert result.control_check_elapsed_ms is not None + + +@pytest.mark.asyncio +async def test_otel_check_uses_existing_provider_and_actual_interval(monkeypatch) -> None: + # Given: a configured provider/exporter and real interval independent of child events. + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + init_observability( + enabled=True, + sink_name="otel", + sink_config={"enabled": True}, + otel_tracer_provider=provider, + ) + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 25_000_000])) + monkeypatch.setattr(control_check.time, "time_ns", lambda: 2_000_000_000) + + # When: completing a check with no control children. + async with AgentControlClient(base_url="https://test.invalid") as client: + result = await check_evaluation_with_local( + client, + "test-agent", + step(), + "post", + [], + trace_id="a" * 32, + span_id="b" * 16, + ) + + # Then: the genuine batch span is correlated and is not mislabeled as a control. + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "agent_control.control_check" + assert span.parent.span_id == int("b" * 16, 16) + assert span.start_time == 2_000_000_000 + assert span.end_time == 2_025_000_000 + assert span.attributes["agent_control.control_check_elapsed_ms"] == 25 + assert span.attributes["agent_control.control_check_id"] == result.control_check_id + assert "agent_control.control_id" not in span.attributes + assert "agent_control.action" not in span.attributes + metadata = json.loads(span.attributes["metadata"]) + assert metadata["check_stage"] == "post" + assert metadata["control_check_elapsed_ms"] == 25 + provider.shutdown() + + +def test_result_timing_rejects_invalid_values_and_server_schema_is_unchanged() -> None: + # Given: malformed elapsed values and a legacy server response. + for value in [-1, float("inf"), float("nan")]: + # When: constructing a typed result with invalid timing. + with pytest.raises(ValidationError): + EvaluationResult(is_safe=True, confidence=1, control_check_elapsed_ms=value) + # Then: the server response has no SDK-only fields. + assert "control_check_elapsed_ms" not in evaluation.EvaluationResponse.model_fields + + +@pytest.mark.asyncio +async def test_old_control_only_sink_remains_compatible() -> None: + # Given: a selected sink implementing only the existing event contract. + class OldSink(BaseControlEventSink): + def write_events(self, events: Sequence[ControlExecutionEvent]) -> SinkResult: + return SinkResult(accepted=len(events)) + + sink = OldSink() + register_control_event_sink(sink) + configure_settings(observability_enabled=True, observability_sink_name="registered") + try: + async with AgentControlClient(base_url="https://test.invalid") as client: + # When: completing a check with the old sink selected. + result = await check_evaluation_with_local(client, "test-agent", step(), "pre", []) + # Then: timing is available without requiring a new sink method or HTTP payload. + assert result.is_safe + assert result.control_check_elapsed_ms is not None + finally: + unregister_control_event_sink(sink) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancelled", [False, True]) +async def test_telemetry_failure_preserves_original_exception(cancelled, monkeypatch, collector): + # Given: a transport failure/cancellation and an independent telemetry failure. + error = asyncio.CancelledError() if cancelled else httpx.ConnectError("original") + monkeypatch.setattr( + control_check, "write_control_check_events", Mock(side_effect=RuntimeError("sink failed")) + ) + + async def fail(request: httpx.Request) -> httpx.Response: + raise error + + # When: calling the public helper. + async with AgentControlClient( + base_url="https://test.invalid", transport=httpx.MockTransport(fail) + ) as client: + with pytest.raises(type(error)) as caught: + await check_evaluation_with_local( + client, "test-agent", step(), "pre", [control_payload("server")] + ) + # Then: telemetry cannot replace the original business failure. + assert caught.value is error + + +@pytest.mark.asyncio +async def test_decorator_server_fallback_preserves_raw_extensions_and_measures( + monkeypatch, collector +): + # Given: an uncached decorator site with a raw server response extension. + # The internal adapter is needed to force the fallback without global init/server setup. + from agent_control import control_decorators + + response = httpx.Response( + 200, + request=httpx.Request("POST", "https://test.invalid/api/v1/evaluation"), + json={"is_safe": True, "confidence": 1, "extension": {"kept": True}}, + ) + from unittest.mock import AsyncMock + + monkeypatch.setattr( + control_decorators, "_post_evaluation_request", AsyncMock(return_value=response) + ) + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 12_000_000])) + # When: the decorator takes its server-only fallback. + result = await control_decorators._evaluate( + "test-agent", + step().model_dump(), + "post", + "https://test.invalid", + trace_id="a" * 32, + span_id="b" * 16, + controls=None, + ) + # Then: raw extensions survive and the genuine check interval is captured. + assert result["extension"] == {"kept": True} + assert result["control_check_elapsed_ms"] == 12 + assert collector.checks[0].control_check_id == result["control_check_id"] + assert collector.checks[0].span_id == "b" * 16 + + +@pytest.mark.asyncio +async def test_server_only_helper_returns_measured_result(monkeypatch, collector) -> None: + # Given: a server response with no duration fields from an older server. + monkeypatch.setattr(control_check.time, "perf_counter_ns", Mock(side_effect=[0, 20_000_000])) + transport = httpx.MockTransport( + lambda request: httpx.Response(200, json={"is_safe": True, "confidence": 1}) + ) + async with AgentControlClient(base_url="https://test.invalid", transport=transport) as client: + # When: using the public server-only helper. + result = await evaluation.check_evaluation(client, "test-agent", step(), "pre") + # Then: this SDK measures caller elapsed without relying on server timing fields. + assert result.control_check_elapsed_ms == 20 + assert collector.checks[0].elapsed_ms == 20 diff --git a/telemetry/src/agent_control_telemetry/__init__.py b/telemetry/src/agent_control_telemetry/__init__.py index a41f0f45..8b88c815 100644 --- a/telemetry/src/agent_control_telemetry/__init__.py +++ b/telemetry/src/agent_control_telemetry/__init__.py @@ -12,6 +12,7 @@ AsyncControlEventSink, BaseAsyncControlEventSink, BaseControlEventSink, + ControlCheckEventSink, ControlEventSink, SinkResult, ) @@ -27,6 +28,7 @@ "AsyncControlEventSink", "BaseAsyncControlEventSink", "BaseControlEventSink", + "ControlCheckEventSink", "ControlEventSink", "SinkResult", "DEFAULT_CONTROL_EVENT_SINK_NAME", diff --git a/telemetry/src/agent_control_telemetry/sinks.py b/telemetry/src/agent_control_telemetry/sinks.py index 17b8322f..75618b0e 100644 --- a/telemetry/src/agent_control_telemetry/sinks.py +++ b/telemetry/src/agent_control_telemetry/sinks.py @@ -4,9 +4,9 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import Protocol +from typing import Protocol, runtime_checkable -from agent_control_models import ControlExecutionEvent +from agent_control_models import ControlCheckEvent, ControlExecutionEvent @dataclass(frozen=True) @@ -74,3 +74,15 @@ async def write_events(self, events: Sequence[ControlExecutionEvent]) -> SinkRes accepted += result.accepted dropped += result.dropped return SinkResult(accepted=accepted, dropped=dropped) + + +@runtime_checkable +class ControlCheckEventSink(Protocol): + """Optional capability for sinks that support whole SDK control checks. + + Existing control-execution sinks need not implement this capability. Check + events are not accepted by the HTTP control-execution ingestion endpoint. + """ + + def write_control_check_events(self, events: Sequence[ControlCheckEvent]) -> SinkResult: + """Write measured SDK check events, separately from individual controls."""