final

Author: ionmincu

src/uipath/core/tracing/manager.py

@@ -13,6 +13,8 @@
 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
@@ -26,7 +28,9 @@ def register_span(self, span: ReadableSpan) -> None:
         self._parent_map[span_id] = parent_id
 
         parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
-        logger.info(f"Registered span: {span.name} (id: {span_id:016x}, parent: {parent_str})")
+        logger.info(
+            f"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."""
@@ -37,7 +41,11 @@ def get_parent_id(self, span_id: int) -> Optional[int]:
         return self._parent_map.get(span_id)
 
     def calculate_depth(self, span_id: int) -> int:
-        """Calculate the depth of a span in the hierarchy."""
+        """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()
@@ -48,32 +56,58 @@ def calculate_depth(self, span_id: int) -> int:
             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."""
+        """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)."""
+        """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
@@ -110,10 +144,53 @@ def register_current_span_provider(
     def get_parent_context() -> context.Context:
         """Get the parent context for span creation.
 
-        Prioritizes:
-        - Bottom-most span when both current_span and external_span exist
-        - Single available span (current_span or external_span)
-        - 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
         """
         current_span = trace.get_current_span()
         has_current_span = (
@@ -152,9 +229,6 @@ def _get_bottom_most_span(
         Returns:
             The span that is deeper (closer to the bottom) in the call hierarchy
         """
-        logger.info("=" * 80)
-        logger.info("Determining bottom-most span...")
-
         # Register both spans in the registry
         _span_registry.register_span(current_span)
         _span_registry.register_span(external_span)
@@ -167,38 +241,16 @@ def _get_bottom_most_span(
         current_span_id = current_span.get_span_context().span_id
         external_span_id = external_span.get_span_context().span_id
 
-        # Get chains from registry
-        current_chain = _span_registry.get_chain(current_span_id)
-        external_chain = _span_registry.get_chain(external_span_id)
-
-        logger.info(f"current_span: '{current_span.name}' depth: {_span_registry.calculate_depth(current_span_id)}")
-        logger.info(f"external_span: '{external_span.name}' depth: {_span_registry.calculate_depth(external_span_id)}")
-
-        # Log the chains
-        logger.info("current_span chain (bottom to top):")
-        for i, span_id in enumerate(current_chain):
-            span = _span_registry.get_span(span_id)
-            if span:
-                parent_id = _span_registry.get_parent_id(span_id)
-                parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
-                logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_str})")
-
-        logger.info("external_span chain (bottom to top):")
-        for i, span_id in enumerate(external_chain):
-            span = _span_registry.get_span(span_id)
-            if span:
-                parent_id = _span_registry.get_parent_id(span_id)
-                parent_str = f"{parent_id:016x}" if parent_id is not None else "None"
-                logger.info(f"  [{i}] {span.name} (id: {span_id:016x}, parent: {parent_str})")
-
         # Check if one span is an ancestor of the other
         if _span_registry.is_ancestor(external_span_id, current_span_id):
-            logger.info("RESULT: current_span is a descendant of external_span -> returning current_span (deeper)")
-            logger.info("=" * 80)
+            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.info("RESULT: external_span is a descendant of current_span -> returning external_span (deeper)")
-            logger.info("=" * 80)
+            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
@@ -207,17 +259,20 @@ def _get_bottom_most_span(
         external_depth = _span_registry.calculate_depth(external_span_id)
 
         if current_depth > external_depth:
-            logger.info(f"RESULT: Different branches, current_span is deeper (depth {current_depth} > {external_depth}) -> returning current_span")
-            logger.info("=" * 80)
+            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.info(f"RESULT: Different branches, external_span is deeper (depth {external_depth} > {current_depth}) -> returning external_span")
-            logger.info("=" * 80)
+            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.info(f"RESULT: Same depth ({current_depth}), different branches -> defaulting to external_span")
-            logger.info("=" * 80)
+            logger.debug(
+                f"Traced Context: Same depth ({current_depth}), different branches -> defaulting to external_span"
+            )
             return external_span
 
     @staticmethod