fix: runtime context as context manager

Author: cristipufu

src/uipath/runtime/base.py

@@ -10,7 +10,6 @@
     TypeVar,
 )
 
-from opentelemetry import trace
 from typing_extensions import override
 from uipath.core import UiPathTraceManager
 
@@ -20,6 +19,7 @@
 )
 from uipath.runtime.logging import UiPathRuntimeExecutionLogHandler
 from uipath.runtime.logging._interceptor import UiPathRuntimeLogsInterceptor
+from uipath.runtime.result import UiPathRuntimeResult
 from uipath.runtime.schema import (
     UiPathRuntimeSchema,
 )
@@ -51,12 +51,13 @@ async def get_schema(self) -> UiPathRuntimeSchema:
         raise NotImplementedError()
 
     @abstractmethod
-    async def execute(self, input: dict[str, Any]) -> Any:
+    async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
         """Produce the agent output."""
         raise NotImplementedError()
 
     async def stream(
         self,
+        input: dict[str, Any],
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Stream execution events in real-time.
 
@@ -120,16 +121,14 @@ def __init__(
         self.trace_manager = trace_manager
         self.root_span = root_span
         self.execution_id = execution_id
-        self.log_handler: Optional[UiPathRuntimeExecutionLogHandler]
+        self.log_handler: Optional[UiPathRuntimeExecutionLogHandler] = None
         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:
+    ) -> UiPathRuntimeResult:
         """Execute runtime with context."""
         if self.log_handler:
             log_interceptor = UiPathRuntimeLogsInterceptor(
@@ -153,8 +152,7 @@ async def execute(
     @override
     async def stream(
         self,
-        root_span: str = "root",
-        execution_id: Optional[str] = None,
+        input: dict[str, Any],
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Stream runtime execution with context.
 
@@ -166,18 +164,24 @@ async def stream(
             UiPathRuntimeEvent instances during execution and final UiPathRuntimeResult
 
         Raises:
-            UiPathRuntimeStreamNotSupportedError: If the runtime doesn't support streaming
+            UiPathStreamNotSupportedError: If the runtime doesn't support streaming
         """
+        if self.log_handler:
+            log_interceptor = UiPathRuntimeLogsInterceptor(
+                execution_id=self.execution_id, log_handler=self.log_handler
+            )
+            log_interceptor.setup()
         try:
-            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
+            if self.execution_id:
+                with self.trace_manager.start_execution_span(
+                    self.root_span, execution_id=self.execution_id
+                ):
+                    async for event in self.delegate.stream(input):
+                        yield event
         finally:
             self.trace_manager.flush_spans()
+            if self.log_handler:
+                log_interceptor.teardown()
 
     def cleanup(self) -> None:
         """Close runtime resources."""

src/uipath/runtime/context.py

@@ -16,7 +16,16 @@
 from pydantic import BaseModel
 from uipath.core.tracing.context import UiPathTraceContext
 
-from uipath.runtime.result import UiPathRuntimeResult
+from uipath.runtime.errors import (
+    UiPathErrorCategory,
+    UiPathErrorCode,
+    UiPathErrorContract,
+    UiPathRuntimeError,
+)
+from uipath.runtime.logging._interceptor import UiPathRuntimeLogsInterceptor
+from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
+
+logger = logging.getLogger(__name__)
 
 C = TypeVar("C", bound="UiPathRuntimeContext")
 
@@ -28,7 +37,6 @@ class UiPathRuntimeContext(BaseModel):
     input: Optional[Any] = None
     resume: bool = False
     job_id: Optional[str] = None
-    execution_id: Optional[str] = None
     trace_context: Optional[UiPathTraceContext] = None
     config_path: str = "uipath.json"
     runtime_dir: Optional[str] = "__uipath"
@@ -38,13 +46,148 @@ class UiPathRuntimeContext(BaseModel):
     output_file: Optional[str] = None
     trace_file: Optional[str] = None
     logs_file: Optional[str] = "execution.log"
-    log_handler: Optional[logging.Handler] = None
     logs_min_level: Optional[str] = "INFO"
     breakpoints: Optional[List[str] | Literal["*"]] = None
     result: Optional[UiPathRuntimeResult] = None
 
     model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
 
+    def __enter__(self):
+        """Async enter method called when entering the 'async with' block.
+
+        Initializes and prepares the runtime contextual environment.
+
+        Returns:
+            The runtime context instance
+        """
+        # Read the input from file if provided
+        if self.input_file:
+            _, file_extension = os.path.splitext(self.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.input_file) as f:
+                self.input = f.read()
+
+        try:
+            if isinstance(self.input, str):
+                if self.input.strip():
+                    self.input = json.loads(self.input)
+                else:
+                    self.input = {}
+            elif self.input is None:
+                self.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
+
+        # Intercept all stdout/stderr/logs
+        # Write to file (runtime), stdout (debug) or log handler (if provided)
+        self.logs_interceptor = UiPathRuntimeLogsInterceptor(
+            min_level=self.logs_min_level,
+            dir=self.runtime_dir,
+            file=self.logs_file,
+            job_id=self.job_id,
+        )
+        self.logs_interceptor.setup()
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Async exit method called when exiting the 'async with' block.
+
+        Cleans up resources and handles any exceptions.
+
+        Always writes output file regardless of whether execution was successful,
+        suspended, or encountered an error.
+        """
+        try:
+            if self.result is None:
+                execution_result = UiPathRuntimeResult()
+            else:
+                execution_result = self.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.job_id:
+                with open(self.result_file_path, "w") as f:
+                    json.dump(content, f, indent=2, default=str)
+
+            # Write the execution output to file if requested
+            if self.output_file:
+                with open(self.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.job_id:
+                    with open(self.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
+        finally:
+            # Restore original logging
+            if hasattr(self, "logs_interceptor"):
+                self.logs_interceptor.teardown()
+
     @cached_property
     def result_file_path(self) -> str:
         """Get the full path to the result file."""

src/uipath/runtime/debug/runtime.py

@@ -1,7 +1,7 @@
 """Debug runtime implementation."""
 
 import logging
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 from uipath.runtime import (
     UiPathBaseRuntime,
@@ -36,7 +36,7 @@ def __init__(
         self.delegate: T = delegate
         self.debug_bridge: UiPathDebugBridge = debug_bridge
 
-    async def execute(self) -> UiPathRuntimeResult:
+    async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
         """Execute the workflow with debug support."""
         try:
             await self.debug_bridge.connect()
@@ -46,14 +46,14 @@ async def execute(self) -> UiPathRuntimeResult:
             result: UiPathRuntimeResult
             # Try to stream events from inner runtime
             try:
-                result = await self._stream_and_debug(self.delegate)
+                result = await self._stream_and_debug(self.delegate, input)
             except UiPathStreamNotSupportedError:
                 # Fallback to regular execute if streaming not supported
                 logger.debug(
                     f"Runtime {self.delegate.__class__.__name__} does not support "
                     "streaming, falling back to execute()"
                 )
-                result = await self.delegate.execute()
+                result = await self.delegate.execute(input)
 
             await self.debug_bridge.emit_execution_completed(result)
 
@@ -71,7 +71,9 @@ async def execute(self) -> UiPathRuntimeResult:
             )
             raise
 
-    async def _stream_and_debug(self, inner_runtime: T) -> UiPathRuntimeResult:
+    async def _stream_and_debug(
+        self, inner_runtime: T, input: dict[str, Any]
+    ) -> UiPathRuntimeResult:
         """Stream events from inner runtime and handle debug interactions."""
         final_result: UiPathRuntimeResult
         execution_completed = False
@@ -82,9 +84,10 @@ async def _stream_and_debug(self, inner_runtime: T) -> UiPathRuntimeResult:
         # Keep streaming until execution completes (not just paused at breakpoint)
         while not execution_completed:
             # Update breakpoints from debug bridge
-            inner_runtime.context.breakpoints = self.debug_bridge.get_breakpoints()
+            if inner_runtime.context is not None:
+                inner_runtime.context.breakpoints = self.debug_bridge.get_breakpoints()
             # Stream events from inner runtime
-            async for event in inner_runtime.stream():
+            async for event in inner_runtime.stream(input):
                 # Handle final result
                 if isinstance(event, UiPathRuntimeResult):
                     final_result = event
@@ -96,7 +99,8 @@ async def _stream_and_debug(self, inner_runtime: T) -> UiPathRuntimeResult:
                             await self.debug_bridge.emit_breakpoint_hit(event)
                             await self.debug_bridge.wait_for_resume()
 
-                            self.delegate.context.resume = True
+                            if inner_runtime.context is not None:
+                                inner_runtime.context.resume = True
 
                         except UiPathDebugQuitError:
                             final_result = UiPathRuntimeResult(
@@ -115,10 +119,6 @@ async def _stream_and_debug(self, inner_runtime: T) -> UiPathRuntimeResult:
 
         return final_result
 
-    async def validate(self) -> None:
-        """Validate runtime configuration."""
-        await self.delegate.validate()
-
     async def cleanup(self) -> None:
         """Cleanup runtime resources."""
         try: