fix(platform): resolve correct span in build_trace_context_headers (#1679)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ctiliescuuipath

packages/uipath-platform/src/uipath/platform/chat/llm_trace_context.py

@@ -5,12 +5,17 @@
 from uipath.core.tracing.span_utils import UiPathSpanUtils
 
 from ..common._config import UiPathConfig
+from ..common._span_utils import _SpanUtils
 
 
 def build_trace_context_headers(
     extra_baggage: list[str] | None = None,
 ) -> dict[str, str]:
-    """Build W3C-style trace context headers from the current OpenTelemetry span.
+    """Build W3C-style trace context headers for LLM Gateway requests.
+
+    Resolves the current span via ``UiPathSpanUtils.get_external_current_span()``
+    (which returns the deepest active span from the LLMOps hierarchy) with a
+    fallback to ``trace.get_current_span()``.
 
     Args:
         extra_baggage: Additional baggage entries (e.g. ``["source=agents"]``)
@@ -25,16 +30,17 @@ def build_trace_context_headers(
     headers: dict[str, str] = {}
     llmops_span = UiPathSpanUtils.get_external_current_span()
     span = llmops_span or trace.get_current_span()
+    config_trace_id = UiPathConfig.trace_id
     ctx = span.get_span_context()
-    if ctx and ctx.trace_id and ctx.span_id:
-        trace_id = format(ctx.trace_id, "032x")
-        span_id = format(ctx.span_id, "016x")
+    if config_trace_id and ctx and ctx.span_id:
+        trace_id = _SpanUtils.normalize_trace_id(config_trace_id)
+        span_id = format(ctx.span_id, "032x")
         headers["x-uipath-traceparent-id"] = f"00-{trace_id}-{span_id}"
 
     baggage_parts: list[str] = list(extra_baggage) if extra_baggage else []
     if folder_key := UiPathConfig.folder_key:
         baggage_parts.append(f"folderKey={folder_key}")
-    if agent_id := UiPathConfig.process_uuid:
+    if agent_id := UiPathConfig.agent_id:
         baggage_parts.append(f"agentId={agent_id}")
     if process_key := UiPathConfig.process_key:
         baggage_parts.append(f"processKey={process_key}")

packages/uipath-platform/tests/services/test_llm_trace_context.py

@@ -3,7 +3,6 @@
 import os
 from unittest.mock import patch
 
-from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags
 from uipath.core.feature_flags import FeatureFlags
@@ -13,6 +12,13 @@
 FEATURE_FLAG = "EnableTraceContextHeaders"
 
 
+def _make_span():
+    """Create a real OTEL span for testing."""
+    provider = TracerProvider()
+    tracer = provider.get_tracer("test")
+    return tracer.start_span("test-span")
+
+
 class TestFeatureFlagDisabled:
     """When the feature flag is off, no headers are returned."""
 
@@ -28,48 +34,74 @@ def test_returns_empty_dict_when_explicitly_disabled(self) -> None:
 
 
 class TestTraceparentHeader:
-    """When enabled, x-uipath-traceparent-id is populated from the active span."""
+    """When enabled, x-uipath-traceparent-id is populated from config + span."""
 
     def setup_method(self) -> None:
         FeatureFlags.reset_flags()
         FeatureFlags.configure_flags({FEATURE_FLAG: True})
 
-    def test_traceparent_from_active_span(self) -> None:
-        provider = TracerProvider()
-        tracer = provider.get_tracer("test")
-        with tracer.start_as_current_span("test-span") as span:
-            ctx = span.get_span_context()
-            expected_trace_id = format(ctx.trace_id, "032x")
-            expected_span_id = format(ctx.span_id, "016x")
-
+    def test_traceparent_from_config_and_span(self) -> None:
+        span = _make_span()
+        ctx = span.get_span_context()
+        expected_span_id = format(ctx.span_id, "032x")
+        config_trace = "abcdef1234567890abcdef1234567890"
+        env = {"UIPATH_TRACE_ID": config_trace}
+        with (
+            patch.dict(os.environ, env),
+            patch(
+                "uipath.platform.chat.llm_trace_context.trace.get_current_span",
+                return_value=span,
+            ),
+        ):
             headers = build_trace_context_headers()
 
         assert "x-uipath-traceparent-id" in headers
         value = headers["x-uipath-traceparent-id"]
-        assert value == f"00-{expected_trace_id}-{expected_span_id}"
-        # Verify format: version (2) + dash + trace_id (32) + dash + span_id (16)
+        assert value == f"00-{config_trace}-{expected_span_id}"
         parts = value.split("-")
         assert len(parts) == 3
         assert parts[0] == "00"
         assert len(parts[1]) == 32
-        assert len(parts[2]) == 16
+        assert len(parts[2]) == 32
+
+    def test_no_traceparent_without_config_trace_id(self) -> None:
+        headers = build_trace_context_headers()
+        assert "x-uipath-traceparent-id" not in headers
+
+    def test_traceparent_strips_dashes_from_config_trace_id(self) -> None:
+        span = _make_span()
+        uuid_trace = "abcdef12-3456-7890-abcd-ef1234567890"
+        env = {"UIPATH_TRACE_ID": uuid_trace}
+        with (
+            patch.dict(os.environ, env),
+            patch(
+                "uipath.platform.chat.llm_trace_context.trace.get_current_span",
+                return_value=span,
+            ),
+        ):
+            headers = build_trace_context_headers()
 
-    def test_no_traceparent_without_active_span(self) -> None:
-        # INVALID_SPAN has trace_id=0 and span_id=0
-        from opentelemetry.context import attach, detach
+        value = headers["x-uipath-traceparent-id"]
+        parts = value.split("-")
+        assert parts[1] == "abcdef1234567890abcdef1234567890"
 
+    def test_no_traceparent_with_invalid_span(self) -> None:
         ctx = SpanContext(
             trace_id=0,
             span_id=0,
             is_remote=False,
             trace_flags=TraceFlags(0),
         )
-        non_recording = NonRecordingSpan(ctx)
-        token = attach(trace.set_span_in_context(non_recording))
-        try:
+        span = NonRecordingSpan(ctx)
+        env = {"UIPATH_TRACE_ID": "abcdef1234567890abcdef1234567890"}
+        with (
+            patch.dict(os.environ, env),
+            patch(
+                "uipath.platform.chat.llm_trace_context.trace.get_current_span",
+                return_value=span,
+            ),
+        ):
             headers = build_trace_context_headers()
-        finally:
-            detach(token)
 
         assert "x-uipath-traceparent-id" not in headers
 
@@ -84,7 +116,7 @@ def setup_method(self) -> None:
     def test_all_env_vars_present(self) -> None:
         env = {
             "UIPATH_FOLDER_KEY": "folder-abc",
-            "UIPATH_PROCESS_UUID": "agent-123",
+            "UIPATH_AGENT_ID": "agent-123",
             "UIPATH_PROCESS_KEY": "process-789",
         }
         with patch.dict(os.environ, env, clear=True):
@@ -103,6 +135,31 @@ def test_partial_env_vars(self) -> None:
         baggage = headers["x-uipath-tracebaggage"]
         assert "folderKey=folder-only" in baggage
 
+    def test_agent_id_from_agent_id_env(self) -> None:
+        env = {"UIPATH_AGENT_ID": "real-agent-id"}
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "agentId=real-agent-id" in baggage
+
+    def test_agent_id_falls_back_to_project_id(self) -> None:
+        env = {"UIPATH_PROJECT_ID": "project-123"}
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "agentId=project-123" in baggage
+
+    def test_no_agent_id_without_env_vars(self) -> None:
+        env = {"UIPATH_FOLDER_KEY": "f1"}
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "agentId" not in baggage
+        assert "folderKey=f1" in baggage
+
     def test_no_baggage_without_env_vars(self) -> None:
         with patch.dict(os.environ, {}, clear=True):
             headers = build_trace_context_headers()
@@ -112,7 +169,7 @@ def test_no_baggage_without_env_vars(self) -> None:
     def test_baggage_comma_separated(self) -> None:
         env = {
             "UIPATH_FOLDER_KEY": "f1",
-            "UIPATH_PROCESS_UUID": "a1",
+            "UIPATH_AGENT_ID": "a1",
         }
         with patch.dict(os.environ, env, clear=True):
             headers = build_trace_context_headers()
@@ -148,14 +205,22 @@ def setup_method(self) -> None:
         FeatureFlags.configure_flags({FEATURE_FLAG: True})
 
     def test_both_headers_present(self) -> None:
-        provider = TracerProvider()
-        tracer = provider.get_tracer("test")
-        env = {"UIPATH_FOLDER_KEY": "folder-abc"}
+        span = _make_span()
+        env = {
+            "UIPATH_FOLDER_KEY": "folder-abc",
+            "UIPATH_TRACE_ID": "abcdef1234567890abcdef1234567890",
+        }
         with (
-            tracer.start_as_current_span("test-span"),
             patch.dict(os.environ, env, clear=True),
+            patch(
+                "uipath.platform.chat.llm_trace_context.trace.get_current_span",
+                return_value=span,
+            ),
         ):
             headers = build_trace_context_headers()
 
         assert "x-uipath-traceparent-id" in headers
+        assert headers["x-uipath-traceparent-id"].startswith(
+            "00-abcdef1234567890abcdef1234567890-"
+        )
         assert "x-uipath-tracebaggage" in headers