feat(tracing): add current span ancestors provider

Author: ionmincu

src/uipath/core/tracing/manager.py

@@ -1,59 +1,319 @@
 """Tracing manager for handling tracer implementations and function registry."""
 
 import logging
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from opentelemetry import context, trace
+from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.trace import set_span_in_context
 
 logger = logging.getLogger(__name__)
 
 
+class SpanRegistry:
+    """Registry to track all spans and their parent relationships."""
+
+    MAX_HIERARCHY_DEPTH = 1000  # Hard limit for hierarchy traversal
+
+    def __init__(self):
+        self._spans: Dict[int, ReadableSpan] = {}  # span_id -> span
+        self._parent_map: Dict[int, Optional[int]] = {}  # span_id -> parent_id
+
+    def register_span(self, span: ReadableSpan) -> None:
+        """Register a span and its parent relationship."""
+        span_id = span.get_span_context().span_id
+        parent_id = span.parent.span_id if span.parent else None
+
+        self._spans[span_id] = span
+        self._parent_map[span_id] = parent_id
+
+        parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
+        logger.debug(
+            f"SpanRegistry: registered span: {span.name} (id: {span_id:016x}, parent: {parent_str})"
+        )
+
+    def get_span(self, span_id: int) -> Optional[ReadableSpan]:
+        """Get a span by ID."""
+        return self._spans.get(span_id)
+
+    def get_parent_id(self, span_id: int) -> Optional[int]:
+        """Get the parent ID of a span."""
+        return self._parent_map.get(span_id)
+
+    def calculate_depth(self, span_id: int) -> int:
+        """Calculate the depth of a span in the hierarchy.
+
+        Returns:
+            The depth of the span, capped at MAX_HIERARCHY_DEPTH.
+        """
+        depth = 0
+        current_id = span_id
+        visited = set()
+
+        while current_id is not None and current_id not in visited:
+            visited.add(current_id)
+            parent_id = self._parent_map.get(current_id)
+            if parent_id is None:
+                break
+            depth += 1
+            if depth >= self.MAX_HIERARCHY_DEPTH:
+                logger.warning(
+                    f"Hit MAX_HIERARCHY_DEPTH ({self.MAX_HIERARCHY_DEPTH}) while calculating depth for span {span_id:016x}"
+                )
+                break
+            current_id = parent_id
+
+        return depth
+
+    def is_ancestor(self, ancestor_id: int, descendant_id: int) -> bool:
+        """Check if ancestor_id is an ancestor of descendant_id.
+
+        Returns:
+            True if ancestor_id is an ancestor of descendant_id, False otherwise.
+            If MAX_HIERARCHY_DEPTH is reached, returns False.
+        """
+        current_id = descendant_id
+        visited = set()
+        steps = 0
+
+        while current_id is not None and current_id not in visited:
+            if current_id == ancestor_id:
+                return True
+            visited.add(current_id)
+            current_id = self._parent_map.get(current_id)
+            steps += 1
+            if steps >= self.MAX_HIERARCHY_DEPTH:
+                logger.warning(
+                    f"Hit MAX_HIERARCHY_DEPTH ({self.MAX_HIERARCHY_DEPTH}) while checking ancestry between {ancestor_id:016x} and {descendant_id:016x}"
+                )
+                return False
+
+        return False
+
+    def get_chain(self, span_id: int) -> List[int]:
+        """Get the complete ancestor chain for a span (bottom to top).
+
+        Returns:
+            List of span IDs from the given span to the root, capped at MAX_HIERARCHY_DEPTH.
+        """
+        chain = []
+        current_id = span_id
+        visited = set()
+
+        while current_id is not None and current_id not in visited:
+            chain.append(current_id)
+            visited.add(current_id)
+            if len(chain) >= self.MAX_HIERARCHY_DEPTH:
+                logger.warning(
+                    f"Hit MAX_HIERARCHY_DEPTH ({self.MAX_HIERARCHY_DEPTH}) while building chain for span {span_id:016x}"
+                )
+                break
+            current_id = self._parent_map.get(current_id)
+
+        return chain
+
+    def clear(self) -> None:
+        """Clear all registered spans."""
+        self._spans.clear()
+        self._parent_map.clear()
+
+
+# Global span registry instance
+_span_registry = SpanRegistry()
+
+
 class UiPathTracingManager:
     """Static utility class to manage tracing implementations and decorated functions."""
 
     _current_span_provider: Optional[Callable[[], Any]] = None
+    _current_span_ancestors_provider: Optional[Callable[[], List[Any]]] = None
 
-    @classmethod
+    @staticmethod
     def register_current_span_provider(
-        cls, current_span_provider: Optional[Callable[[], Any]]
+        current_span_provider: Optional[Callable[[], Any]],
     ):
         """Register a custom current span provider function.
 
         Args:
             current_span_provider: A function that returns the current span from an external
                                  tracing framework. If None, no custom span parenting will be used.
         """
-        cls._current_span_provider = current_span_provider
+        UiPathTracingManager._current_span_provider = current_span_provider
 
     @staticmethod
-    def get_parent_context():
+    def get_parent_context() -> context.Context:
         """Get the parent context for span creation.
 
-        Prioritizes:
-        1. Currently active OTel span (for recursion/children)
-        2. External span provider (for top-level calls)
-        3. Current context as fallback
+        This method determines the correct parent context when creating a new traced span.
+        It handles scenarios where spans may exist in both OpenTelemetry's context (current_span)
+        and in an external tracing system (external_span), such as LangGraph.
+
+        The algorithm follows this priority:
+
+        1. **No spans available**: Returns the current OpenTelemetry context (empty context)
+
+        2. **Only current_span exists**: Returns a context with current_span set as parent
+           - This is the standard OpenTelemetry behavior for nested traced functions
+
+        3. **Only external_span exists**: Returns a context with external_span set as parent
+           - This occurs when an external tracing system (like LangGraph) has an active span
+             but there's no OTel span in the current call stack
+
+        4. **Both spans exist**: Calls `_get_bottom_most_span()` to determine which is deeper
+           - Uses the SpanRegistry to build parent-child relationships
+           - Returns the span that is closer to the "bottom" (leaf) of the trace tree
+           - This ensures new spans are always attached to the deepest/most specific parent
+
+        Context Sources:
+        - **current_span**: Retrieved from OpenTelemetry's `trace.get_current_span()`
+          - Represents the active OTel span in the current execution context
+          - Created by `@traced` decorators or manual span creation
+
+        - **external_span**: Retrieved from the registered custom span provider
+          - Set via `register_current_span_provider()`
+          - Typically provided by external frameworks (LangGraph, LangChain, etc.)
+          - Allows integration with tracing systems outside of OpenTelemetry
+
+        Returns:
+            context.Context: An OpenTelemetry context containing the appropriate parent span,
+                           or the current empty context if no spans are available
+
+        Example:
+            ```python
+            # Called by the @traced decorator when creating a new span:
+            ctx = UiPathTracingManager.get_parent_context()
+            with tracer.start_as_current_span("my_span", context=ctx) as span:
+                # New span will have the correct parent based on the logic above
+                pass
+            ```
+
+        See Also:
+            - `_get_bottom_most_span()`: Logic for choosing between two available spans
+            - `register_current_span_provider()`: Register external span provider
+            - `get_external_current_span()`: Retrieve span from external provider
         """
-        # Always use the currently active OTel span if valid (recursion / children)
         current_span = trace.get_current_span()
+        has_current_span = (
+            current_span is not None and current_span.get_span_context().is_valid
+        )
+
+        external_span = UiPathTracingManager.get_external_current_span()
+        has_external_span = external_span is not None
 
-        if current_span is not None and current_span.get_span_context().is_valid:
+        # Only one or no spans available
+        if not has_current_span:
+            return (
+                set_span_in_context(external_span)
+                if has_external_span
+                else context.get_current()
+            )
+        if not has_external_span:
             return set_span_in_context(current_span)
 
-        # Only for the very top-level call, fallback to external span provider
+        # Both spans exist - find the bottom-most one
+        bottom_span = UiPathTracingManager._get_bottom_most_span(
+            current_span, external_span
+        )
+        return set_span_in_context(bottom_span)
+
+    @staticmethod
+    def _get_bottom_most_span(
+        current_span: ReadableSpan, external_span: ReadableSpan
+    ) -> ReadableSpan:
+        """Determine which span is deeper in the ancestor tree.
+
+        Args:
+            current_span: The OTel current span
+            external_span: The external span from the provider
+
+        Returns:
+            The span that is deeper (closer to the bottom) in the call hierarchy
+        """
+        # Register both spans in the registry
+        _span_registry.register_span(current_span)
+        _span_registry.register_span(external_span)
+
+        # Also register external ancestors
+        external_ancestors = UiPathTracingManager.get_ancestor_spans() or []
+        for ancestor in external_ancestors:
+            _span_registry.register_span(ancestor)
+
+        current_span_id = current_span.get_span_context().span_id
+        external_span_id = external_span.get_span_context().span_id
+
+        # Check if one span is an ancestor of the other
+        if _span_registry.is_ancestor(external_span_id, current_span_id):
+            logger.debug(
+                "Traced Context: current_span is a descendant of external_span -> returning current_span (deeper)"
+            )
+            return current_span
+        elif _span_registry.is_ancestor(current_span_id, external_span_id):
+            logger.debug(
+                "Traced Context: external_span is a descendant of current_span -> returning external_span (deeper)"
+            )
+            return external_span
+
+        # Neither is an ancestor of the other - they're in different branches
+        # Use depth as tiebreaker
+        current_depth = _span_registry.calculate_depth(current_span_id)
+        external_depth = _span_registry.calculate_depth(external_span_id)
+
+        if current_depth > external_depth:
+            logger.debug(
+                f"Traced Context: Different branches, current_span is deeper (depth {current_depth} > {external_depth}) -> returning current_span"
+            )
+            return current_span
+        elif external_depth > current_depth:
+            logger.debug(
+                f"Traced Context: Different branches, external_span is deeper (depth {external_depth} > {current_depth}) -> returning external_span"
+            )
+            return external_span
+        else:
+            # Same depth, different branches - default to external
+            logger.debug(
+                f"Traced Context: Same depth ({current_depth}), different branches -> defaulting to external_span"
+            )
+            return external_span
+
+    @staticmethod
+    def get_external_current_span():
+        """Get the current span from the external provider, if any."""
         if UiPathTracingManager._current_span_provider is not None:
             try:
-                external_span = UiPathTracingManager._current_span_provider()
-                if external_span is not None:
-                    return set_span_in_context(external_span)
+                return UiPathTracingManager._current_span_provider()
             except Exception as e:
                 logger.warning(f"Error getting current span from provider: {e}")
+        return None
 
-        # Last fallback
-        ctx = context.get_current()
+    @staticmethod
+    def get_ancestor_spans() -> List[Any]:
+        """Get the ancestor spans from the registered provider, if any."""
+        if UiPathTracingManager._current_span_ancestors_provider is not None:
+            try:
+                return UiPathTracingManager._current_span_ancestors_provider()
+            except Exception as e:
+                logger.warning(f"Error getting ancestor spans from provider: {e}")
+        return []
 
-        return ctx
+    @staticmethod
+    def register_current_span_ancestors_provider(
+        current_span_ancestors_provider: Optional[Callable[[], List[Any]]],
+    ):
+        """Register a custom current span ancestors provider function.
+
+        Args:
+            current_span_ancestors_provider: A function that returns a list of ancestor spans
+                                           from an external tracing framework. If None, no custom
+                                           span ancestor information will be used.
+        """
+        UiPathTracingManager._current_span_ancestors_provider = (
+            current_span_ancestors_provider
+        )
+
+    @staticmethod
+    def get_current_span_ancestors_provider():
+        """Get the currently set custom span ancestors provider."""
+        return UiPathTracingManager._current_span_ancestors_provider
 
 
 __all__ = ["UiPathTracingManager"]