address comments

Author: namrataghadi-galileo

sdks/python/src/agent_control/evaluation.py

@@ -20,6 +20,7 @@
 from ._state import state
 from .client import AgentControlClient
 from .observability import add_event, get_logger, is_observability_enabled
+from .tracing import get_trace_and_span_ids
 from .validation import ensure_agent_name
 
 _logger = get_logger(__name__)
@@ -291,6 +292,13 @@ async def check_evaluation_with_local(
         httpx.HTTPError: If server request fails
     """
     normalized_name = ensure_agent_name(agent_name)
+    resolved_trace_id = trace_id
+    resolved_span_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
+
     # Partition controls by local flag
     local_controls: list[_ControlAdapter] = []
     parse_errors: list[ControlMatch] = []
@@ -389,8 +397,8 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
             local_result,
             request,
             applicable_local_controls,
-            trace_id,
-            span_id,
+            resolved_trace_id,
+            resolved_span_id,
             agent_name=event_agent_name,
         )
 
@@ -409,10 +417,10 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
     if _has_applicable_prefiltered_server_controls(server_control_payloads, request):
         request_payload = request.model_dump(mode="json", exclude_none=True)
         headers: dict[str, str] = {}
-        if trace_id:
-            headers["X-Trace-Id"] = trace_id
-        if span_id:
-            headers["X-Span-Id"] = span_id
+        if resolved_trace_id:
+            headers["X-Trace-Id"] = resolved_trace_id
+        if resolved_span_id:
+            headers["X-Span-Id"] = resolved_span_id
 
         response = await client.http_client.post(
             "/api/v1/evaluation",

sdks/python/src/agent_control/telemetry/trace_context.py

@@ -40,6 +40,8 @@ def get_trace_context_from_provider() -> TraceContext | None:
     span_id = trace_context.get("span_id")
     if not isinstance(trace_id, str) or not isinstance(span_id, str):
         return None
+    if not trace_id or not span_id:
+        return None
 
     return {
         "trace_id": trace_id,

sdks/python/tests/test_observability_updates.py

@@ -10,6 +10,10 @@
     _map_applies_to,
     _merge_results,
 )
+from agent_control.telemetry.trace_context import (
+    clear_trace_context_provider,
+    set_trace_context_provider,
+)
 from agent_control_models import ControlDefinition
 
 # =============================================================================
@@ -326,6 +330,9 @@ def test_fallback_warning_logged_only_once(self):
 class TestCheckEvaluationWithLocal:
     """Tests for check_evaluation_with_local event emission and non_matches."""
 
+    def teardown_method(self) -> None:
+        clear_trace_context_provider()
+
     @pytest.mark.asyncio
     async def test_emits_events_when_trace_context_provided(self):
         """Should emit observability events when trace_id and span_id are passed."""
@@ -398,7 +405,7 @@ async def test_emits_events_when_trace_context_provided(self):
 
     @pytest.mark.asyncio
     async def test_emits_events_without_trace_context(self):
-        """Should still emit events when trace_id/span_id not provided (fallback IDs)."""
+        """Should resolve trace context from the provider when IDs are omitted."""
         from agent_control_models import EvaluationResponse, Step
 
         mock_response = EvaluationResponse(
@@ -424,6 +431,12 @@ async def test_emits_events_without_trace_context(self):
         client = MagicMock()
         client.http_client = AsyncMock()
         step = Step(type="llm", name="test-step", input="hello")
+        set_trace_context_provider(
+            lambda: {
+                "trace_id": "a" * 32,
+                "span_id": "b" * 16,
+            }
+        )
 
         with patch("agent_control.evaluation.ControlEngine", return_value=mock_engine), \
              patch("agent_control.evaluation.list_evaluators", return_value=["regex"]), \
@@ -438,8 +451,8 @@ async def test_emits_events_without_trace_context(self):
             )
             mock_emit.assert_called_once()
             call_args = mock_emit.call_args
-            assert call_args[0][3] is None  # trace_id passed as None
-            assert call_args[0][4] is None  # span_id passed as None
+            assert call_args[0][3] == "a" * 32
+            assert call_args[0][4] == "b" * 16
 
     @pytest.mark.asyncio
     async def test_forwards_trace_headers_to_server(self):
@@ -492,6 +505,59 @@ async def test_forwards_trace_headers_to_server(self):
         assert headers["X-Trace-Id"] == "aaaa1111bbbb2222cccc3333dddd4444"
         assert headers["X-Span-Id"] == "eeee5555ffff6666"
 
+    @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 = [{
+            "id": 1,
+            "name": "server-ctrl",
+            "control": {
+                "condition": {
+                    "evaluator": {"name": "regex", "config": {"pattern": "test"}},
+                    "selector": {"path": "input"},
+                },
+                "action": {"decision": "deny"},
+                "execution": "server",
+            },
+        }]
+
+        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()
+
+        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(
+                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
+
 
 # =============================================================================
 # control_decorators non_matches dict conversion

sdks/python/tests/test_trace_context.py

@@ -52,3 +52,14 @@ def test_get_trace_context_from_provider_returns_none_for_invalid_shape() -> Non
     )
 
     assert get_trace_context_from_provider() is None
+
+
+def test_get_trace_context_from_provider_returns_none_for_empty_ids() -> None:
+    set_trace_context_provider(
+        lambda: {
+            "trace_id": "",
+            "span_id": "",
+        }
+    )
+
+    assert get_trace_context_from_provider() is None