fixup

Author: cristipufu

src/uipath/_cli/_dev/_terminal/__init__.py

@@ -7,8 +7,6 @@
 from uuid import uuid4
 
 from dotenv import load_dotenv
-from opentelemetry import trace
-from opentelemetry.sdk.trace import Tracer, TracerProvider
 from textual.app import App, ComposeResult
 from textual.binding import Binding
 from textual.containers import Container, Horizontal
@@ -17,7 +15,6 @@
 from ..._runtime._contracts import (
     UiPathRuntimeContext,
     UiPathRuntimeFactory,
-    UiPathTraceContext,
 )
 from ..._utils._console import ConsoleLogger
 from ._components._details import RunDetailsPanel
@@ -27,7 +24,6 @@
 from ._models._messages import LogMessage, TraceMessage
 from ._traces._exporter import RunContextExporter
 from ._traces._logger import RunContextLogHandler
-from ._traces._processor import RunContextProcessor
 
 console = ConsoleLogger()
 load_dotenv(override=True)
@@ -57,10 +53,12 @@ def __init__(
         self.initial_input: str = '{\n  "message": "Hello World"\n}'
         self.runs: Dict[str, ExecutionRun] = {}
         self.runtime_factory = runtime_factory
-        self.trace_provider: TracerProvider = TracerProvider()
-        self.trace_processor: RunContextProcessor = RunContextProcessor()
-        self.trace_provider.add_span_processor(self.trace_processor)
-        trace.set_tracer_provider(self.trace_provider)
+        self.runtime_factory.add_span_exporter(
+            RunContextExporter(
+                on_trace=self._handle_trace_message,
+                on_log=self._handle_log_message,
+            )
+        )
 
     def compose(self) -> ComposeResult:
         with Horizontal():
@@ -145,10 +143,7 @@ async def _execute_runtime(self, run: ExecutionRun):
                 entrypoint=run.entrypoint,
                 input=run.input_data,
                 trace_id=str(uuid4()),
-                tracing_enabled=False,
-                trace_context=UiPathTraceContext(
-                    enabled=False,
-                ),
+                execution_id=run.id,
                 logs_min_level=env.get("LOG_LEVEL", "INFO"),
                 log_handler=RunContextLogHandler(
                     run_id=run.id, on_log=self._handle_log_message
@@ -157,21 +152,11 @@ async def _execute_runtime(self, run: ExecutionRun):
 
             self._add_info_log(run, f"Starting execution: {run.entrypoint}")
 
-            trace_exporter = RunContextExporter(
-                run_id=run.id,
-                on_trace=self._handle_trace_message,
-                on_log=self._handle_log_message,
-            )
-            self.trace_processor.register_exporter(run.id, trace_exporter)
-            tracer: Tracer = trace.get_tracer("uipath-dev-terminal")
-
-            with tracer.start_as_current_span("root", attributes={"run.id": run.id}):
-                result = await self.runtime_factory.execute(context)
-                run.output_data = result.output
-                if run.output_data:
-                    self._add_info_log(run, f"Execution result: {run.output_data}")
+            result = await self.runtime_factory.execute_in_root_span(context)
 
-            self.trace_processor.unregister_exporter(run.id)
+            run.output_data = result.output
+            if run.output_data:
+                self._add_info_log(run, f"Execution result: {run.output_data}")
 
             self._add_info_log(run, "✅ Execution completed successfully")
             run.status = "completed"

src/uipath/_cli/_dev/_terminal/_traces/_exporter.py

@@ -13,11 +13,9 @@ class RunContextExporter(SpanExporter):
 
     def __init__(
         self,
-        run_id: str,
         on_trace: Callable[[TraceMessage], None],
         on_log: Callable[[LogMessage], None],
     ):
-        self.run_id = run_id
         self.on_trace = on_trace
         self.on_log = on_log
         self.logger = logging.getLogger(__name__)
@@ -54,14 +52,15 @@ def _export_span(self, span: Span):
         span_id = f"{span_context.span_id:016x}"  # 16-char hex string
         trace_id = f"{span_context.trace_id:032x}"  # 32-char hex string
 
+        run_id = span.attributes.get("execution.id")
         # Get parent span ID if available
         parent_span_id = None
         if hasattr(span, "parent") and span.parent:
             parent_span_id = f"{span.parent.span_id:016x}"
 
         # Create trace message with all required fields
         trace_msg = TraceMessage(
-            run_id=self.run_id,
+            run_id=run_id,
             span_name=span.name,
             span_id=span_id,
             parent_span_id=parent_span_id,
@@ -80,7 +79,7 @@ def _export_span(self, span: Span):
             for event in span.events:
                 log_level = self._determine_log_level(event, span.status)
                 log_msg = LogMessage(
-                    run_id=self.run_id,
+                    run_id=run_id,
                     level=log_level,
                     message=event.name,
                     timestamp=datetime.fromtimestamp(event.timestamp / 1_000_000_000),

src/uipath/_cli/_dev/_terminal/_traces/_processor.py

@@ -8,8 +8,13 @@
 from abc import ABC, abstractmethod
 from enum import Enum
 from functools import cached_property
-from typing import Any, Dict, Generic, Optional, Type, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
 
+from opentelemetry import context as context_api
+from opentelemetry import trace
+from opentelemetry.sdk.trace import Span, SpanProcessor, TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+from opentelemetry.trace import Tracer
 from pydantic import BaseModel, Field
 
 from ._logging import LogsInterceptor
@@ -136,6 +141,8 @@ class UiPathTraceContext(BaseModel):
     enabled: Union[bool, str] = False
     reference_id: Optional[str] = None
 
+    model_config = {"arbitrary_types_allowed": True}
+
 
 class UiPathRuntimeContext(BaseModel):
     """Context information passed throughout the runtime execution."""
@@ -144,6 +151,7 @@ class UiPathRuntimeContext(BaseModel):
     input: Optional[str] = None
     input_json: Optional[Any] = None
     job_id: Optional[str] = None
+    execution_id: Optional[str] = None
     trace_id: Optional[str] = None
     trace_context: Optional[UiPathTraceContext] = None
     tracing_enabled: Union[bool, str] = False
@@ -488,6 +496,15 @@ def __init__(self, runtime_class: Type[T], context_class: Type[C]):
 
         self.runtime_class = runtime_class
         self.context_class = context_class
+        self.tracer_provider: TracerProvider = TracerProvider()
+        self.tracer_span_processors: List[SpanProcessor] = []
+        trace.set_tracer_provider(self.tracer_provider)
+
+    def add_span_exporter(self, span_exporter: SpanExporter) -> None:
+        """Add a span processor to the tracer provider."""
+        span_processor = UiPathExecutionTraceProcessor(span_exporter)
+        self.tracer_span_processors.append(span_processor)
+        self.tracer_provider.add_span_processor(span_processor)
 
     def new_context(self, **kwargs) -> C:
         """Create a new context instance."""
@@ -500,4 +517,39 @@ def from_context(self, context: C) -> T:
     async def execute(self, context: C) -> Optional[UiPathRuntimeResult]:
         """Execute runtime with context."""
         async with self.from_context(context) as runtime:
-            return await runtime.execute()
+            result = await runtime.execute()
+            for span_processor in self.tracer_span_processors:
+                span_processor.force_flush()
+            return result
+
+    async def execute_in_root_span(
+        self, context: C, root_span: str = "root"
+    ) -> Optional[UiPathRuntimeResult]:
+        """Execute runtime with context."""
+        async with self.from_context(context) as runtime:
+            tracer: Tracer = trace.get_tracer("uipath-runtime")
+            with tracer.start_as_current_span(
+                root_span, attributes={"execution.id": context.execution_id}
+            ):
+                result = await runtime.execute()
+            for span_processor in self.tracer_span_processors:
+                span_processor.force_flush()
+            return result
+
+
+class UiPathExecutionTraceProcessor(BatchSpanProcessor):
+    def on_start(
+        self, span: Span, parent_context: Optional[context_api.Context] = None
+    ):
+        """Called when a span is started."""
+        if parent_context:
+            parent_span = trace.get_current_span(parent_context)
+        else:
+            parent_span = trace.get_current_span()
+
+        if parent_span and parent_span.is_recording():
+            run_id = parent_span.attributes.get("execution.id")
+            if run_id:
+                span.set_attribute("execution.id", run_id)
+
+        super().on_start(span, parent_context)

src/uipath/_cli/_runtime/_contracts.py

@@ -8,13 +8,8 @@
 from dataclasses import asdict, is_dataclass
 from typing import Any, Dict, Optional, Type, TypeVar, cast, get_type_hints
 
-from opentelemetry import trace
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from pydantic import BaseModel
 
-from uipath.tracing import LlmOpsHttpExporter
-
 from ._contracts import (
     UiPathBaseRuntime,
     UiPathErrorCategory,
@@ -42,19 +37,6 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
         """
         await self.validate()
 
-        shutdown_provider = False
-        trace_provider: Optional[TracerProvider] = None
-        span_processor: Optional[BatchSpanProcessor] = None
-
-        if self.context.trace_context.enabled:
-            trace_provider = trace.get_tracer_provider()
-            if not trace_provider:
-                trace_provider = TracerProvider()
-                trace.set_tracer_provider(trace_provider)
-                shutdown_provider = True
-            span_processor = BatchSpanProcessor(LlmOpsHttpExporter())
-            trace_provider.add_span_processor(span_processor)
-
         try:
             if self.context.entrypoint is None:
                 return None
@@ -83,13 +65,6 @@ async def execute(self) -> Optional[UiPathRuntimeResult]:
                 UiPathErrorCategory.SYSTEM,
             ) from e
 
-        finally:
-            if self.context.trace_context.enabled:
-                await span_processor.force_flush()
-                # Only shutdown if we created the provider
-                if shutdown_provider:
-                    await trace_provider.shutdown()
-
     async def validate(self) -> None:
         """Validate runtime inputs."""
         if not self.context.entrypoint:
@@ -127,12 +102,6 @@ async def cleanup(self) -> None:
 
     async def _execute_python_script(self, script_path: str, input_data: Any) -> Any:
         """Execute the Python script with the given input."""
-        # parent_span = trace.get_current_span()
-        # ctx = trace.set_span_in_context(parent_span)
-
-        # print(f"Before module execution - Current span: {parent_span}")
-        # print(f"Current span attributes: {getattr(parent_span, 'attributes', {})}")
-
         spec = importlib.util.spec_from_file_location("dynamic_module", script_path)
         if not spec or not spec.loader:
             raise UiPathRuntimeError(
@@ -144,16 +113,9 @@ async def _execute_python_script(self, script_path: str, input_data: Any) -> Any
 
         module = importlib.util.module_from_spec(spec)
 
-        # Attach the context BEFORE any module operations
-        # token = context_api.attach(ctx)
-
         try:
-            # print("Executing module with context attached")
             spec.loader.exec_module(module)
 
-            # active_span = trace.get_current_span()
-            # print(f"After module execution - Active span: {active_span}")
-
             # Execute the function while context is still attached
             for func_name in ["main", "run", "execute"]:
                 if hasattr(module, func_name):
@@ -209,26 +171,26 @@ async def _execute_python_script(self, script_path: str, input_data: Any) -> Any
                             UiPathErrorCategory.USER,
                         ) from e
 
-                    # Case 3: Dict parameter
-                    else:
-                        try:
-                            result = (
-                                await main_func(input_data)
-                                if is_async
-                                else main_func(input_data)
-                            )
-                            return (
-                                self._convert_from_class(result)
-                                if result is not None
-                                else {}
-                            )
-                        except Exception as e:
-                            raise UiPathRuntimeError(
-                                "FUNCTION_EXECUTION_ERROR",
-                                f"Error executing {func_name} function with dictionary input",
-                                f"Error: {str(e)}",
-                                UiPathErrorCategory.USER,
-                            ) from e
+                # Case 3: Dict parameter
+                else:
+                    try:
+                        result = (
+                            await main_func(input_data)
+                            if is_async
+                            else main_func(input_data)
+                        )
+                        return (
+                            self._convert_from_class(result)
+                            if result is not None
+                            else {}
+                        )
+                    except Exception as e:
+                        raise UiPathRuntimeError(
+                            "FUNCTION_EXECUTION_ERROR",
+                            f"Error executing {func_name} function with dictionary input",
+                            f"Error: {str(e)}",
+                            UiPathErrorCategory.USER,
+                        ) from e
 
             # If we get here, no main function was found
             raise UiPathRuntimeError(
@@ -249,10 +211,6 @@ async def _execute_python_script(self, script_path: str, input_data: Any) -> Any
                 UiPathErrorCategory.USER,
             ) from e
 
-        # finally:
-        # Only detach the context at the very end, after all operations
-        # context_api.detach(token)
-
     def _convert_to_class(self, data: Dict[str, Any], cls: Type[T]) -> T:
         """Convert a dictionary to either a dataclass, Pydantic model, or regular class instance."""
         # Handle Pydantic models