fix: stamp reference hierarchy as span attribute to survive BatchSpanProcessor thread boundary

ContextVar values are not propagated to background threads. The
BatchSpanProcessor exports spans in a worker thread where
ReferenceContextAccessor.get() always returns None, so Context.referenceHierarchy
was never populated in the wire payload.

Fix: register a span-start hook (_inject_reference_hierarchy) that runs
synchronously in on_start — same thread/context as the span creator — reads
the ambient ReferenceContext and stamps it as uipath.reference_hierarchy on the
span. The attribute travels with the span to the export thread.
otel_span_to_uipath_span now reads from the attribute instead of the ContextVar,
and pops it from attributes_dict so it does not appear in the wire Attributes field.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: jepadil23

packages/uipath-core/src/uipath/core/tracing/processors.py

@@ -1,6 +1,6 @@
 """Custom span processors for UiPath execution tracing."""
 
-from typing import cast
+from typing import Callable, cast
 
 from opentelemetry import context as context_api
 from opentelemetry import trace
@@ -13,6 +13,21 @@
 
 from uipath.core.tracing.types import UiPathTraceSettings
 
+# Hooks called synchronously in on_start (same thread as span creation, where
+# ContextVar values are still live). Use register_span_start_hook to stamp
+# span attributes that must survive the BatchSpanProcessor thread boundary.
+_span_start_hooks: list[Callable[[Span], None]] = []
+
+
+def register_span_start_hook(hook: Callable[[Span], None]) -> None:
+    """Register a callable invoked for every span at creation time.
+
+    The hook receives the live, writable Span so it can call set_attribute.
+    Runs in the same thread/context as the span creator — ContextVar values
+    are available here but NOT in the BatchSpanProcessor export thread.
+    """
+    _span_start_hooks.append(hook)
+
 
 class UiPathExecutionTraceProcessorMixin:
     """Mixin that propagates execution.id and optionally filters spans."""
@@ -32,6 +47,9 @@ def on_start(self, span: Span, parent_context: context_api.Context | None = None
             if execution_id:
                 span.set_attribute("execution.id", execution_id)
 
+        for hook in _span_start_hooks:
+            hook(span)
+
     def on_end(self, span: ReadableSpan):
         """Called when a span ends. Filters before delegating to parent."""
         span_filter = self._settings.span_filter if self._settings else None
@@ -73,4 +91,5 @@ def __init__(
 __all__ = [
     "UiPathExecutionBatchTraceProcessor",
     "UiPathExecutionSimpleTraceProcessor",
+    "register_span_start_hook",
 ]

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

@@ -8,13 +8,25 @@
 from os import environ as env
 from typing import Any, Dict, List, Optional
 
-from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace import ReadableSpan, Span
 from opentelemetry.trace import StatusCode
 from pydantic import BaseModel, ConfigDict, Field
 from uipath.core.serialization import serialize_json
+from uipath.core.tracing.processors import register_span_start_hook
 
 from ._reference_context import ReferenceContextAccessor
 
+
+def _inject_reference_hierarchy(span: Span) -> None:
+    ref_ctx = ReferenceContextAccessor.get()
+    if ref_ctx:
+        wire = ref_ctx.to_wire_list()
+        if wire:
+            span.set_attribute("uipath.reference_hierarchy", json.dumps(wire))
+
+
+register_span_start_hook(_inject_reference_hierarchy)
+
 logger = logging.getLogger(__name__)
 
 # SourceEnum.CodedAgents = 10 (default for Python SDK / coded agents)
@@ -232,6 +244,11 @@ def otel_span_to_uipath_span(
         # Only copy if we need to modify - we'll build attributes_dict lazily
         attributes_dict: dict[str, Any] = dict(otel_attrs) if otel_attrs else {}
 
+        # Pull the reference hierarchy stamped by the span-start hook (runs in the
+        # correct thread/context; BatchSpanProcessor exports in a background thread
+        # where ContextVar values are not available).
+        ref_hierarchy_json = attributes_dict.pop("uipath.reference_hierarchy", None)
+
         # Map status
         status = 1  # Default to OK
         if otel_span.status.status_code == StatusCode.ERROR:
@@ -331,14 +348,12 @@ def otel_span_to_uipath_span(
             except Exception as e:
                 logger.warning(f"Error processing attachments: {e}")
 
-        # Build Context.referenceHierarchy from the ambient ReferenceContext
-        # (set by the agent runtime at run start via ReferenceContextAccessor).
         context: Optional[Dict[str, Any]] = None
-        ref_ctx = ReferenceContextAccessor.get()
-        if ref_ctx:
-            wire_list = ref_ctx.to_wire_list()
-            if wire_list:
-                context = {"referenceHierarchy": wire_list}
+        if ref_hierarchy_json:
+            try:
+                context = {"referenceHierarchy": json.loads(ref_hierarchy_json)}
+            except (json.JSONDecodeError, TypeError):
+                pass
 
         # Create UiPathSpan from OpenTelemetry span
         start_time = datetime.fromtimestamp(

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

@@ -1,5 +1,6 @@
 """Tests for ReferenceContext, ReferenceContextAccessor, and span Context wiring."""
 
+import json
 import os
 from datetime import datetime
 from unittest.mock import Mock
@@ -247,14 +248,9 @@ def test_context_absent_when_no_reference_context_set(
         self, monkeypatch: pytest.MonkeyPatch
     ) -> None:
         monkeypatch.setenv("UIPATH_ORGANIZATION_ID", "test-org")
-        # Ensure accessor is clear
-        token = ReferenceContextAccessor.set(None)
-        try:
-            span = _SpanUtils.otel_span_to_uipath_span(_make_mock_span())
-            assert span.context is None
-            assert "Context" not in span.to_dict()
-        finally:
-            ReferenceContextAccessor.reset(token)
+        span = _SpanUtils.otel_span_to_uipath_span(_make_mock_span())
+        assert span.context is None
+        assert "Context" not in span.to_dict()
 
     def test_context_present_when_reference_context_set(
         self, monkeypatch: pytest.MonkeyPatch
@@ -263,23 +259,22 @@ def test_context_present_when_reference_context_set(
         ref_ctx = ReferenceContext.Empty.add(
             "agent", "550e8400-e29b-41d4-a716-446655440001", "1.0"
         )
-        token = ReferenceContextAccessor.set(ref_ctx)
-        try:
-            span = _SpanUtils.otel_span_to_uipath_span(_make_mock_span())
-            assert span.context == {
-                "referenceHierarchy": [
-                    {
-                        "serviceType": "agent",
-                        "referenceId": "550e8400-e29b-41d4-a716-446655440001",
-                        "version": "1.0",
-                    }
-                ]
-            }
-            wire = span.to_dict()
-            assert "Context" in wire
-            assert wire["Context"]["referenceHierarchy"][0]["serviceType"] == "agent"
-        finally:
-            ReferenceContextAccessor.reset(token)
+        mock = _make_mock_span(
+            attributes={"uipath.reference_hierarchy": json.dumps(ref_ctx.to_wire_list())}
+        )
+        span = _SpanUtils.otel_span_to_uipath_span(mock)
+        assert span.context == {
+            "referenceHierarchy": [
+                {
+                    "serviceType": "agent",
+                    "referenceId": "550e8400-e29b-41d4-a716-446655440001",
+                    "version": "1.0",
+                }
+            ]
+        }
+        wire = span.to_dict()
+        assert "Context" in wire
+        assert wire["Context"]["referenceHierarchy"][0]["serviceType"] == "agent"
 
     def test_context_carries_full_hierarchy(
         self, monkeypatch: pytest.MonkeyPatch
@@ -290,14 +285,66 @@ def test_context_carries_full_hierarchy(
             .add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.0")
             .add("agent", "550e8400-e29b-41d4-a716-446655440011")
         )
+        mock = _make_mock_span(
+            attributes={"uipath.reference_hierarchy": json.dumps(ref_ctx.to_wire_list())}
+        )
+        wire = _SpanUtils.otel_span_to_uipath_span(mock).to_dict()
+        hierarchy = wire["Context"]["referenceHierarchy"]
+        assert len(hierarchy) == 2
+        assert hierarchy[0]["serviceType"] == "maestro"
+        assert hierarchy[0]["version"] == "2.0"
+        assert hierarchy[1]["serviceType"] == "agent"
+        assert "version" not in hierarchy[1]
+
+    def test_reference_hierarchy_not_in_attributes_field(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setenv("UIPATH_ORGANIZATION_ID", "test-org")
+        ref_ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001")
+        mock = _make_mock_span(
+            attributes={"uipath.reference_hierarchy": json.dumps(ref_ctx.to_wire_list())}
+        )
+        wire = _SpanUtils.otel_span_to_uipath_span(mock).to_dict()
+        attributes = json.loads(wire["Attributes"])
+        assert "uipath.reference_hierarchy" not in attributes
+
+
+# ---------------------------------------------------------------------------
+# _inject_reference_hierarchy hook
+# ---------------------------------------------------------------------------
+
+class TestReferenceHierarchyHook:
+    def setup_method(self) -> None:
+        token = ReferenceContextAccessor.set(None)
+        ReferenceContextAccessor.reset(token)
+
+    def test_hook_stamps_attribute_when_context_set(self) -> None:
+        from uipath.platform.common._span_utils import _inject_reference_hierarchy
+
+        ref_ctx = ReferenceContext.Empty.add(
+            "agent", "550e8400-e29b-41d4-a716-446655440001"
+        )
         token = ReferenceContextAccessor.set(ref_ctx)
         try:
-            wire = _SpanUtils.otel_span_to_uipath_span(_make_mock_span()).to_dict()
-            hierarchy = wire["Context"]["referenceHierarchy"]
-            assert len(hierarchy) == 2
-            assert hierarchy[0]["serviceType"] == "maestro"
-            assert hierarchy[0]["version"] == "2.0"
-            assert hierarchy[1]["serviceType"] == "agent"
-            assert "version" not in hierarchy[1]
+            mock_span = Mock()
+            _inject_reference_hierarchy(mock_span)
+            mock_span.set_attribute.assert_called_once()
+            key, value = mock_span.set_attribute.call_args[0]
+            assert key == "uipath.reference_hierarchy"
+            hierarchy = json.loads(value)
+            assert len(hierarchy) == 1
+            assert hierarchy[0]["serviceType"] == "agent"
+            assert hierarchy[0]["referenceId"] == "550e8400-e29b-41d4-a716-446655440001"
+        finally:
+            ReferenceContextAccessor.reset(token)
+
+    def test_hook_noop_when_context_not_set(self) -> None:
+        from uipath.platform.common._span_utils import _inject_reference_hierarchy
+
+        token = ReferenceContextAccessor.set(None)
+        try:
+            mock_span = Mock()
+            _inject_reference_hierarchy(mock_span)
+            mock_span.set_attribute.assert_not_called()
         finally:
             ReferenceContextAccessor.reset(token)