From 132305194db08c880cb2d0f088389e150eb04f5f Mon Sep 17 00:00:00 2001 From: Jesus Date: Mon, 8 Jun 2026 21:12:48 -0700 Subject: [PATCH 1/3] feat: add ReferenceContext for span reference hierarchy propagation Introduces an immutable, copy-on-write ReferenceContext (modelled after service-common BaggageContext) that propagates the reference hierarchy through all spans emitted to LLMOps/Traceview. - `_reference_context.py`: ReferenceEntry, ReferenceContext (add / to_wire_list / from_baggage_header / to_baggage_header_value), ReferenceContextAccessor (ContextVar-backed, async-safe) - `_span_utils.py`: UiPathSpan gains `context` field; to_dict() emits `Context.referenceHierarchy` when set; otel_span_to_uipath_span() reads ReferenceContextAccessor and builds the wire payload - `uipath.tracing`: re-exports ReferenceEntry, ReferenceContext, ReferenceContextAccessor for downstream consumers - Tests: 30 new tests covering immutability, wire format, baggage header round-trip, accessor lifecycle, and span wiring Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../src/uipath/platform/common/__init__.py | 8 + .../platform/common/_reference_context.py | 254 +++++++++++++++ .../src/uipath/platform/common/_span_utils.py | 15 + .../tests/services/test_reference_context.py | 303 ++++++++++++++++++ .../uipath/src/uipath/tracing/__init__.py | 8 + 5 files changed, 588 insertions(+) create mode 100644 packages/uipath-platform/src/uipath/platform/common/_reference_context.py create mode 100644 packages/uipath-platform/tests/services/test_reference_context.py diff --git a/packages/uipath-platform/src/uipath/platform/common/__init__.py b/packages/uipath-platform/src/uipath/platform/common/__init__.py index cefd92075..24fc01645 100644 --- a/packages/uipath-platform/src/uipath/platform/common/__init__.py +++ b/packages/uipath-platform/src/uipath/platform/common/__init__.py @@ -22,6 +22,11 @@ from ._http_config import get_ca_bundle_path, get_httpx_client_kwargs from ._models import Endpoint, RequestSpec from ._service_url_overrides import inject_routing_headers, resolve_service_url +from ._reference_context import ( + ReferenceContext, + ReferenceContextAccessor, + ReferenceEntry, +) from ._span_utils import UiPathSpan, _SpanUtils from ._url import UiPathUrl from ._user_agent import user_agent_value @@ -108,6 +113,9 @@ "ResourceOverwrite", "ResourceOverwriteParser", "ResourceOverwritesContext", + "ReferenceEntry", + "ReferenceContext", + "ReferenceContextAccessor", "UiPathSpan", "_SpanUtils", "resolve_service_url", diff --git a/packages/uipath-platform/src/uipath/platform/common/_reference_context.py b/packages/uipath-platform/src/uipath/platform/common/_reference_context.py new file mode 100644 index 000000000..f8f92d58d --- /dev/null +++ b/packages/uipath-platform/src/uipath/platform/common/_reference_context.py @@ -0,0 +1,254 @@ +"""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 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) + """ + + __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): + id_str = reference_id + 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]: + """Serialize to the ``referenceHierarchy`` wire format. + + Returns: + A list of dicts suitable for JSON serialization as + ``Context.referenceHierarchy`` in the span payload. + """ + result = [] + for e in self._entries: + item: dict = { + "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=;ref.v=1.0,ref.type=maestro;ref.id="`` + + 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() # type: ignore[attr-defined] + + +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: + """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) -> None: + """Restore the ambient context to its prior value.""" + cls._current.reset(token) diff --git a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py b/packages/uipath-platform/src/uipath/platform/common/_span_utils.py index ab91b3623..b3f7d142c 100644 --- a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py +++ b/packages/uipath-platform/src/uipath/platform/common/_span_utils.py @@ -13,6 +13,8 @@ from pydantic import BaseModel, ConfigDict, Field from uipath.core.serialization import serialize_json +from ._reference_context import ReferenceContextAccessor + logger = logging.getLogger(__name__) # SourceEnum.CodedAgents = 10 (default for Python SDK / coded agents) @@ -99,6 +101,7 @@ class UiPathSpan: agent_version: Optional[str] = None verbosity_level: Optional[int] = 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. @@ -151,6 +154,8 @@ def to_dict(self, serialize_attributes: bool = True) -> Dict[str, Any]: } if self.verbosity_level is not None: result["VerbosityLevel"] = self.verbosity_level + if self.context is not None: + result["Context"] = self.context return result @@ -326,6 +331,15 @@ 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} + # Create UiPathSpan from OpenTelemetry span start_time = datetime.fromtimestamp( (otel_span.start_time or 0) / 1e9 @@ -357,6 +371,7 @@ def otel_span_to_uipath_span( reference_id=reference_id, source=source, attachments=attachments, + context=context, ) @staticmethod diff --git a/packages/uipath-platform/tests/services/test_reference_context.py b/packages/uipath-platform/tests/services/test_reference_context.py new file mode 100644 index 000000000..3379a784c --- /dev/null +++ b/packages/uipath-platform/tests/services/test_reference_context.py @@ -0,0 +1,303 @@ +"""Tests for ReferenceContext, ReferenceContextAccessor, and span Context wiring.""" + +import os +from datetime import datetime +from unittest.mock import Mock + +import pytest +from opentelemetry.sdk.trace import Span as OTelSpan +from opentelemetry.trace import SpanContext, StatusCode + +from uipath.platform.common import _SpanUtils +from uipath.platform.common._reference_context import ( + ReferenceContext, + ReferenceContextAccessor, + ReferenceEntry, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_mock_span(attributes: dict | None = None) -> Mock: + mock = Mock(spec=OTelSpan) + mock.get_span_context.return_value = SpanContext( + trace_id=0x123456789ABCDEF0123456789ABCDEF0, + span_id=0x0123456789ABCDEF, + is_remote=False, + ) + mock.name = "test-span" + mock.parent = None + mock.status.status_code = StatusCode.OK + mock.attributes = attributes or {} + mock.events = [] + mock.links = [] + now_ns = int(datetime.now().timestamp() * 1e9) + mock.start_time = now_ns + mock.end_time = now_ns + 1_000_000 + return mock + + +# --------------------------------------------------------------------------- +# ReferenceContext — immutability & copy-on-write +# --------------------------------------------------------------------------- + +class TestReferenceContextImmutability: + def test_empty_singleton_is_falsy(self) -> None: + assert not ReferenceContext.Empty + + def test_add_returns_new_instance(self) -> None: + base = ReferenceContext.Empty + child = base.add("agent", "550e8400-e29b-41d4-a716-446655440001") + assert child is not base + + def test_original_unmodified_after_add(self) -> None: + base = ReferenceContext.Empty + base.add("agent", "550e8400-e29b-41d4-a716-446655440001") + assert len(base) == 0 + + def test_siblings_do_not_share_entries(self) -> None: + base = ReferenceContext.Empty.add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.0") + child_a = base.add("agent", "550e8400-e29b-41d4-a716-446655440011") + child_b = base.add("agent", "550e8400-e29b-41d4-a716-446655440012") + + assert len(base) == 1 + assert len(child_a) == 2 + assert len(child_b) == 2 + assert child_a.entries[1].reference_id != child_b.entries[1].reference_id + + def test_equality_and_hash(self) -> None: + a = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + b = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + assert a == b + assert hash(a) == hash(b) + + def test_different_entries_not_equal(self) -> None: + a = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + b = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440002") + assert a != b + + +# --------------------------------------------------------------------------- +# ReferenceContext.add — validation & UUID coercion +# --------------------------------------------------------------------------- + +class TestReferenceContextAdd: + def test_add_with_string_uuid(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001", "1.0") + assert len(ctx) == 1 + e = ctx.entries[0] + assert e.service_type == "agent" + assert e.reference_id == "550e8400-e29b-41d4-a716-446655440001" + assert e.version == "1.0" + + def test_add_with_uuid_object(self) -> None: + import uuid + uid = uuid.UUID("550e8400-e29b-41d4-a716-446655440001") + ctx = ReferenceContext.Empty.add("agent", uid) + assert ctx.entries[0].reference_id == str(uid) + + def test_add_without_version(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + assert ctx.entries[0].version is None + + def test_add_blank_version_normalised_to_none(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001", " ") + assert ctx.entries[0].version is None + + def test_add_empty_service_type_raises(self) -> None: + with pytest.raises(ValueError, match="service_type"): + ReferenceContext.Empty.add("", "550e8400-e29b-41d4-a716-446655440001") + + def test_add_invalid_reference_id_type_raises(self) -> None: + with pytest.raises(TypeError, match="reference_id"): + ReferenceContext.Empty.add("agent", 12345) # type: ignore[arg-type] + + def test_entries_ordered_oldest_first(self) -> None: + ctx = ( + ReferenceContext.Empty + .add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.0") + .add("agent", "550e8400-e29b-41d4-a716-446655440011") + ) + assert ctx.entries[0].service_type == "maestro" + assert ctx.entries[1].service_type == "agent" + + +# --------------------------------------------------------------------------- +# ReferenceContext.to_wire_list +# --------------------------------------------------------------------------- + +class TestToWireList: + def test_empty_produces_empty_list(self) -> None: + assert ReferenceContext.Empty.to_wire_list() == [] + + def test_single_entry_with_version(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001", "1.0.0") + wire = ctx.to_wire_list() + assert wire == [ + {"serviceType": "agent", "referenceId": "550e8400-e29b-41d4-a716-446655440001", "version": "1.0.0"} + ] + + def test_single_entry_without_version_omits_key(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + wire = ctx.to_wire_list() + assert "version" not in wire[0] + + def test_multiple_entries_order_preserved(self) -> None: + ctx = ( + ReferenceContext.Empty + .add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.1.0") + .add("agent", "550e8400-e29b-41d4-a716-446655440011") + ) + wire = ctx.to_wire_list() + assert len(wire) == 2 + assert wire[0]["serviceType"] == "maestro" + assert wire[0]["version"] == "2.1.0" + assert wire[1]["serviceType"] == "agent" + assert "version" not in wire[1] + + +# --------------------------------------------------------------------------- +# ReferenceContext.from_baggage_header / to_baggage_header_value round-trip +# --------------------------------------------------------------------------- + +class TestBaggageHeaderRoundTrip: + def test_round_trip_single_entry_with_version(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001", "1.0") + assert ReferenceContext.from_baggage_header(ctx.to_baggage_header_value()) == ctx + + def test_round_trip_multiple_entries(self) -> None: + ctx = ( + ReferenceContext.Empty + .add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.0") + .add("agent", "550e8400-e29b-41d4-a716-446655440011") + ) + assert ReferenceContext.from_baggage_header(ctx.to_baggage_header_value()) == ctx + + def test_empty_header_returns_empty(self) -> None: + assert ReferenceContext.from_baggage_header("") == ReferenceContext.Empty + assert ReferenceContext.from_baggage_header(None) == ReferenceContext.Empty + + def test_malformed_entry_skipped_silently(self) -> None: + # Only the second entry is valid + header = "not-a-valid-entry,ref.type=agent;ref.id=550e8400-e29b-41d4-a716-446655440001" + ctx = ReferenceContext.from_baggage_header(header) + assert len(ctx) == 1 + assert ctx.entries[0].service_type == "agent" + + def test_entry_with_invalid_uuid_skipped(self) -> None: + header = "ref.type=agent;ref.id=not-a-uuid" + ctx = ReferenceContext.from_baggage_header(header) + assert ctx == ReferenceContext.Empty + + def test_entry_missing_type_skipped(self) -> None: + header = "ref.id=550e8400-e29b-41d4-a716-446655440001" + ctx = ReferenceContext.from_baggage_header(header) + assert ctx == ReferenceContext.Empty + + def test_empty_context_produces_empty_header(self) -> None: + assert ReferenceContext.Empty.to_baggage_header_value() == "" + + +# --------------------------------------------------------------------------- +# ReferenceContextAccessor — ContextVar semantics +# --------------------------------------------------------------------------- + +class TestReferenceContextAccessor: + def setup_method(self) -> None: + # Ensure clean state before each test + current = ReferenceContextAccessor.get() + if current is not None: + token = ReferenceContextAccessor.set(None) + # immediately reset to avoid polluting other tests + ReferenceContextAccessor.reset(token) + + def test_default_is_none(self) -> None: + assert ReferenceContextAccessor.get() is None + + def test_set_and_get(self) -> None: + ctx = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + token = ReferenceContextAccessor.set(ctx) + try: + assert ReferenceContextAccessor.get() == ctx + finally: + ReferenceContextAccessor.reset(token) + + def test_reset_restores_prior_value(self) -> None: + ctx_a = ReferenceContext.Empty.add("agent", "550e8400-e29b-41d4-a716-446655440001") + ctx_b = ctx_a.add("langgraph", "550e8400-e29b-41d4-a716-446655440002") + + token_a = ReferenceContextAccessor.set(ctx_a) + token_b = ReferenceContextAccessor.set(ctx_b) + + assert ReferenceContextAccessor.get() == ctx_b + ReferenceContextAccessor.reset(token_b) + assert ReferenceContextAccessor.get() == ctx_a + ReferenceContextAccessor.reset(token_a) + assert ReferenceContextAccessor.get() is None + + +# --------------------------------------------------------------------------- +# UiPathSpan.context wiring via otel_span_to_uipath_span +# --------------------------------------------------------------------------- + +class TestContextWiring: + 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) + + def test_context_present_when_reference_context_set( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("UIPATH_ORGANIZATION_ID", "test-org") + 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) + + def test_context_carries_full_hierarchy( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("UIPATH_ORGANIZATION_ID", "test-org") + ref_ctx = ( + ReferenceContext.Empty + .add("maestro", "550e8400-e29b-41d4-a716-446655440010", "2.0") + .add("agent", "550e8400-e29b-41d4-a716-446655440011") + ) + 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] + finally: + ReferenceContextAccessor.reset(token) diff --git a/packages/uipath/src/uipath/tracing/__init__.py b/packages/uipath/src/uipath/tracing/__init__.py index e6c37bc99..1b0541a55 100644 --- a/packages/uipath/src/uipath/tracing/__init__.py +++ b/packages/uipath/src/uipath/tracing/__init__.py @@ -1,6 +1,11 @@ """Tracing utilities and OpenTelemetry exporters.""" from uipath.core import traced +from uipath.platform.common._reference_context import ( + ReferenceContext, + ReferenceContextAccessor, + ReferenceEntry, +) from uipath.platform.common._span_utils import ( AttachmentDirection, AttachmentProvider, @@ -25,4 +30,7 @@ "AttachmentProvider", "SpanAttachment", "VerbosityLevel", + "ReferenceEntry", + "ReferenceContext", + "ReferenceContextAccessor", ] From 84045a2258f785b1264f315036121dadc4975ad7 Mon Sep 17 00:00:00 2001 From: Jesus Date: Mon, 8 Jun 2026 21:46:51 -0700 Subject: [PATCH 2/3] fix: normalize empty-string env vars to None for Guid fields in UiPathSpan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit UIPATH_FOLDER_KEY, UIPATH_ORGANIZATION_ID, and UIPATH_TENANT_ID default to "" when unset. SpanReq.FolderKey is Guid? on the server — sending "" causes a JSON deserialization exception ("could not convert to System.Nullable[Guid]") which cascades into "The spans field is required" from ASP.NET's model binding. Normalize all three to None so unset vars serialize as null instead of an unparseable empty string. Co-Authored-By: Claude Sonnet 4.6 (1M context) --- .../src/uipath/platform/common/_span_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py b/packages/uipath-platform/src/uipath/platform/common/_span_utils.py index b3f7d142c..1b5f8c7ee 100644 --- a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py +++ b/packages/uipath-platform/src/uipath/platform/common/_span_utils.py @@ -76,14 +76,14 @@ class UiPathSpan: created_at: str = field(default_factory=lambda: datetime.now().isoformat() + "Z") updated_at: str = field(default_factory=lambda: datetime.now().isoformat() + "Z") organization_id: Optional[str] = field( - default_factory=lambda: env.get("UIPATH_ORGANIZATION_ID", "") + default_factory=lambda: env.get("UIPATH_ORGANIZATION_ID") or None ) tenant_id: Optional[str] = field( - default_factory=lambda: env.get("UIPATH_TENANT_ID", "") + default_factory=lambda: env.get("UIPATH_TENANT_ID") or None ) expiry_time_utc: Optional[str] = None folder_key: Optional[str] = field( - default_factory=lambda: env.get("UIPATH_FOLDER_KEY", "") + default_factory=lambda: env.get("UIPATH_FOLDER_KEY") or None ) source: int = DEFAULT_SOURCE span_type: str = "Coded Agents" From ce989778f8500365081cd3243280730af8674209 Mon Sep 17 00:00:00 2001 From: Jesus Date: Tue, 9 Jun 2026 14:37:30 -0700 Subject: [PATCH 3/3] fix: stamp reference hierarchy as span attribute to survive BatchSpanProcessor thread boundary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../src/uipath/core/tracing/processors.py | 21 +++- .../src/uipath/platform/common/_span_utils.py | 31 +++-- .../tests/services/test_reference_context.py | 111 +++++++++++++----- 3 files changed, 122 insertions(+), 41 deletions(-) diff --git a/packages/uipath-core/src/uipath/core/tracing/processors.py b/packages/uipath-core/src/uipath/core/tracing/processors.py index 2fd10f639..9a1343c29 100644 --- a/packages/uipath-core/src/uipath/core/tracing/processors.py +++ b/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", ] diff --git a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py b/packages/uipath-platform/src/uipath/platform/common/_span_utils.py index 1b5f8c7ee..bf32e9d59 100644 --- a/packages/uipath-platform/src/uipath/platform/common/_span_utils.py +++ b/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( diff --git a/packages/uipath-platform/tests/services/test_reference_context.py b/packages/uipath-platform/tests/services/test_reference_context.py index 3379a784c..a2e7c1a52 100644 --- a/packages/uipath-platform/tests/services/test_reference_context.py +++ b/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)