fix(sdk): forward target context from decorators

Author: abhinav-galileo

sdks/python/src/agent_control/control_decorators.py

@@ -39,7 +39,8 @@ async def chat(message: str) -> str:
 from agent_control_telemetry import get_trace_context_from_provider
 
 from agent_control import AgentControlClient
-from agent_control.evaluation import check_evaluation_with_local
+from agent_control._state import state
+from agent_control.evaluation import _resolve_session_target, check_evaluation_with_local
 from agent_control.observability import (
     get_logger,
     log_control_evaluation,
@@ -269,7 +270,14 @@ async def _evaluate(
     if span_id:
         headers["X-Span-Id"] = span_id
 
-    async with AgentControlClient(base_url=server_url) as client:
+    target_type, target_id = _resolve_session_target(None, None)
+
+    async with AgentControlClient(
+        base_url=server_url,
+        api_key=state.api_key,
+        api_key_header=state.api_key_header,
+        runtime_token_cache=state.runtime_token_cache,
+    ) as client:
         # If we have controls, use local evaluation which handles both SDK and server controls
         if controls is not None:
             try:
@@ -282,6 +290,8 @@ async def _evaluate(
                     step=step_obj,
                     stage=stage,  # type: ignore
                     controls=controls,
+                    target_type=target_type,
+                    target_id=target_id,
                     trace_id=trace_id,
                     span_id=span_id,
                     event_agent_name=event_agent_name,
@@ -369,13 +379,18 @@ async def _evaluate(
                 )
 
         # Fallback: server-only evaluation
+        payload = {
+            "agent_name": str(agent_name),
+            "step": step,
+            "stage": stage,
+        }
+        if target_type is not None and target_id is not None:
+            payload["target_type"] = target_type
+            payload["target_id"] = target_id
+
         response = await client.http_client.post(
             "/api/v1/evaluation",
-            json={
-                "agent_name": str(agent_name),
-                "step": step,
-                "stage": stage
-            },
+            json=payload,
             headers=headers,
         )
         response.raise_for_status()

sdks/python/tests/test_control_decorators.py

@@ -793,6 +793,65 @@ async def search_tool(query: str) -> str:
             assert captured_steps[0]["type"] == "tool"
 
 
+class TestEvaluateSessionContext:
+    """Tests for session context forwarding in the decorator evaluation path."""
+
+    @pytest.mark.asyncio
+    async def test_evaluate_forwards_session_target_and_client_state(self):
+        """The decorator path should use the same session target as evaluate_controls()."""
+        from agent_control.control_decorators import _evaluate
+        from agent_control.runtime_auth import RuntimeTokenCache
+        from agent_control_models import EvaluationResult
+
+        captured_client_kwargs = {}
+        captured_eval_kwargs = {}
+        runtime_token_cache = RuntimeTokenCache()
+
+        class FakeClient:
+            def __init__(self, **kwargs):
+                captured_client_kwargs.update(kwargs)
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, exc_type, exc, tb):
+                return None
+
+        async def fake_check_evaluation_with_local(**kwargs):
+            captured_eval_kwargs.update(kwargs)
+            return EvaluationResult(is_safe=True, confidence=1.0)
+
+        with (
+            patch("agent_control.control_decorators.AgentControlClient", FakeClient),
+            patch(
+                "agent_control.control_decorators.check_evaluation_with_local",
+                side_effect=fake_check_evaluation_with_local,
+            ),
+            patch("agent_control.control_decorators.state.api_key", "session-key"),
+            patch("agent_control.control_decorators.state.api_key_header", "X-Custom-Key"),
+            patch(
+                "agent_control.control_decorators.state.runtime_token_cache",
+                runtime_token_cache,
+            ),
+            patch("agent_control.control_decorators.state.target_type", "env"),
+            patch("agent_control.control_decorators.state.target_id", "prod"),
+        ):
+            result = await _evaluate(
+                agent_name="agent",
+                step={"type": "llm", "name": "chat", "input": "hello"},
+                stage="pre",
+                server_url="http://server.test",
+                controls=[{"id": 1, "name": "control", "control": {}}],
+            )
+
+        assert result["is_safe"] is True
+        assert captured_client_kwargs["api_key"] == "session-key"
+        assert captured_client_kwargs["api_key_header"] == "X-Custom-Key"
+        assert captured_client_kwargs["runtime_token_cache"] is runtime_token_cache
+        assert captured_eval_kwargs["target_type"] == "env"
+        assert captured_eval_kwargs["target_id"] == "prod"
+
+
 # =============================================================================
 # STEERING CONTEXT TESTS
 # =============================================================================