feat(tracing): add a run_id-keyed span provider to UiPathSpanUtils [PLT-112227]

packages/uipath-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-core"
-version = "0.5.33"
+version = "0.5.34"
 description = "UiPath Core abstractions"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

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

@@ -2,6 +2,7 @@
 
 import logging
 from typing import Callable, Optional
+from uuid import UUID
 
 from opentelemetry import context, trace
 from opentelemetry.trace import NonRecordingSpan, Span, set_span_in_context
@@ -139,6 +140,7 @@ class UiPathSpanUtils:
 
     _current_span_provider: Optional[Callable[[], Optional[Span]]] = None
     _current_span_ancestors_provider: Optional[Callable[[], list[Span]]] = None
+    _run_span_provider: Optional[Callable[[UUID], Optional[Span]]] = None
 
     @staticmethod
     def register_current_span_provider(
@@ -330,5 +332,32 @@ def get_current_span_ancestors_provider():
         """Get the currently set custom span ancestors provider."""
         return UiPathSpanUtils._current_span_ancestors_provider
 
+    @staticmethod
+    def register_run_span_provider(
+        run_span_provider: Optional[Callable[[UUID], Optional[Span]]],
+    ):
+        """Register a provider for the span an external framework opened for a run.
+
+        Unlike the current span provider, which answers for whatever is executing now,
+        this answers for a specific run, so a caller holding a framework run id (e.g. a
+        LangChain callback's ``run_id``) can parent its own spans to that run's span.
+
+        Args:
+            run_span_provider: A function that takes a framework run id and returns the
+                               span opened for that run, or None if it has none. If None,
+                               run span lookups return None.
+        """
+        UiPathSpanUtils._run_span_provider = run_span_provider
+
+    @staticmethod
+    def get_run_span(run_id: UUID) -> Optional[Span]:
+        """Get the span the registered provider opened for ``run_id``, if any."""
+        if UiPathSpanUtils._run_span_provider is not None:
+            try:
+                return UiPathSpanUtils._run_span_provider(run_id)
+            except Exception as e:
+                logger.warning("Error getting run span from provider: %s", e)
+        return None
+
 
 __all__ = ["ParentedNonRecordingSpan", "UiPathSpanUtils"]

packages/uipath-core/tests/tracing/test_external_integration.py

@@ -87,3 +87,49 @@ def test_function():
 
     spans = span_capture.get_spans()
     assert len(spans) == 1
+
+
+def test_run_span_provider_returns_the_span_for_the_run():
+    """Test that get_run_span resolves through the registered provider."""
+    from uuid import uuid4
+
+    from uipath.core.tracing.span_utils import UiPathSpanUtils
+
+    run_id, other_run_id = uuid4(), uuid4()
+    span = trace.get_tracer("external").start_span("llm_call")
+    UiPathSpanUtils.register_run_span_provider({run_id: span}.get)
+
+    assert UiPathSpanUtils.get_run_span(run_id) is span
+    assert UiPathSpanUtils.get_run_span(other_run_id) is None
+
+    # Clean up
+    UiPathSpanUtils.register_run_span_provider(None)
+    span.end()
+
+
+def test_run_span_without_provider_returns_none():
+    """Test that get_run_span returns None when no provider is registered."""
+    from uuid import uuid4
+
+    from uipath.core.tracing.span_utils import UiPathSpanUtils
+
+    UiPathSpanUtils.register_run_span_provider(None)
+
+    assert UiPathSpanUtils.get_run_span(uuid4()) is None
+
+
+def test_run_span_provider_raises_exception():
+    """Test that exceptions from the run span provider are caught."""
+    from uuid import UUID, uuid4
+
+    from uipath.core.tracing.span_utils import UiPathSpanUtils
+
+    def failing_provider(run_id: UUID):
+        raise RuntimeError("Provider failed!")
+
+    UiPathSpanUtils.register_run_span_provider(failing_provider)
+
+    assert UiPathSpanUtils.get_run_span(uuid4()) is None
+
+    # Clean up
+    UiPathSpanUtils.register_run_span_provider(None)