fix: add execution options

Author: cristipufu

src/uipath/runtime/__init__.py

@@ -1,6 +1,11 @@
 """UiPath Runtime Package."""
 
-from uipath.runtime.base import UiPathBaseRuntime, UiPathStreamNotSupportedError
+from uipath.runtime.base import (
+    UiPathBaseRuntime,
+    UiPathExecutionOptions,
+    UiPathExecutionRuntime,
+    UiPathStreamNotSupportedError,
+)
 from uipath.runtime.context import UiPathRuntimeContext
 from uipath.runtime.events import UiPathRuntimeEvent
 from uipath.runtime.factory import UiPathRuntimeFactory
@@ -16,6 +21,8 @@
 __all__ = [
     "UiPathRuntimeContext",
     "UiPathBaseRuntime",
+    "UiPathExecutionOptions",
+    "UiPathExecutionRuntime",
     "UiPathRuntimeFactory",
     "UiPathRuntimeResult",
     "UiPathRuntimeStatus",

src/uipath/runtime/base.py

@@ -6,10 +6,13 @@
     Any,
     AsyncGenerator,
     Generic,
+    List,
+    Literal,
     Optional,
     TypeVar,
 )
 
+from pydantic import BaseModel, Field
 from typing_extensions import override
 from uipath.core import UiPathTraceManager
 
@@ -33,6 +36,21 @@ class UiPathStreamNotSupportedError(NotImplementedError):
     pass
 
 
+class UiPathExecutionOptions(BaseModel):
+    """Execution-time options controlling runtime behavior."""
+
+    resume: bool = Field(
+        default=False,
+        description="Indicates whether to resume a suspended execution.",
+    )
+    breakpoints: Optional[List[str] | Literal["*"]] = Field(
+        default=None,
+        description="List of nodes or '*' to break on all steps.",
+    )
+
+    model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
+
+
 class UiPathBaseRuntime(ABC):
     """Base runtime class implementing the async context manager protocol.
 
@@ -51,13 +69,14 @@ async def get_schema(self) -> UiPathRuntimeSchema:
         raise NotImplementedError()
 
     @abstractmethod
-    async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
+    async def execute(
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
+    ) -> UiPathRuntimeResult:
         """Produce the agent output."""
         raise NotImplementedError()
 
     async def stream(
-        self,
-        input: dict[str, Any],
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Stream execution events in real-time.
 
@@ -126,8 +145,7 @@ def __init__(
             self.log_handler = UiPathRuntimeExecutionLogHandler(execution_id)
 
     async def execute(
-        self,
-        input: dict[str, Any],
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
     ) -> UiPathRuntimeResult:
         """Execute runtime with context."""
         if self.log_handler:
@@ -141,18 +159,17 @@ async def execute(
                 with self.trace_manager.start_execution_span(
                     self.root_span, execution_id=self.execution_id
                 ):
-                    return await self.delegate.execute(input)
+                    return await self.delegate.execute(input, options=options)
             else:
-                return await self.delegate.execute(input)
+                return await self.delegate.execute(input, options=options)
         finally:
             self.trace_manager.flush_spans()
             if self.log_handler:
                 log_interceptor.teardown()
 
     @override
     async def stream(
-        self,
-        input: dict[str, Any],
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Stream runtime execution with context.
 
@@ -176,7 +193,7 @@ async def stream(
                 with self.trace_manager.start_execution_span(
                     self.root_span, execution_id=self.execution_id
                 ):
-                    async for event in self.delegate.stream(input):
+                    async for event in self.delegate.stream(input, options=options):
                         yield event
         finally:
             self.trace_manager.flush_spans()

src/uipath/runtime/context.py

@@ -6,8 +6,6 @@
 from functools import cached_property
 from typing import (
     Any,
-    List,
-    Literal,
     Optional,
     TypeVar,
 )
@@ -35,7 +33,6 @@ class UiPathRuntimeContext(BaseModel):
 
     entrypoint: Optional[str] = None
     input: Optional[Any] = None
-    resume: bool = False
     job_id: Optional[str] = None
     trace_context: Optional[UiPathTraceContext] = None
     config_path: str = "uipath.json"
@@ -47,7 +44,6 @@ class UiPathRuntimeContext(BaseModel):
     trace_file: Optional[str] = None
     logs_file: Optional[str] = "execution.log"
     logs_min_level: Optional[str] = "INFO"
-    breakpoints: Optional[List[str] | Literal["*"]] = None
     result: Optional[UiPathRuntimeResult] = None
 
     model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
@@ -111,9 +107,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         """
         try:
             if self.result is None:
-                execution_result = UiPathRuntimeResult()
-            else:
-                execution_result = self.result
+                self.result = UiPathRuntimeResult()
 
             if exc_type:
                 # Create error info from exception
@@ -128,10 +122,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                         category=UiPathErrorCategory.UNKNOWN,
                     )
 
-                execution_result.status = UiPathRuntimeStatus.FAULTED
-                execution_result.error = error_info
+                self.result.status = UiPathRuntimeStatus.FAULTED
+                self.result.error = error_info
 
-            content = execution_result.to_dict()
+            content = self.result.to_dict()
 
             # Always write output file at runtime, except for inner runtimes
             # Inner runtimes have execution_id

src/uipath/runtime/debug/runtime.py

@@ -1,12 +1,12 @@
 """Debug runtime implementation."""
 
 import logging
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Optional, TypeVar
 
 from uipath.runtime import (
     UiPathBaseRuntime,
     UiPathBreakpointResult,
-    UiPathRuntimeContext,
+    UiPathExecutionOptions,
     UiPathRuntimeResult,
     UiPathRuntimeStatus,
     UiPathStreamNotSupportedError,
@@ -26,53 +26,48 @@ class UiPathDebugRuntime(UiPathBaseRuntime, Generic[T]):
 
     def __init__(
         self,
-        context: UiPathRuntimeContext,
         delegate: T,
         debug_bridge: UiPathDebugBridge,
     ):
         """Initialize the UiPathDebugRuntime."""
-        super().__init__(context)
-        self.context: UiPathRuntimeContext = context
+        super().__init__()
         self.delegate: T = delegate
         self.debug_bridge: UiPathDebugBridge = debug_bridge
 
-    async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
+    async def execute(
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
+    ) -> UiPathRuntimeResult:
         """Execute the workflow with debug support."""
         try:
             await self.debug_bridge.connect()
 
             await self.debug_bridge.emit_execution_started()
 
             result: UiPathRuntimeResult
+
             # Try to stream events from inner runtime
             try:
-                result = await self._stream_and_debug(self.delegate, input)
+                result = await self._stream_and_debug(input, options=options)
             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(input)
+                result = await self.delegate.execute(input, options=options)
 
             await self.debug_bridge.emit_execution_completed(result)
 
-            self.context.result = result
-
             return result
 
         except Exception as e:
-            # Emit execution error
-            self.context.result = UiPathRuntimeResult(
-                status=UiPathRuntimeStatus.FAULTED,
-            )
             await self.debug_bridge.emit_execution_error(
                 error=str(e),
             )
             raise
 
     async def _stream_and_debug(
-        self, inner_runtime: T, input: dict[str, Any]
+        self, input: dict[str, Any], options: Optional[UiPathExecutionOptions] = None
     ) -> UiPathRuntimeResult:
         """Stream events from inner runtime and handle debug interactions."""
         final_result: UiPathRuntimeResult
@@ -81,13 +76,17 @@ async def _stream_and_debug(
         # Starting in paused state - wait for breakpoints and resume
         await self.debug_bridge.wait_for_resume()
 
+        debug_options = UiPathExecutionOptions(
+            resume=options.resume if options else False,
+            breakpoints=options.breakpoints if options else None,
+        )
+
         # Keep streaming until execution completes (not just paused at breakpoint)
         while not execution_completed:
             # Update breakpoints from debug bridge
-            if inner_runtime.context is not None:
-                inner_runtime.context.breakpoints = self.debug_bridge.get_breakpoints()
+            debug_options.breakpoints = self.debug_bridge.get_breakpoints()
             # Stream events from inner runtime
-            async for event in inner_runtime.stream(input):
+            async for event in self.delegate.stream(input, options=debug_options):
                 # Handle final result
                 if isinstance(event, UiPathRuntimeResult):
                     final_result = event
@@ -99,8 +98,8 @@ async def _stream_and_debug(
                             await self.debug_bridge.emit_breakpoint_hit(event)
                             await self.debug_bridge.wait_for_resume()
 
-                            if inner_runtime.context is not None:
-                                inner_runtime.context.resume = True
+                            # Tell inner runtime we're resuming
+                            debug_options.resume = True
 
                         except UiPathDebugQuitError:
                             final_result = UiPathRuntimeResult(
@@ -122,9 +121,6 @@ async def _stream_and_debug(
     async def cleanup(self) -> None:
         """Cleanup runtime resources."""
         try:
-            await self.delegate.cleanup()
-        finally:
-            try:
-                await self.debug_bridge.disconnect()
-            except Exception as e:
-                logger.warning(f"Error disconnecting debug bridge: {e}")
+            await self.debug_bridge.disconnect()
+        except Exception as e:
+            logger.warning(f"Error disconnecting debug bridge: {e}")