Hide breakpoint logic from runtimes

Author: akshaylive

src/uipath/runtime/debug/runtime.py

@@ -1,12 +1,13 @@
 """Debug runtime implementation."""
 
 import logging
-from typing import Any, Generic, TypeVar
+from typing import Any, AsyncGenerator, Generic, TypeVar
 
 from uipath.runtime import (
     UiPathBaseRuntime,
     UiPathBreakpointResult,
     UiPathRuntimeContext,
+    UiPathRuntimeEvent,
     UiPathRuntimeResult,
     UiPathRuntimeStatus,
     UiPathStreamNotSupportedError,
@@ -46,7 +47,7 @@ async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
             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)
             except UiPathStreamNotSupportedError:
                 # Fallback to regular execute if streaming not supported
                 logger.debug(
@@ -71,9 +72,38 @@ async def execute(self, input: dict[str, Any]) -> UiPathRuntimeResult:
             )
             raise
 
-    async def _stream_and_debug(
-        self, inner_runtime: T, input: dict[str, Any]
-    ) -> UiPathRuntimeResult:
+    async def stream(
+        self,
+        input: dict[str, Any],
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        """Stream events from inner runtime and handle debug interactions."""
+        async for event in self.delegate.stream(input):
+            if isinstance(event, UiPathRuntimeStateEvent):
+                # Move these to constructor and remove from context
+                if self.context:
+                    breakpoints = self.context.breakpoints
+                else:
+                    breakpoints = None
+                hit_breakpoint = False
+
+                if breakpoints == "*":
+                    hit_breakpoint = True
+                elif isinstance(breakpoints, list) and event.node_name in breakpoints:
+                    hit_breakpoint = True
+
+                if hit_breakpoint:
+                    yield UiPathBreakpointResult(
+                        breakpoint_node=event.node_name,
+                        breakpoint_type="before",
+                        next_nodes=event.next_nodes,
+                        current_state={"node": event.node_name},
+                    )
+                else:
+                    yield event
+            else:
+                yield event
+
+    async def _stream_and_debug(self, input: dict[str, Any]) -> UiPathRuntimeResult:
         """Stream events from inner runtime and handle debug interactions."""
         final_result: UiPathRuntimeResult
         execution_completed = False
@@ -84,10 +114,10 @@ async def _stream_and_debug(
         # 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()
+            if self.delegate.context is not None:
+                self.delegate.context.breakpoints = self.debug_bridge.get_breakpoints()
             # Stream events from inner runtime
-            async for event in inner_runtime.stream(input):
+            async for event in self.stream(input):
                 # Handle final result
                 if isinstance(event, UiPathRuntimeResult):
                     final_result = event
@@ -99,8 +129,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
+                            if self.delegate.context is not None:
+                                self.delegate.context.resume = True
 
                         except UiPathDebugQuitError:
                             final_result = UiPathRuntimeResult(

src/uipath/runtime/events/state.py

@@ -1,6 +1,6 @@
 """Events related to agent messages and state updates."""
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List
 
 from pydantic import Field
 
@@ -62,8 +62,12 @@ class UiPathRuntimeStateEvent(UiPathRuntimeEvent):
     """
 
     payload: Dict[str, Any] = Field(description="Framework-specific state update")
-    node_name: Optional[str] = Field(
-        default=None, description="Name of the node/agent that caused this update"
+    node_name: str = Field(
+        ..., description="Name of the node/agent that caused this update"
+    )
+    next_nodes: List[str] = Field(
+        default_factory=list,
+        description="List of node names that the agent should follow",
     )
     event_type: UiPathRuntimeEventType = Field(
         default=UiPathRuntimeEventType.RUNTIME_STATE, frozen=True

tests/test_debugger.py

@@ -9,7 +9,6 @@
 
 from uipath.runtime import (
     UiPathBaseRuntime,
-    UiPathBreakpointResult,
     UiPathRuntimeContext,
     UiPathRuntimeResult,
     UiPathRuntimeStatus,
@@ -90,29 +89,8 @@ async def stream(
             yield UiPathRuntimeStateEvent(
                 node_name=node,
                 payload={"index": idx, "node": node},
+                next_nodes=self.node_sequence[idx + 1 : idx + 2],  # at most one
             )
-
-            # 2) Check for breakpoints on this node
-            if self.context:
-                breakpoints = self.context.breakpoints
-            else:
-                breakpoints = None
-            hit_breakpoint = False
-
-            if breakpoints == "*":
-                hit_breakpoint = True
-            elif isinstance(breakpoints, list) and node in breakpoints:
-                hit_breakpoint = True
-
-            if hit_breakpoint:
-                next_nodes = self.node_sequence[idx + 1 : idx + 2]  # at most one
-                yield UiPathBreakpointResult(
-                    breakpoint_node=node,
-                    breakpoint_type="before",
-                    next_nodes=next_nodes,
-                    current_state={"node": node, "index": idx},
-                )
-
         # 3) Final result at the end of streaming
         yield UiPathRuntimeResult(
             status=UiPathRuntimeStatus.SUCCESSFUL,