feat: add ReferenceContext for span reference hierarchy propagation

packages/uipath-platform/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-platform"
-version = "0.1.91"
+version = "0.1.92"
 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/common/__init__.py

@@ -20,9 +20,15 @@
 from ._folder_context import FolderContext, header_folder
 from ._http_config import get_ca_bundle_path, get_httpx_client_kwargs
 from ._models import Endpoint, RequestSpec
+from ._reference_context import (
+    ReferenceContext,
+    ReferenceContextAccessor,
+    ReferenceEntry,
+)
 from ._service_url_overrides import inject_routing_headers, resolve_service_url
 from ._span_utils import (
     ExecutionType,
+    ReferenceHierarchySpanProcessor,
     SpanSource,
     SpanStatus,
     UiPathSpan,
@@ -117,6 +123,10 @@
     "ResourceOverwrite",
     "ResourceOverwriteParser",
     "ResourceOverwritesContext",
+    "ReferenceEntry",
+    "ReferenceContext",
+    "ReferenceContextAccessor",
+    "ReferenceHierarchySpanProcessor",
     "SpanSource",
     "SpanStatus",
     "UiPathSpan",

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

@@ -0,0 +1,260 @@
+"""Immutable reference-hierarchy context for span propagation.
+
+Follows the same design as service-common BaggageContext:
+- Immutable, copy-on-write — each mutating call returns a NEW instance so
+  sibling spans cannot bleed context into each other.
+- ContextVar-backed accessor — flows across await boundaries without
+  threading the value through every function signature.
+- Wire format compatible with the ``ref.*`` keys in ``x-uipath-tracebaggage``
+  so context parsed by service-common middleware is understood here and
+  vice-versa.
+"""
+from __future__ import annotations
+
+import contextvars
+import uuid
+from dataclasses import dataclass
+from typing import ClassVar, Dict, Iterator, List, Optional, Tuple
+
+
+__all__ = [
+    "ReferenceEntry",
+    "ReferenceContext",
+    "ReferenceContextAccessor",
+    "BAGGAGE_HEADER_NAME",
+    "BAGGAGE_KEY_TYPE",
+    "BAGGAGE_KEY_ID",
+    "BAGGAGE_KEY_VERSION",
+]
+
+BAGGAGE_HEADER_NAME = "x-uipath-tracebaggage"
+
+# Key names — matches service-common ReferenceHierarchyKeys
+BAGGAGE_KEY_TYPE = "ref.type"
+BAGGAGE_KEY_ID = "ref.id"
+BAGGAGE_KEY_VERSION = "ref.v"
+
+
+@dataclass(frozen=True)
+class ReferenceEntry:
+    """A single node in the reference hierarchy call chain."""
+
+    service_type: str
+    reference_id: str  # UUID string
+    version: Optional[str] = None
+
+
+class ReferenceContext:
+    """Immutable, copy-on-write ordered list of reference entries.
+
+    Outermost caller first, current service appended last.
+    Each mutating call returns a new instance — the original is never
+    modified, preventing sibling spans from sharing context.
+
+    Usage::
+
+        ctx = ReferenceContext.Empty
+        ctx = ctx.add("maestro", process_id, "2.1.0")
+        ctx = ctx.add("agent", agent_id)
+        token = ReferenceContextAccessor.set(ctx)
+        try:
+            ...
+        finally:
+            ReferenceContextAccessor.reset(token)
+    """
+
+    Empty: ClassVar["ReferenceContext"]
+    __slots__ = ("_entries",)
+
+    def __init__(self, entries: Tuple[ReferenceEntry, ...] = ()) -> None:
+        self._entries: Tuple[ReferenceEntry, ...] = entries
+
+    @property
+    def entries(self) -> Tuple[ReferenceEntry, ...]:
+        return self._entries
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def __iter__(self) -> Iterator[ReferenceEntry]:
+        return iter(self._entries)
+
+    def __bool__(self) -> bool:
+        return len(self._entries) > 0
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ReferenceContext):
+            return NotImplemented
+        return self._entries == other._entries
+
+    def __hash__(self) -> int:
+        return hash(self._entries)
+
+    def add(
+        self,
+        service_type: str,
+        reference_id: str | uuid.UUID,
+        version: Optional[str] = None,
+    ) -> "ReferenceContext":
+        """Returns a new context with this entry appended (copy-on-write).
+
+        Args:
+            service_type: Identifier for the calling service (e.g. ``"agent"``,
+                ``"maestro"``).
+            reference_id: UUID of the referenced entity (UUID object or string).
+            version: Optional version string.
+
+        Returns:
+            A new :class:`ReferenceContext` with the entry appended.
+        """
+        if not service_type or not service_type.strip():
+            raise ValueError("service_type must be a non-empty string.")
+        if isinstance(reference_id, uuid.UUID):
+            id_str = str(reference_id)
+        elif isinstance(reference_id, str):
+            try:
+                id_str = str(uuid.UUID(reference_id))
+            except ValueError:
+                raise ValueError(
+                    f"reference_id {reference_id!r} is not a valid UUID."
+                )
+        else:
+            raise TypeError("reference_id must be a UUID or string.")
+        entry = ReferenceEntry(
+            service_type=service_type,
+            reference_id=id_str,
+            version=version if version and version.strip() else None,
+        )
+        return ReferenceContext(self._entries + (entry,))
+
+    def to_wire_list(self) -> List[Dict[str, str]]:
+        """Serialize to the ``referenceHierarchy`` wire format.
+
+        Returns:
+            A list of dicts suitable for JSON serialization as
+            ``Context.referenceHierarchy`` in the span payload.
+        """
+        result: List[Dict[str, str]] = []
+        for e in self._entries:
+            item: Dict[str, str] = {
+                "serviceType": e.service_type,
+                "referenceId": e.reference_id,
+            }
+            if e.version:
+                item["version"] = e.version
+            result.append(item)
+        return result
+
+    @staticmethod
+    def from_baggage_header(header_value: Optional[str]) -> "ReferenceContext":
+        """Parse ``x-uipath-tracebaggage`` header value into a ReferenceContext.
+
+        Only entries that carry the ``ref.*`` shape (type + valid UUID id) are
+        included. Malformed or plain-KV entries are silently skipped so a bad
+        header from an upstream service cannot crash this one.
+
+        Args:
+            header_value: Raw header string, e.g.
+                ``"ref.type=agent;ref.id=<uuid>;ref.v=1.0,ref.type=maestro;ref.id=<uuid>"``
+
+        Returns:
+            Parsed :class:`ReferenceContext`, or :attr:`ReferenceContext.Empty`
+            if the header is absent, empty, or contains no valid ref entries.
+        """
+        if not header_value or not header_value.strip():
+            return ReferenceContext.Empty
+
+        entries: List[ReferenceEntry] = []
+        for raw_entry in header_value.split(","):
+            entry_text = raw_entry.strip()
+            if not entry_text:
+                continue
+            props: dict[str, str] = {}
+            for raw_pair in entry_text.split(";"):
+                pair_text = raw_pair.strip()
+                eq = pair_text.find("=")
+                if eq <= 0 or eq >= len(pair_text) - 1:
+                    continue
+                key = pair_text[:eq].strip()
+                value = pair_text[eq + 1:].strip()
+                if key and value:
+                    props[key] = value
+
+            type_v = props.get(BAGGAGE_KEY_TYPE)
+            id_v = props.get(BAGGAGE_KEY_ID)
+            if not type_v or not id_v:
+                continue
+            try:
+                parsed_uuid = uuid.UUID(id_v)
+            except (ValueError, AttributeError):
+                continue
+            entries.append(
+                ReferenceEntry(
+                    service_type=type_v,
+                    reference_id=str(parsed_uuid),
+                    version=props.get(BAGGAGE_KEY_VERSION) or None,
+                )
+            )
+
+        if not entries:
+            return ReferenceContext.Empty
+        return ReferenceContext(tuple(entries))
+
+    def to_baggage_header_value(self) -> str:
+        """Serialize to ``x-uipath-tracebaggage`` header value.
+
+        Returns:
+            Comma-separated entries; each is a semicolon-separated list of
+            ``key=value`` pairs. Empty context returns ``""``.
+        """
+        if not self._entries:
+            return ""
+        parts: List[str] = []
+        for e in self._entries:
+            kv = f"{BAGGAGE_KEY_TYPE}={e.service_type};{BAGGAGE_KEY_ID}={e.reference_id}"
+            if e.version:
+                kv += f";{BAGGAGE_KEY_VERSION}={e.version}"
+            parts.append(kv)
+        return ",".join(parts)
+
+
+# Assigned after class body so ReferenceContext is fully bound.
+ReferenceContext.Empty = ReferenceContext()
+
+
+class ReferenceContextAccessor:
+    """Ambient accessor for the current :class:`ReferenceContext`.
+
+    Backed by :mod:`contextvars` so the value propagates across ``await``
+    boundaries without being threaded through every call signature.
+
+    Usage::
+
+        token = ReferenceContextAccessor.set(ctx)
+        try:
+            ...  # code here sees ReferenceContextAccessor.get() == ctx
+        finally:
+            ReferenceContextAccessor.reset(token)
+    """
+
+    _current: contextvars.ContextVar[Optional[ReferenceContext]] = (
+        contextvars.ContextVar("uipath_reference_context", default=None)
+    )
+
+    @classmethod
+    def get(cls) -> Optional[ReferenceContext]:
+        """Return the current ambient context, or ``None`` if not set."""
+        return cls._current.get()
+
+    @classmethod
+    def set(cls, value: Optional[ReferenceContext]) -> contextvars.Token[Optional[ReferenceContext]]:
+        """Set the ambient context. Returns a token for restoration.
+
+        Pass the token to :meth:`reset` in a ``finally`` block.
+        """
+        return cls._current.set(value)
+
+    @classmethod
+    def reset(cls, token: contextvars.Token[Optional[ReferenceContext]]) -> None:
+        """Restore the ambient context to its prior value."""
+        cls._current.reset(token)

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

@@ -10,11 +10,11 @@
 from os import environ as env
 from typing import Any, Dict, List, Optional, TypeVar
 
-from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry import context as context_api
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.trace import StatusCode
 from pydantic import BaseModel, ConfigDict, Field
 from uipath.core.serialization import serialize_json
-
 from uipath.platform.constants import (
     ENV_FOLDER_KEY,
     ENV_JOB_KEY,
@@ -26,6 +26,30 @@
     ENV_UIPATH_TRACE_ID,
 )
 
+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))
+
+
+class ReferenceHierarchySpanProcessor(SpanProcessor):
+    """Stamps uipath.reference_hierarchy on every span at creation time.
+
+    Runs on_start in the span-creating thread so ContextVar values are live.
+    Register this before any processor that reads the attribute on on_start
+    (e.g. LiveTrackingSpanProcessor).
+    """
+
+    def on_start(
+        self, span: Span, parent_context: context_api.Context | None = None
+    ) -> None:
+        _inject_reference_hierarchy(span)
+
 logger = logging.getLogger(__name__)
 
 
@@ -245,6 +269,7 @@ class UiPathSpan:
     agent_version: Optional[str] = None
     verbosity_level: Optional[VerbosityLevel] = None
     attachments: Optional[List[SpanAttachment]] = None
+    context: Optional[Dict[str, Any]] = None
 
     def to_dict(self, serialize_attributes: bool = True) -> Dict[str, Any]:
         """Convert the Span to a dictionary suitable for JSON serialization.
@@ -304,6 +329,8 @@ def to_dict(self, serialize_attributes: bool = True) -> Dict[str, Any]:
                 del result[guid_key]
         if self.verbosity_level is not None:
             result["VerbosityLevel"] = self.verbosity_level
+        if self.context is not None:
+            result["Context"] = self.context
         return result
 
 
@@ -380,6 +407,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 = SpanStatus.OK
         if otel_span.status.status_code == StatusCode.ERROR:
@@ -496,6 +528,13 @@ def otel_span_to_uipath_span(
             except Exception as e:
                 logger.warning(f"Error processing attachments: {e}")
 
+        context: Optional[Dict[str, Any]] = None
+        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(
             (otel_span.start_time or 0) / 1e9
@@ -527,6 +566,7 @@ def otel_span_to_uipath_span(
             reference_id=reference_id,
             source=source,
             attachments=attachments,
+            context=context,
         )
 
     @staticmethod