feat(trace): add span filtering

Author: ionmincu

src/uipath/tracing/_otel_exporters.py

@@ -2,7 +2,7 @@
 import logging
 import os
 import time
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Any, Callable, Dict, List, Optional, Sequence
 
 import httpx
 from opentelemetry.sdk.trace import ReadableSpan
@@ -97,9 +97,17 @@ class Status:
     def __init__(
         self,
         trace_id: Optional[str] = None,
+        span_filter: Optional[Callable[[Dict[str, Any]], bool]] = None,
         **kwargs,
     ):
-        """Initialize the exporter with the base URL and authentication token."""
+        """Initialize the exporter with the base URL and authentication token.
+
+        Args:
+            trace_id: Optional custom trace ID to use for all spans
+            span_filter: Optional filter function that takes a span dict and returns True
+                        if the span should be filtered out (dropped). Children of filtered
+                        spans will be reparented to the filtered span's parent.
+        """
         super().__init__(**kwargs)
         self.base_url = self._get_base_url()
         self.auth_token = os.environ.get("UIPATH_ACCESS_TOKEN")
@@ -112,6 +120,10 @@ def __init__(
 
         self.http_client = httpx.Client(**client_kwargs, headers=self.headers)
         self.trace_id = trace_id
+        self.span_filter = span_filter
+
+        # Track filtered span IDs across batches: filtered_id -> new_parent_id
+        self._reparent_mapping: Dict[str, str] = {}
 
     def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
         """Export spans to UiPath LLM Ops."""
@@ -132,6 +144,15 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
             for span in spans
         ]
 
+        # Apply filtering and reparenting if filter is configured
+        filter_enabled = os.environ.get("UIPATH_FILTER_PARENT_SPAN")
+        if filter_enabled:
+            span_list = self._filter_and_reparent_spans(span_list)
+
+        if len(span_list) == 0:
+            logger.debug("No spans to export after filtering")
+            return SpanExportResult.SUCCESS
+
         url = self._build_url(span_list)
 
         # Process spans in-place - work directly with dict
@@ -149,6 +170,138 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
 
         return self._send_with_retries(url, span_list)
 
+    def _filter_and_reparent_spans(
+        self, span_list: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """Filter out spans and reparent their children.
+
+        Rules:
+        1. Root spans (uipath.is_root=True) are DROPPED, children reparented to UIPATH_PARENT_SPAN_ID
+        2. Spans matching span_filter are DROPPED, children reparented to filtered span's parent
+
+        Args:
+            span_list: List of span dicts to filter
+
+        Returns:
+            Filtered list of spans with updated ParentIds
+        """
+        new_parent_id = os.environ.get("UIPATH_PARENT_SPAN_ID")
+        if not new_parent_id:
+            logger.info("[Filter] UIPATH_PARENT_SPAN_ID not set, skipping filtering")
+            return span_list
+
+        logger.info(
+            f"[Filter] Starting filter with {len(span_list)} spans, "
+            f"UIPATH_PARENT_SPAN_ID={new_parent_id}, "
+            f"span_filter={'set' if self.span_filter else 'not set'}"
+        )
+
+        # First pass: identify spans to filter and build reparent mapping
+        logger.info("[Filter] === FIRST PASS: Identifying spans to filter ===")
+        for span in span_list:
+            span_id = span.get("Id")
+            span_name = span.get("Name")
+            span_parent_id = span.get("ParentId")
+            attributes = span.get("Attributes", {})
+
+            logger.info(
+                f"[Filter] Checking span: Id={span_id}, Name={span_name}, "
+                f"ParentId={span_parent_id}, attributes_type={type(attributes).__name__}"
+            )
+
+            if not isinstance(attributes, dict):
+                logger.info("[Filter]   -> Skipping (attributes not a dict)")
+                continue
+
+            is_root = attributes.get("uipath.is_root", False)
+            original_parent_id = attributes.get("uipath.original_parent_id")
+
+            logger.info(
+                f"[Filter]   -> is_root={is_root}, original_parent_id={original_parent_id}"
+            )
+
+            # Rule 1: Root spans are dropped, children go to UIPATH_PARENT_SPAN_ID
+            if is_root:
+                self._reparent_mapping[span_id] = new_parent_id
+                logger.info(
+                    f"[Filter] Root span marked for filtering: "
+                    f"Id={span_id}, Name={span.get('Name')}, "
+                    f"children will be reparented to {new_parent_id}"
+                )
+                continue
+
+            # Rule 2: Check custom filter function
+            if not self.span_filter:
+                logger.info("[Filter]   -> KEEP (no custom filter set)")
+                continue
+
+            filter_result = self.span_filter(span)
+            logger.info(f"[Filter]   -> Custom filter result: {filter_result}")
+
+            if not filter_result:
+                logger.info("[Filter]   -> KEEP (custom filter returned False)")
+                continue
+
+            # Filtered span's children go to this span's parent
+            # Use original_parent_id if available, otherwise use current ParentId
+            parent = original_parent_id or span.get("ParentId")
+            if parent:
+                # Check if parent itself was filtered (transitive reparenting)
+                while parent in self._reparent_mapping:
+                    parent = self._reparent_mapping[parent]
+                self._reparent_mapping[span_id] = parent
+            else:
+                self._reparent_mapping[span_id] = new_parent_id
+            logger.info(
+                f"[Filter]   -> WILL FILTER (custom filter matched), "
+                f"children will be reparented to {self._reparent_mapping[span_id]}"
+            )
+
+        logger.info(
+            f"[Filter] After first pass, reparent_mapping has {len(self._reparent_mapping)} entries: "
+            f"{self._reparent_mapping}"
+        )
+
+        # Second pass: filter spans and reparent children
+        logger.info("[Filter] === SECOND PASS: Filtering and reparenting ===")
+        filtered_spans = []
+        for span in span_list:
+            span_id = span.get("Id")
+            span_name = span.get("Name")
+            parent_id = span.get("ParentId")
+
+            # Skip filtered spans
+            if span_id in self._reparent_mapping:
+                logger.info(
+                    f"[Filter] DROPPING span: Id={span_id}, Name={span_name}"
+                )
+                continue
+
+            # Reparent if parent was filtered
+            if parent_id and parent_id in self._reparent_mapping:
+                old_parent = parent_id
+                # Follow the chain for transitive reparenting
+                while parent_id in self._reparent_mapping:
+                    parent_id = self._reparent_mapping[parent_id]
+                span["ParentId"] = parent_id
+                logger.info(
+                    f"[Filter] REPARENTING span: Id={span_id}, Name={span_name}, "
+                    f"ParentId: {old_parent} -> {parent_id}"
+                )
+            else:
+                logger.info(
+                    f"[Filter] KEEPING span unchanged: Id={span_id}, Name={span_name}, "
+                    f"ParentId={parent_id}"
+                )
+
+            filtered_spans.append(span)
+
+        logger.info(
+            f"[Filter] Complete: {len(span_list)} input -> {len(filtered_spans)} output spans, "
+            f"mapping size: {len(self._reparent_mapping)}"
+        )
+        return filtered_spans
+
     def force_flush(self, timeout_millis: int = 30000) -> bool:
         """Force flush the exporter."""
         return True

src/uipath/tracing/_utils.py

@@ -212,8 +212,13 @@ def otel_span_to_uipath_span(
 
         # Get parent span ID if it exists
         parent_id = None
+        is_root = otel_span.parent is None
+        original_parent_id: Optional[str] = None
+
         if otel_span.parent is not None:
             parent_id = _SpanUtils.span_id_to_uuid4(otel_span.parent.span_id)
+            # Store original parent ID for potential reparenting later
+            original_parent_id = str(parent_id)
         else:
             # Only set UIPATH_PARENT_SPAN_ID for root spans (spans without a parent)
             parent_span_id_str = env.get("UIPATH_PARENT_SPAN_ID")
@@ -226,6 +231,12 @@ 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 {}
 
+        # Add markers for filtering/reparenting in the exporter
+        if is_root:
+            attributes_dict["uipath.is_root"] = True
+        if original_parent_id:
+            attributes_dict["uipath.original_parent_id"] = original_parent_id
+
         # Map status
         status = 1  # Default to OK
         if otel_span.status.status_code == StatusCode.ERROR: