feat(sdk): measure whole control-check elapsed time

Author: josjeon

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",

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.

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

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".

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",

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)

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: