feat(platform): add W3C trace context headers to LLM Gateway requests

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.43"
+version = "0.1.44"
 description = "HTTP client library for programmatic access to UiPath Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

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

@@ -35,6 +35,7 @@
     ToolDefinition,
 )
 from .llm_throttle import get_llm_semaphore
+from .llm_trace_context import build_trace_context_headers
 
 # Common constants
 API_VERSION = "2024-10-21"  # Standard API version for OpenAI-compatible endpoints
@@ -224,7 +225,7 @@ async def embeddings(
                 endpoint,
                 json={"input": input},
                 params={"api-version": API_VERSION},
-                headers=self._llm_headers,
+                headers={**self._llm_headers, **build_trace_context_headers()},
             )
 
         return TextEmbedding.model_validate(response.json())
@@ -355,7 +356,7 @@ class Country(BaseModel):
                 endpoint,
                 json=request_body,
                 params={"api-version": API_VERSION},
-                headers=self._llm_headers,
+                headers={**self._llm_headers, **build_trace_context_headers()},
             )
 
         return ChatCompletion.model_validate(response.json())
@@ -599,6 +600,7 @@ class Country(BaseModel):
 
         headers = {
             **self._llm_headers,
+            **build_trace_context_headers(),
             "X-UiPath-LlmGateway-NormalizedApi-ModelName": model,
             "X-UiPath-LLMGateway-AllowFull4xxResponse": "true",
         }

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

@@ -0,0 +1,42 @@
+"""W3C-style trace context headers for LLM Gateway requests."""
+
+from opentelemetry import trace
+from uipath.core.feature_flags import FeatureFlags
+
+from ..common._config import UiPathConfig
+
+
+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.
+
+    Args:
+        extra_baggage: Additional baggage entries (e.g. ``["source=agents"]``)
+            that callers can inject alongside the platform-level entries.
+
+    Returns an empty dict when the ``EnableTraceContextHeaders`` feature flag
+    is not enabled, or when no active span is present.
+    """
+    if not FeatureFlags.is_flag_enabled("EnableTraceContextHeaders"):
+        return {}
+
+    headers: dict[str, str] = {}
+    span = trace.get_current_span()
+    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")
+        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:
+        baggage_parts.append(f"agentId={agent_id}")
+    if process_key := UiPathConfig.process_key:
+        baggage_parts.append(f"processKey={process_key}")
+    if baggage_parts:
+        headers["x-uipath-tracebaggage"] = ",".join(baggage_parts)
+
+    return headers

packages/uipath-platform/src/uipath/platform/common/_config.py

@@ -121,6 +121,12 @@ def folder_path(self) -> str | None:
 
         return os.getenv(ENV_FOLDER_PATH, None)
 
+    @property
+    def process_key(self) -> str | None:
+        from uipath.platform.common.constants import ENV_PROCESS_KEY
+
+        return os.getenv(ENV_PROCESS_KEY, None)
+
     @property
     def process_uuid(self) -> str | None:
         from uipath.platform.common.constants import ENV_UIPATH_PROCESS_UUID

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

@@ -0,0 +1,161 @@
+"""Tests for build_trace_context_headers."""
+
+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
+
+from uipath.platform.chat.llm_trace_context import build_trace_context_headers
+
+FEATURE_FLAG = "EnableTraceContextHeaders"
+
+
+class TestFeatureFlagDisabled:
+    """When the feature flag is off, no headers are returned."""
+
+    def setup_method(self) -> None:
+        FeatureFlags.reset_flags()
+
+    def test_returns_empty_dict_by_default(self) -> None:
+        assert build_trace_context_headers() == {}
+
+    def test_returns_empty_dict_when_explicitly_disabled(self) -> None:
+        FeatureFlags.configure_flags({FEATURE_FLAG: False})
+        assert build_trace_context_headers() == {}
+
+
+class TestTraceparentHeader:
+    """When enabled, x-uipath-traceparent-id is populated from the active 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")
+
+            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)
+        parts = value.split("-")
+        assert len(parts) == 3
+        assert parts[0] == "00"
+        assert len(parts[1]) == 32
+        assert len(parts[2]) == 16
+
+    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
+
+        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:
+            headers = build_trace_context_headers()
+        finally:
+            detach(token)
+
+        assert "x-uipath-traceparent-id" not in headers
+
+
+class TestBaggageHeader:
+    """When enabled, x-uipath-tracebaggage is populated from UiPathConfig."""
+
+    def setup_method(self) -> None:
+        FeatureFlags.reset_flags()
+        FeatureFlags.configure_flags({FEATURE_FLAG: True})
+
+    def test_all_env_vars_present(self) -> None:
+        env = {
+            "UIPATH_FOLDER_KEY": "folder-abc",
+            "UIPATH_PROCESS_UUID": "agent-123",
+            "UIPATH_PROCESS_KEY": "process-789",
+        }
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "folderKey=folder-abc" in baggage
+        assert "agentId=agent-123" in baggage
+        assert "processKey=process-789" in baggage
+
+    def test_partial_env_vars(self) -> None:
+        env = {"UIPATH_FOLDER_KEY": "folder-only"}
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "folderKey=folder-only" in baggage
+
+    def test_no_baggage_without_env_vars(self) -> None:
+        with patch.dict(os.environ, {}, clear=True):
+            headers = build_trace_context_headers()
+
+        assert "x-uipath-tracebaggage" not in headers
+
+    def test_baggage_comma_separated(self) -> None:
+        env = {
+            "UIPATH_FOLDER_KEY": "f1",
+            "UIPATH_PROCESS_UUID": "a1",
+        }
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers()
+
+        baggage = headers["x-uipath-tracebaggage"]
+        parts = baggage.split(",")
+        assert len(parts) == 2  # folderKey + agentId
+
+    def test_extra_baggage_included(self) -> None:
+        env = {"UIPATH_FOLDER_KEY": "f1"}
+        with patch.dict(os.environ, env, clear=True):
+            headers = build_trace_context_headers(extra_baggage=["source=agents"])
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert "source=agents" in baggage
+        assert "folderKey=f1" in baggage
+
+    def test_extra_baggage_only(self) -> None:
+        with patch.dict(os.environ, {}, clear=True):
+            headers = build_trace_context_headers(
+                extra_baggage=["source=agents", "custom=value"]
+            )
+
+        baggage = headers["x-uipath-tracebaggage"]
+        assert baggage == "source=agents,custom=value"
+
+
+class TestBothHeaders:
+    """When enabled with an active span and env vars, both headers are present."""
+
+    def setup_method(self) -> None:
+        FeatureFlags.reset_flags()
+        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"}
+        with (
+            tracer.start_as_current_span("test-span"),
+            patch.dict(os.environ, env, clear=True),
+        ):
+            headers = build_trace_context_headers()
+
+        assert "x-uipath-traceparent-id" in headers
+        assert "x-uipath-tracebaggage" in headers