address comments

Author: namrataghadi-galileo

sdks/python/src/agent_control/evaluation.py

@@ -18,6 +18,7 @@
 from ._state import state
 from .client import AgentControlClient
 from .evaluation_events import build_control_execution_events, enqueue_observability_events
+from .observability import is_observability_enabled
 from .telemetry import emit_control_events, has_control_event_sink
 from .tracing import get_trace_and_span_ids
 from .validation import ensure_agent_name
@@ -173,8 +174,36 @@ async def check_evaluation(
     step: Step,
     stage: Literal["pre", "post"],
 ) -> EvaluationResult:
-    """Check if agent interaction is safe."""
+    """Check if agent interaction is safe through the public SDK helper.
+
+    This helper preserves the default server-only evaluation path, but it also
+    participates in the optional merged-event flow when a control event sink is
+    registered. In that mode, the SDK asks the server to skip final event
+    ingestion, reconstructs server events from the lightweight response, and
+    emits them through the registered sink before returning the parsed result.
+
+    Args:
+        client: Configured AgentControl client.
+        agent_name: Agent name to evaluate against.
+        step: Step payload to evaluate.
+        stage: Evaluation stage, ``pre`` or ``post``.
+
+    Returns:
+        The parsed evaluation result returned by the server.
+    """
     normalized_name = ensure_agent_name(agent_name)
+    merged_emission_enabled = has_control_event_sink()
+    trace_id = None
+    span_id = None
+    headers: dict[str, str] | None = None
+
+    if merged_emission_enabled:
+        trace_id, span_id = get_trace_and_span_ids()
+        headers = {
+            "X-Trace-Id": trace_id,
+            "X-Span-Id": span_id,
+            "X-Agent-Control-Merge-Events": "true",
+        }
 
     request = EvaluationRequest(
         agent_name=normalized_name,
@@ -183,10 +212,28 @@ async def check_evaluation(
     )
     request_payload = request.model_dump(mode="json")
 
-    response = await client.http_client.post("/api/v1/evaluation", json=request_payload)
+    response = await client.http_client.post(
+        "/api/v1/evaluation",
+        json=request_payload,
+        headers=headers,
+    )
     response.raise_for_status()
 
-    return cast(EvaluationResult, EvaluationResult.from_dict(response.json()))
+    evaluation_response = EvaluationResponse.model_validate(response.json())
+
+    if merged_emission_enabled:
+        server_control_lookup = _build_server_control_lookup(state.server_controls or [])
+        server_events = build_control_execution_events(
+            evaluation_response,
+            request,
+            server_control_lookup,
+            trace_id,
+            span_id,
+            normalized_name,
+        )
+        emit_control_events(server_events)
+
+    return cast(EvaluationResult, EvaluationResult.from_dict(evaluation_response.model_dump()))
 
 
 async def check_evaluation_with_local(
@@ -305,6 +352,7 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
         return result.model_copy(update={"errors": combined_errors})
 
     merged_emission_enabled = has_control_event_sink()
+    should_reconstruct_local_events = merged_emission_enabled or is_observability_enabled()
 
     local_result: EvaluationResponse | None = None
     local_events = []
@@ -316,20 +364,21 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
     if applicable_local_controls:
         engine = ControlEngine(applicable_local_controls, context="sdk")
         local_result = await engine.process(request)
-        local_control_lookup = {
-            control.id: control.control for control in applicable_local_controls
-        }
-        local_events = build_control_execution_events(
-            local_result,
-            request,
-            local_control_lookup,
-            resolved_trace_id,
-            resolved_span_id,
-            event_agent_name,
-        )
+        if should_reconstruct_local_events:
+            local_control_lookup = {
+                control.id: control.control for control in applicable_local_controls
+            }
+            local_events = build_control_execution_events(
+                local_result,
+                request,
+                local_control_lookup,
+                resolved_trace_id,
+                resolved_span_id,
+                event_agent_name,
+            )
 
-        if not merged_emission_enabled:
-            enqueue_observability_events(local_events)
+            if not merged_emission_enabled:
+                enqueue_observability_events(local_events)
 
         if not local_result.is_safe:
             result = _with_parse_errors(EvaluationResult.model_validate(local_result.model_dump()))
@@ -354,15 +403,17 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
         )
         response.raise_for_status()
         server_result = EvaluationResponse.model_validate(response.json())
-        server_control_lookup = _build_server_control_lookup(server_control_payloads)
-        server_events = build_control_execution_events(
-            server_result,
-            request,
-            server_control_lookup,
-            resolved_trace_id,
-            resolved_span_id,
-            event_agent_name,
-        )
+        server_events = []
+        if merged_emission_enabled:
+            server_control_lookup = _build_server_control_lookup(server_control_payloads)
+            server_events = build_control_execution_events(
+                server_result,
+                request,
+                server_control_lookup,
+                resolved_trace_id,
+                resolved_span_id,
+                event_agent_name,
+            )
 
         if local_result is not None:
             result = _with_parse_errors(_merge_results(local_result, server_result))

sdks/python/tests/test_evaluation.py

@@ -66,6 +66,7 @@ def json(self) -> dict[str, object]:
             },
             "stage": "pre",
         },
+        headers=None,
     )
 
 

sdks/python/tests/test_observability_updates.py

@@ -326,6 +326,7 @@ async def test_resolves_provider_trace_context_for_local_events(self):
 
     @pytest.mark.asyncio
     async def test_forwards_provider_trace_headers_to_server_when_ids_omitted(self):
+        """Server POST should resolve trace headers from the provider when omitted."""
         from agent_control_models import Step
 
         controls = [{
@@ -355,7 +356,12 @@ async def test_forwards_provider_trace_headers_to_server_when_ids_omitted(self):
         client.http_client = AsyncMock()
         client.http_client.post = AsyncMock(return_value=mock_http_response)
         step = Step(type="llm", name="test-step", input="hello")
-        set_trace_context_provider(lambda: {"trace_id": "c" * 32, "span_id": "d" * 16})
+        set_trace_context_provider(
+            lambda: {
+                "trace_id": "c" * 32,
+                "span_id": "d" * 16,
+            }
+        )
 
         with patch("agent_control.evaluation.list_evaluators", return_value=["regex"]):
             await evaluation.check_evaluation_with_local(
@@ -366,67 +372,158 @@ async def test_forwards_provider_trace_headers_to_server_when_ids_omitted(self):
                 controls=controls,
             )
 
-        headers = client.http_client.post.call_args.kwargs["headers"]
-        assert headers["X-Trace-Id"] == "c" * 32
-        assert headers["X-Span-Id"] == "d" * 16
-        # Verify POST was called with headers
         call_kwargs = client.http_client.post.call_args
         headers = call_kwargs.kwargs.get("headers", {})
-        assert headers["X-Trace-Id"] == "aaaa1111bbbb2222cccc3333dddd4444"
-        assert headers["X-Span-Id"] == "eeee5555ffff6666"
+        assert headers["X-Trace-Id"] == "c" * 32
+        assert headers["X-Span-Id"] == "d" * 16
 
+
+class TestCheckEvaluation:
     @pytest.mark.asyncio
-    async def test_forwards_provider_trace_headers_to_server_when_ids_omitted(self):
-        """Server POST should resolve trace headers from the provider when omitted."""
+    async def test_default_path_keeps_server_only_behavior(self):
         from agent_control_models import Step
 
+        mock_http_response = MagicMock()
+        mock_http_response.raise_for_status = MagicMock()
+        mock_http_response.json.return_value = {
+            "is_safe": True,
+            "confidence": 0.9,
+            "matches": None,
+            "errors": None,
+            "non_matches": None,
+        }
+
+        client = MagicMock()
+        client.http_client = AsyncMock()
+        client.http_client.post = AsyncMock(return_value=mock_http_response)
+        step = Step(type="llm", name="test-step", input="hello")
+
+        with patch("agent_control.evaluation.has_control_event_sink", return_value=False), \
+             patch("agent_control.evaluation.emit_control_events") as mock_emit:
+            result = await evaluation.check_evaluation(
+                client=client,
+                agent_name="agent-000000000001",
+                step=step,
+                stage="pre",
+            )
+
+        call_kwargs = client.http_client.post.call_args.kwargs
+        assert call_kwargs["headers"] is None
+        mock_emit.assert_not_called()
+        assert result.is_safe is True
+        assert result.confidence == 0.9
+
+    @pytest.mark.asyncio
+    async def test_merged_event_sink_emits_reconstructed_server_events(self):
+        from agent_control_models import Step
+
+        controls = [
+            {
+                "id": 2,
+                "name": "server-ctrl",
+                "control": {
+                    "condition": {
+                        "evaluator": {"name": "regex", "config": {"pattern": "test"}},
+                        "selector": {"path": "input"},
+                    },
+                    "action": {"decision": "allow"},
+                    "execution": "server",
+                },
+            }
+        ]
+
+        mock_http_response = MagicMock()
+        mock_http_response.raise_for_status = MagicMock()
+        mock_http_response.json.return_value = {
+            "is_safe": True,
+            "confidence": 0.9,
+            "matches": [
+                {
+                    "control_id": 2,
+                    "control_name": "server-ctrl",
+                    "action": "allow",
+                    "control_execution_id": "ce-server",
+                    "result": {"matched": True, "confidence": 0.4},
+                }
+            ],
+            "errors": None,
+            "non_matches": None,
+        }
+
+        client = MagicMock()
+        client.http_client = AsyncMock()
+        client.http_client.post = AsyncMock(return_value=mock_http_response)
+        step = Step(type="llm", name="test-step", input="hello")
+
+        with patch("agent_control.evaluation.has_control_event_sink", return_value=True), \
+             patch("agent_control.evaluation.get_trace_and_span_ids", return_value=("a" * 32, "b" * 16)), \
+             patch("agent_control.evaluation.emit_control_events") as mock_emit, \
+             patch.object(evaluation.state, "server_controls", controls):
+            result = await evaluation.check_evaluation(
+                client=client,
+                agent_name="agent-000000000001",
+                step=step,
+                stage="pre",
+            )
+
+        headers = client.http_client.post.call_args.kwargs["headers"]
+        assert headers["X-Trace-Id"] == "a" * 32
+        assert headers["X-Span-Id"] == "b" * 16
+        assert headers["X-Agent-Control-Merge-Events"] == "true"
+
+        mock_emit.assert_called_once()
+        emitted_events = mock_emit.call_args.args[0]
+        assert len(emitted_events) == 1
+        assert emitted_events[0].control_id == 2
+        assert emitted_events[0].trace_id == "a" * 32
+        assert emitted_events[0].span_id == "b" * 16
+        assert emitted_events[0].metadata["primary_evaluator"] == "regex"
+        assert result.matches is not None
+        assert len(result.matches) == 1
+
+    @pytest.mark.asyncio
+    async def test_skips_local_event_reconstruction_when_nothing_consumes_events(self):
+        from agent_control_models import EvaluationResponse, Step
+
         controls = [{
             "id": 1,
-            "name": "server-ctrl",
+            "name": "local-ctrl",
             "control": {
                 "condition": {
                     "evaluator": {"name": "regex", "config": {"pattern": "test"}},
                     "selector": {"path": "input"},
                 },
-                "action": {"decision": "deny"},
-                "execution": "server",
+                "action": {"decision": "allow"},
+                "execution": "sdk",
             },
         }]
 
-        mock_http_response = MagicMock()
-        mock_http_response.json.return_value = {
-            "is_safe": True,
-            "confidence": 1.0,
-            "matches": None,
-            "errors": None,
-            "non_matches": None,
-        }
-        mock_http_response.raise_for_status = MagicMock()
+        mock_response = EvaluationResponse(is_safe=True, confidence=1.0)
+        mock_engine = MagicMock()
+        mock_engine.process = AsyncMock(return_value=mock_response)
 
         client = MagicMock()
         client.http_client = AsyncMock()
-        client.http_client.post = AsyncMock(return_value=mock_http_response)
         step = Step(type="llm", name="test-step", input="hello")
-        set_trace_context_provider(
-            lambda: {
-                "trace_id": "c" * 32,
-                "span_id": "d" * 16,
-            }
-        )
 
-        with patch("agent_control.evaluation.list_evaluators", return_value=["regex"]):
-            await evaluation.check_evaluation_with_local(
+        with patch("agent_control.evaluation.ControlEngine", return_value=mock_engine), \
+             patch("agent_control.evaluation.list_evaluators", return_value=["regex"]), \
+             patch("agent_control.evaluation.has_control_event_sink", return_value=False), \
+             patch("agent_control.evaluation.is_observability_enabled", return_value=False), \
+             patch("agent_control.evaluation.build_control_execution_events") as mock_build, \
+             patch("agent_control.evaluation.enqueue_observability_events") as mock_enqueue:
+            result = await evaluation.check_evaluation_with_local(
                 client=client,
                 agent_name="agent-000000000001",
                 step=step,
                 stage="pre",
                 controls=controls,
             )
 
-        call_kwargs = client.http_client.post.call_args
-        headers = call_kwargs.kwargs.get("headers", {})
-        assert headers["X-Trace-Id"] == "c" * 32
-        assert headers["X-Span-Id"] == "d" * 16
+        mock_build.assert_not_called()
+        mock_enqueue.assert_not_called()
+        assert result.is_safe is True
+        assert result.confidence == 1.0
 
 
 # =============================================================================