ensure control sink exists for merged mode

Author: namrataghadi-galileo

sdks/python/src/agent_control/__init__.py

@@ -477,6 +477,7 @@ async def handle(message: str):
 
     # Re-init behavior: always stop existing loop before mutating shared agent/session globals.
     _stop_policy_refresh_loop()
+    clear_control_event_sink()
 
     # Configure logging if provided (do this early before any logging happens)
     if log_config:
@@ -647,6 +648,7 @@ def _reset_state() -> None:
         state.server_controls = None
         state.server_url = None
         state.api_key = None
+        clear_control_event_sink()
 
 
 async def ashutdown() -> None:

sdks/python/src/agent_control/evaluation.py

@@ -117,6 +117,32 @@ def _has_applicable_prefiltered_server_controls(
     )
 
 
+def _is_merged_event_mode_enabled(agent_name: str) -> bool:
+    """Return whether SDK-side merged event emission is safe for this request.
+
+    Merged reconstruction depends on initialized SDK session state:
+    a registered sink, an initialized agent, and a cached server-control
+    snapshot for the same agent. If any of those prerequisites are missing,
+    evaluation falls back to the default behavior where the server remains the
+    final emitter for server-side events.
+
+    Args:
+        agent_name: Normalized agent name for the current request.
+
+    Returns:
+        ``True`` when the current SDK session has enough state to reconstruct
+        and emit merged events safely.
+    """
+    if not has_control_event_sink():
+        return False
+
+    current_agent = state.current_agent
+    if current_agent is None or current_agent.agent_name != agent_name:
+        return False
+
+    return state.server_controls is not None
+
+
 def _merge_results(
     local_result: EvaluationResponse,
     server_result: EvaluationResponse,
@@ -192,7 +218,7 @@ async def check_evaluation(
         The parsed evaluation result returned by the server.
     """
     normalized_name = ensure_agent_name(agent_name)
-    merged_emission_enabled = has_control_event_sink()
+    merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name)
     trace_id = None
     span_id = None
     headers: dict[str, str] | None = None
@@ -351,7 +377,7 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
         combined_errors = (result.errors or []) + parse_errors
         return result.model_copy(update={"errors": combined_errors})
 
-    merged_emission_enabled = has_control_event_sink()
+    merged_emission_enabled = _is_merged_event_mode_enabled(normalized_name)
     should_reconstruct_local_events = merged_emission_enabled or is_observability_enabled()
 
     local_result: EvaluationResponse | None = None

sdks/python/tests/test_init_step_merge.py

@@ -20,9 +20,11 @@ class DoesNotExist: ...
 @pytest.fixture(autouse=True)
 def _clean_registry() -> Generator[None, None, None]:
     """Ensure each test starts with an empty step registry."""
+    agent_control._reset_state()
     clear()
     yield
     clear()
+    agent_control._reset_state()
 
 
 def test_init_passes_merged_steps_to_register_agent(

sdks/python/tests/test_observability_updates.py

@@ -398,7 +398,7 @@ async def test_default_path_keeps_server_only_behavior(self):
         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), \
+        with patch("agent_control.evaluation._is_merged_event_mode_enabled", return_value=False), \
              patch("agent_control.evaluation.emit_control_events") as mock_emit:
             result = await evaluation.check_evaluation(
                 client=client,
@@ -455,7 +455,7 @@ async def test_merged_event_sink_emits_reconstructed_server_events(self):
         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), \
+        with patch("agent_control.evaluation._is_merged_event_mode_enabled", 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):
@@ -508,7 +508,7 @@ async def test_skips_local_event_reconstruction_when_nothing_consumes_events(sel
 
         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_merged_event_mode_enabled", 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:
@@ -525,6 +525,42 @@ async def test_skips_local_event_reconstruction_when_nothing_consumes_events(sel
         assert result.is_safe is True
         assert result.confidence == 1.0
 
+    @pytest.mark.asyncio
+    async def test_check_evaluation_falls_back_when_initialized_agent_does_not_match(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=True), \
+             patch("agent_control.evaluation.emit_control_events") as mock_emit, \
+             patch.object(evaluation.state, "current_agent", MagicMock(agent_name="agent-000000000002")), \
+             patch.object(evaluation.state, "server_controls", []):
+            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
+
 
 # =============================================================================
 # control_decorators non_matches dict conversion
@@ -606,7 +642,7 @@ async def test_merged_event_sink_emits_reconstructed_local_and_server_events_onc
 
         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=True), \
+             patch("agent_control.evaluation._is_merged_event_mode_enabled", return_value=True), \
              patch("agent_control.evaluation.emit_control_events") as mock_emit, \
              patch("agent_control.evaluation.enqueue_observability_events") as mock_enqueue:
             result = await evaluation.check_evaluation_with_local(

sdks/python/tests/test_shutdown.py

@@ -15,6 +15,7 @@
 import agent_control.observability as obs_mod
 from agent_control._state import state
 from agent_control.observability import EventBatcher
+from agent_control.telemetry.event_sink import has_control_event_sink, set_control_event_sink
 
 
 def _make_started_batcher() -> EventBatcher:
@@ -64,6 +65,7 @@ def test_shutdown_resets_state(self):
         state.server_controls = [{"name": "test"}]
         state.server_url = "http://localhost:8000"
         state.api_key = "key"
+        set_control_event_sink(lambda events: None)
 
         agent_control.shutdown()
 
@@ -72,6 +74,7 @@ def test_shutdown_resets_state(self):
         assert state.server_controls is None
         assert state.server_url is None
         assert state.api_key is None
+        assert has_control_event_sink() is False
 
     def test_shutdown_idempotent(self):
         agent_control.shutdown()