refactor(RuntimeFactory): refactor runtime and factory

This is done in order to follow single responsibility principles (SRP).

Author: akshaylive

src/uipath/runtime/__init__.py

@@ -3,7 +3,7 @@
 from uipath.runtime.base import UiPathBaseRuntime
 from uipath.runtime.context import UiPathRuntimeContext
 from uipath.runtime.events import UiPathRuntimeEvent
-from uipath.runtime.factory import UiPathRuntimeExecutor, UiPathRuntimeFactory
+from uipath.runtime.factory import UiPathRuntimeFactory
 from uipath.runtime.result import (
     UiPathApiTrigger,
     UiPathBreakpointResult,
@@ -16,7 +16,6 @@
     "UiPathRuntimeContext",
     "UiPathBaseRuntime",
     "UiPathRuntimeFactory",
-    "UiPathRuntimeExecutor",
     "UiPathRuntimeResult",
     "UiPathRuntimeEvent",
     "UiPathBreakpointResult",

src/uipath/runtime/base.py

@@ -1,26 +1,38 @@
 """Base runtime class and async context manager implementation."""
 
-import json
 import logging
-import os
 from abc import ABC, abstractmethod
-from typing import AsyncGenerator
+from typing import (
+    Any,
+    AsyncGenerator,
+    Generic,
+    List,
+    Optional,
+    TypeVar,
+)
 
-from uipath.runtime.context import UiPathRuntimeContext
-from uipath.runtime.errors import (
-    UiPathErrorCategory,
-    UiPathErrorCode,
-    UiPathErrorContract,
-    UiPathRuntimeError,
+from opentelemetry import trace
+from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
+    BaseInstrumentor,
 )
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider
+from opentelemetry.sdk.trace.export import SpanExporter
+from typing_extensions import override
+
+from uipath.runtime.context import UiPathRuntimeContext
 from uipath.runtime.events import (
     UiPathRuntimeEvent,
 )
+from uipath.runtime.logging import UiPathRuntimeExecutionLogHandler
 from uipath.runtime.logging._interceptor import UiPathRuntimeLogsInterceptor
-from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
 from uipath.runtime.schema import (
     UiPathRuntimeSchema,
 )
+from uipath.runtime.tracing import (
+    UiPathExecutionBatchTraceProcessor,
+    UiPathExecutionSimpleTraceProcessor,
+    UiPathRuntimeExecutionSpanExporter,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +49,7 @@ class UiPathBaseRuntime(ABC):
     This allows using the class with 'async with' statements.
     """
 
-    def __init__(self, context: UiPathRuntimeContext):
+    def __init__(self, context: Optional[UiPathRuntimeContext] = None):
         """Initialize the runtime with the provided context."""
         self.context = context
 
@@ -48,70 +60,10 @@ async def get_schema(self) -> UiPathRuntimeSchema:
         """
         raise NotImplementedError()
 
-    async def __aenter__(self):
-        """Async enter method called when entering the 'async with' block.
-
-        Initializes and prepares the runtime environment.
-
-        Returns:
-            The runtime instance
-        """
-        # Read the input from file if provided
-        if self.context.input_file:
-            _, file_extension = os.path.splitext(self.context.input_file)
-            if file_extension != ".json":
-                raise UiPathRuntimeError(
-                    code=UiPathErrorCode.INVALID_INPUT_FILE_EXTENSION,
-                    title="Invalid Input File Extension",
-                    detail="The provided input file must be in JSON format.",
-                )
-            with open(self.context.input_file) as f:
-                self.context.input = f.read()
-
-        try:
-            if isinstance(self.context.input, str):
-                if self.context.input.strip():
-                    self.context.input = json.loads(self.context.input)
-                else:
-                    self.context.input = {}
-            elif self.context.input is None:
-                self.context.input = {}
-            # else: leave it as-is (already a dict, list, bool, etc.)
-        except json.JSONDecodeError as e:
-            raise UiPathRuntimeError(
-                UiPathErrorCode.INPUT_INVALID_JSON,
-                "Invalid JSON input",
-                f"The input data is not valid JSON: {str(e)}",
-                UiPathErrorCategory.USER,
-            ) from e
-
-        await self.validate()
-
-        # Intercept all stdout/stderr/logs
-        # Write to file (runtime), stdout (debug) or log handler (if provided)
-        self.logs_interceptor = UiPathRuntimeLogsInterceptor(
-            min_level=self.context.logs_min_level,
-            dir=self.context.runtime_dir,
-            file=self.context.logs_file,
-            job_id=self.context.job_id,
-            execution_id=self.context.execution_id,
-            log_handler=self.context.log_handler,
-        )
-        self.logs_interceptor.setup()
-
-        return self
-
     @abstractmethod
-    async def execute(self) -> UiPathRuntimeResult:
-        """Execute with the provided context.
-
-        Returns:
-            Dictionary with execution results
-
-        Raises:
-            RuntimeError: If execution fails
-        """
-        pass
+    async def execute(self, input: dict[str, Any]) -> Any:
+        """Produce the agent output."""
+        raise NotImplementedError()
 
     async def stream(
         self,
@@ -154,101 +106,141 @@ async def stream(
         # Without it, the function wouldn't match the AsyncGenerator return type
         yield
 
-    @abstractmethod
-    async def validate(self):
-        """Validate runtime inputs."""
-        pass
+    def get_instrumentor(self) -> Optional[BaseInstrumentor]:
+        """Get instrumentor for this runtime. If no instrumentor is available, return None."""
+        return None
 
     @abstractmethod
     async def cleanup(self):
         """Cleaup runtime resources."""
         pass
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Async exit method called when exiting the 'async with' block.
 
-        Cleans up resources and handles any exceptions.
+T = TypeVar("T", bound=UiPathBaseRuntime)
+
+
+class TraceManager:
+    """Trace manager."""
+
+    def __init__(self):
+        """Initialize a trace manager."""
+        self.tracer_provider: TracerProvider = TracerProvider()
+        trace.set_tracer_provider(self.tracer_provider)
+        self.tracer_span_processors: List[SpanProcessor] = []
+        self.execution_span_exporter = UiPathRuntimeExecutionSpanExporter()
+        self.add_span_exporter(self.execution_span_exporter)
 
-        Always writes output file regardless of whether execution was successful,
-        suspended, or encountered an error.
+    def add_span_exporter(
+        self,
+        span_exporter: SpanExporter,
+        batch: bool = True,
+    ) -> "TraceManager":
+        """Add a span processor to the tracer provider."""
+        span_processor: SpanProcessor
+        if batch:
+            span_processor = UiPathExecutionBatchTraceProcessor(span_exporter)
+        else:
+            span_processor = UiPathExecutionSimpleTraceProcessor(span_exporter)
+        self.tracer_span_processors.append(span_processor)
+        self.tracer_provider.add_span_processor(span_processor)
+        return self
+
+    def get_execution_spans(
+        self,
+        execution_id: str,
+    ) -> List[ReadableSpan]:
+        """Retrieve spans for a given execution id."""
+        return self.execution_span_exporter.get_spans(execution_id)
+
+    def flush_spans(self) -> None:
+        """Flush all span processors."""
+        for span_processor in self.tracer_span_processors:
+            span_processor.force_flush()
+
+
+class UiPathExecutionRuntime(UiPathBaseRuntime, Generic[T]):
+    """Handles runtime execution with tracing/telemetry."""
+
+    def __init__(
+        self,
+        delegate: T,
+        trace_manager: TraceManager,
+        root_span: str = "root",
+        execution_id: Optional[str] = None,
+    ):
+        """Initialize the executor."""
+        self.delegate = delegate
+        self.trace_manager = trace_manager
+        self.root_span = root_span
+        self.execution_id = execution_id
+        self.log_handler: Optional[UiPathRuntimeExecutionLogHandler]
+        if execution_id is not None:
+            self.log_handler = UiPathRuntimeExecutionLogHandler(execution_id)
+        else:
+            self.log_handler = None
+
+    async def execute(
+        self,
+        input: dict[str, Any],
+    ) -> Any:
+        """Execute runtime with context."""
+        instrumentor = self.delegate.get_instrumentor()
+        log_interceptor = UiPathRuntimeLogsInterceptor(
+            execution_id=self.execution_id, log_handler=self.log_handler
+        )
+        log_interceptor.setup()
+
+        if instrumentor is not None:
+            instrumentor.instrument(tracer_provider=self.trace_manager.tracer_provider)
+        try:
+            tracer = trace.get_tracer("uipath-runtime")
+            span_attributes: dict[str, Any] = {}
+            if self.execution_id:
+                span_attributes["execution.id"] = self.execution_id
+            with tracer.start_as_current_span(
+                self.root_span, attributes=span_attributes
+            ):
+                return await self.delegate.execute(input)
+        finally:
+            self.trace_manager.flush_spans()
+            if instrumentor is not None:
+                instrumentor.uninstrument()
+            log_interceptor.teardown()
+
+    @override
+    async def stream(
+        self,
+        root_span: str = "root",
+        execution_id: Optional[str] = None,
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        """Stream runtime execution with context.
+
+        Args:
+            runtime: The runtime instance
+            context: The runtime context
+
+        Yields:
+            UiPathRuntimeEvent instances during execution and final UiPathRuntimeResult
+
+        Raises:
+            UiPathRuntimeStreamNotSupportedError: If the runtime doesn't support streaming
         """
+        instrumentor = self.delegate.get_instrumentor()
+        if instrumentor is not None:
+            instrumentor.instrument(tracer_provider=self.trace_manager.tracer_provider)
         try:
-            if self.context.result is None:
-                execution_result = UiPathRuntimeResult()
-            else:
-                execution_result = self.context.result
-
-            if exc_type:
-                # Create error info from exception
-                if isinstance(exc_val, UiPathRuntimeError):
-                    error_info = exc_val.error_info
-                else:
-                    # Generic error
-                    error_info = UiPathErrorContract(
-                        code=f"ERROR_{exc_type.__name__}",
-                        title=f"Runtime error: {exc_type.__name__}",
-                        detail=str(exc_val),
-                        category=UiPathErrorCategory.UNKNOWN,
-                    )
-
-                execution_result.status = UiPathRuntimeStatus.FAULTED
-                execution_result.error = error_info
-
-            content = execution_result.to_dict()
-
-            # Always write output file at runtime, except for inner runtimes
-            # Inner runtimes have execution_id
-            if self.context.job_id and not self.context.execution_id:
-                with open(self.context.result_file_path, "w") as f:
-                    json.dump(content, f, indent=2, default=str)
-
-            # Write the execution output to file if requested
-            if self.context.output_file:
-                with open(self.context.output_file, "w") as f:
-                    f.write(content.get("output", "{}"))
-
-            # Don't suppress exceptions
-            return False
-
-        except Exception as e:
-            logger.error(f"Error during runtime shutdown: {str(e)}")
-
-            # Create a fallback error result if we fail during cleanup
-            if not isinstance(e, UiPathRuntimeError):
-                error_info = UiPathErrorContract(
-                    code="RUNTIME_SHUTDOWN_ERROR",
-                    title="Runtime shutdown failed",
-                    detail=f"Error: {str(e)}",
-                    category=UiPathErrorCategory.SYSTEM,
-                )
-            else:
-                error_info = e.error_info
-
-            # Last-ditch effort to write error output
-            try:
-                error_result = UiPathRuntimeResult(
-                    status=UiPathRuntimeStatus.FAULTED, error=error_info
-                )
-                error_result_content = error_result.to_dict()
-                if self.context.job_id:
-                    with open(self.context.result_file_path, "w") as f:
-                        json.dump(error_result_content, f, indent=2, default=str)
-            except Exception as write_error:
-                logger.error(f"Failed to write error output file: {str(write_error)}")
-                raise
-
-            # Re-raise as RuntimeError if it's not already a UiPathRuntimeError
-            if not isinstance(e, UiPathRuntimeError):
-                raise RuntimeError(
-                    error_info.code,
-                    error_info.title,
-                    error_info.detail,
-                    error_info.category,
-                ) from e
-            raise
+            tracer = trace.get_tracer("uipath-runtime")
+            span_attributes: dict[str, Any] = {}
+            if execution_id:
+                span_attributes["execution.id"] = "exec-a"
+            with tracer.start_as_current_span(root_span, attributes=span_attributes):
+                async for event in self.delegate.stream():
+                    yield event
         finally:
-            # Restore original logging
-            if hasattr(self, "logs_interceptor"):
-                self.logs_interceptor.teardown()
+            self.trace_manager.flush_spans()
+            if instrumentor is not None:
+                instrumentor.uninstrument()
 
-            await self.cleanup()
+    def cleanup(self) -> None:
+        """Close runtime resources."""
+        pass