From 9ef7a5951b2823121b47c5c4d226927d313d6494 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Mon, 5 Oct 2026 14:54:12 +0300 Subject: [PATCH 1/3] feat(runtime): end a run the host stopped as stopped A run cancelled because its host asked it to stop was reported as faulted with ERROR_CancelledError, indistinguishable from a crash. A host can now install a stop-requested probe (uipath.runtime.stop). When the run ends in a cancellation, or in an error raised because of one, and the probe says a stop was requested, the context writes the new stopped status with no error, and the result sink receives it like any final result. Without a probe, or when the probe says no, nothing changes. Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 2 +- src/uipath/runtime/context.py | 34 ++++++- src/uipath/runtime/result.py | 1 + src/uipath/runtime/stop.py | 28 ++++++ tests/test_stop.py | 162 ++++++++++++++++++++++++++++++++++ uv.lock | 2 +- 6 files changed, 226 insertions(+), 3 deletions(-) create mode 100644 src/uipath/runtime/stop.py create mode 100644 tests/test_stop.py diff --git a/pyproject.toml b/pyproject.toml index ebe7206f..a6db49ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-runtime" -version = "0.13.5" +version = "0.13.6" description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/src/uipath/runtime/context.py b/src/uipath/runtime/context.py index cdccf045..9290d942 100644 --- a/src/uipath/runtime/context.py +++ b/src/uipath/runtime/context.py @@ -1,5 +1,6 @@ """Context information passed throughout the runtime execution.""" +import asyncio import json import logging import os @@ -21,6 +22,7 @@ from uipath.runtime.logging._interceptor import UiPathRuntimeLogsInterceptor from uipath.runtime.output_sinks import ResultSink, get_log_handler, get_result_sink from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus +from uipath.runtime.stop import StopRequestedProbe, get_stop_requested_probe logger = logging.getLogger(__name__) @@ -127,6 +129,7 @@ class UiPathRuntimeContext(BaseModel): # installed when the context started is the one that receives the result — even if the registry is # changed in between or __exit__ runs on a different event loop. _result_sink: ResultSink | None = PrivateAttr(default=None) + _stop_requested_probe: StopRequestedProbe | None = PrivateAttr(default=None) def _apply_execution_source(self) -> None: """Derive execution_source from the command, if not already set. @@ -247,6 +250,7 @@ def __enter__(self): # the result sink is stashed for __exit__ so the two are read at the same moment (same context). log_handler = get_log_handler() self._result_sink = get_result_sink() + self._stop_requested_probe = get_stop_requested_probe() self.logs_interceptor = UiPathRuntimeLogsInterceptor( min_level=self.logs_min_level, @@ -271,7 +275,10 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self.result is None: self.result = UiPathRuntimeResult() - if exc_type and exc_val: + if exc_type and exc_val and self._stopped_on_request(exc_val): + self.result.status = UiPathRuntimeStatus.STOPPED + self.result.error = None + elif exc_type and exc_val: # Create error info from exception match exc_val: case UiPathFaultedTriggerError(): @@ -327,6 +334,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if result_sink is not None and self.result.status in ( UiPathRuntimeStatus.SUCCESSFUL, UiPathRuntimeStatus.FAULTED, + UiPathRuntimeStatus.STOPPED, ): try: self._deliver_result(result_sink, output_payload) @@ -377,6 +385,19 @@ def __exit__(self, exc_type, exc_val, exc_tb): if hasattr(self, "logs_interceptor"): self.logs_interceptor.teardown() + def _stopped_on_request(self, exc: BaseException) -> bool: + """Whether the run ended by a cancellation the host asked for.""" + probe = self._stop_requested_probe + if probe is None or not _caused_by_cancellation(exc): + return False + try: + return bool(probe()) + except Exception: + logger.exception( + "Stop-requested probe failed; the run is reported as faulted" + ) + return False + def _deliver_result(self, sink: ResultSink, output_payload: Any) -> None: """Spill the output arguments to a file, then hand the result + that path to the sink.""" args_path = self.resolved_output_arguments_file_path @@ -527,3 +548,14 @@ def from_config( setattr(instance, attr_name, kwargs[attr_name]) return instance + + +def _caused_by_cancellation(exc: BaseException | None) -> bool: + """Whether ``exc`` is a cancellation, or was raised while handling or because of one.""" + seen: set[int] = set() + while exc is not None and id(exc) not in seen: + if isinstance(exc, asyncio.CancelledError): + return True + seen.add(id(exc)) + exc = exc.__cause__ or exc.__context__ + return False diff --git a/src/uipath/runtime/result.py b/src/uipath/runtime/result.py index 66908f12..e692199f 100644 --- a/src/uipath/runtime/result.py +++ b/src/uipath/runtime/result.py @@ -16,6 +16,7 @@ class UiPathRuntimeStatus(str, Enum): SUCCESSFUL = "successful" FAULTED = "faulted" SUSPENDED = "suspended" + STOPPED = "stopped" class UiPathRuntimeResult(UiPathRuntimeEvent): diff --git a/src/uipath/runtime/stop.py b/src/uipath/runtime/stop.py new file mode 100644 index 00000000..6a820eb4 --- /dev/null +++ b/src/uipath/runtime/stop.py @@ -0,0 +1,28 @@ +"""Per-execution slot through which a host says whether it asked the execution to stop. + +Unset, a cancelled execution ends like any other failure. Set, a cancellation the host asked for +ends the execution as stopped. A contextvar, so concurrent jobs in one process each see their own +host's answer. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import Callable + +# () -> whether the host has asked this execution to stop +StopRequestedProbe = Callable[[], bool] + +_stop_requested_probe: ContextVar[StopRequestedProbe | None] = ContextVar( + "uipath_stop_requested_probe", default=None +) + + +def set_stop_requested_probe(probe: StopRequestedProbe | None) -> None: + """Install the probe (``None`` clears).""" + _stop_requested_probe.set(probe) + + +def get_stop_requested_probe() -> StopRequestedProbe | None: + """The installed probe, or None.""" + return _stop_requested_probe.get() diff --git a/tests/test_stop.py b/tests/test_stop.py new file mode 100644 index 00000000..51f328d9 --- /dev/null +++ b/tests/test_stop.py @@ -0,0 +1,162 @@ +"""A run the host asked to stop ends as stopped rather than faulted.""" + +import asyncio +import json +from pathlib import Path +from typing import Any + +import pytest + +from uipath.runtime.context import UiPathRuntimeContext +from uipath.runtime.output_sinks import set_result_sink +from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus +from uipath.runtime.stop import get_stop_requested_probe, set_stop_requested_probe + + +class _DummyInterceptor: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def setup(self) -> None: + pass + + def teardown(self) -> None: + pass + + +@pytest.fixture(autouse=True) +def _isolate(monkeypatch: pytest.MonkeyPatch) -> Any: + monkeypatch.setattr( + "uipath.runtime.context.UiPathRuntimeLogsInterceptor", _DummyInterceptor + ) + set_stop_requested_probe(None) + set_result_sink(None) + yield + set_stop_requested_probe(None) + set_result_sink(None) + + +def _ctx(tmp_path: Path) -> UiPathRuntimeContext: + return UiPathRuntimeContext( + job_id="job-key", runtime_dir=str(tmp_path / "rt"), result_file="output.json" + ) + + +def _written(ctx: UiPathRuntimeContext) -> dict[str, Any]: + return json.loads(Path(ctx.resolved_result_file_path).read_text()) + + +def _run_cancelled(ctx: UiPathRuntimeContext) -> None: + with pytest.raises(asyncio.CancelledError): + with ctx: + raise asyncio.CancelledError() + + +def test_probe_is_unset_by_default() -> None: + assert get_stop_requested_probe() is None + + +def test_a_cancellation_the_host_asked_for_ends_as_stopped(tmp_path: Path) -> None: + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + _run_cancelled(ctx) + + content = _written(ctx) + assert content["status"] == UiPathRuntimeStatus.STOPPED.value + assert content.get("error") is None + + +def test_a_cancellation_without_a_probe_still_faults(tmp_path: Path) -> None: + ctx = _ctx(tmp_path) + + _run_cancelled(ctx) + + content = _written(ctx) + assert content["status"] == UiPathRuntimeStatus.FAULTED.value + assert content["error"]["code"] == "ERROR_CancelledError" + + +def test_a_cancellation_the_host_did_not_ask_for_faults(tmp_path: Path) -> None: + set_stop_requested_probe(lambda: False) + ctx = _ctx(tmp_path) + + _run_cancelled(ctx) + + assert _written(ctx)["status"] == UiPathRuntimeStatus.FAULTED.value + + +def test_an_error_while_stopping_ends_as_stopped(tmp_path: Path) -> None: + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + with pytest.raises(RuntimeError): + with ctx: + try: + raise asyncio.CancelledError() + except asyncio.CancelledError: + raise RuntimeError("cleanup failed while stopping") from None + + assert _written(ctx)["status"] == UiPathRuntimeStatus.STOPPED.value + + +def test_an_unrelated_error_still_faults_after_a_stop_request(tmp_path: Path) -> None: + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + with pytest.raises(ValueError): + with ctx: + raise ValueError("bad input") + + content = _written(ctx) + assert content["status"] == UiPathRuntimeStatus.FAULTED.value + assert content["error"]["code"] == "ERROR_ValueError" + + +def test_a_run_that_completes_after_a_stop_request_keeps_its_result( + tmp_path: Path, +) -> None: + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + with ctx: + ctx.result = UiPathRuntimeResult(output={"done": True}) + + content = _written(ctx) + assert content["status"] == UiPathRuntimeStatus.SUCCESSFUL.value + assert content["output"] == {"done": True} + + +def test_a_probe_that_raises_reports_the_run_as_faulted(tmp_path: Path) -> None: + def _broken() -> bool: + raise RuntimeError("probe broke") + + set_stop_requested_probe(_broken) + ctx = _ctx(tmp_path) + + _run_cancelled(ctx) + + assert _written(ctx)["status"] == UiPathRuntimeStatus.FAULTED.value + + +def test_the_result_sink_receives_a_stopped_result(tmp_path: Path) -> None: + delivered: list[UiPathRuntimeResult] = [] + set_result_sink(lambda result, _path: delivered.append(result)) + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + _run_cancelled(ctx) + + assert [r.status for r in delivered] == [UiPathRuntimeStatus.STOPPED] + + +def test_the_probe_is_read_when_the_context_starts(tmp_path: Path) -> None: + set_stop_requested_probe(lambda: True) + ctx = _ctx(tmp_path) + + with pytest.raises(asyncio.CancelledError): + with ctx: + set_stop_requested_probe(None) + raise asyncio.CancelledError() + + assert _written(ctx)["status"] == UiPathRuntimeStatus.STOPPED.value diff --git a/uv.lock b/uv.lock index 7106406e..f4a278a1 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.13.5" +version = "0.13.6" source = { editable = "." } dependencies = [ { name = "chardet" }, From 06b82fd80aaee4fdb11b79af0665a04a0920f0e6 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Mon, 5 Oct 2026 14:58:19 +0300 Subject: [PATCH 2/3] fix(debug): stop a run when the debugger quits while it is going The debug runtime looked at the bridge's terminate signal only while it was waiting: before the first resume, at a breakpoint, or while polling a trigger. A Quit from Studio Web while the agent was running was logged and then ignored, and the run went on to completion. The delegate's stream is now watched for a quit. One that arrives while the delegate is working cancels it where it waits, as a stop would; one that arrives while an event is being handled ends the stream before the next. The delegate is closed either way and the run ends with the result a quit has always produced. A cancellation from anywhere else still propagates. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/uipath/runtime/debug/runtime.py | 67 +++++++++++- tests/test_debugger.py | 163 ++++++++++++++++++++++++++++ 2 files changed, 228 insertions(+), 2 deletions(-) diff --git a/src/uipath/runtime/debug/runtime.py b/src/uipath/runtime/debug/runtime.py index 57ef5ec6..4637533f 100644 --- a/src/uipath/runtime/debug/runtime.py +++ b/src/uipath/runtime/debug/runtime.py @@ -60,6 +60,7 @@ def __init__( if trigger_poll_interval < 0: raise ValueError("trigger_poll_interval must be >= 0") self.trigger_poll_interval = trigger_poll_interval + self._quit_requested = False async def execute( self, @@ -120,6 +121,7 @@ async def _stream_and_debug( """Stream events from inner runtime and handle debug interactions.""" final_result: UiPathRuntimeResult execution_completed = False + self._quit_requested = False # Starting in paused state - wait for breakpoints and resume try: @@ -172,8 +174,8 @@ async def _stream_and_debug( debug_options.breakpoints = self.debug_bridge.get_breakpoints() # Stream events from inner runtime - async for event in self.delegate.stream( - current_input, options=debug_options + async for event in self._until_quit( + self.delegate.stream(current_input, options=debug_options) ): yield event @@ -260,6 +262,61 @@ async def _stream_and_debug( elif isinstance(event, UiPathRuntimeStateEvent): await self.debug_bridge.emit_state_update(event) + if self._quit_requested and not execution_completed: + logger.info("Debug session quit by user during execution") + final_result = UiPathRuntimeResult( + status=UiPathRuntimeStatus.SUCCESSFUL, + ) + yield final_result + execution_completed = True + + async def _until_quit( + self, events: AsyncGenerator[UiPathRuntimeEvent, None] + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + """Pass the delegate's events through until the debugger quits. + + A quit while the delegate is working cancels it where it waits, as a stop would; one + between events ends the stream before the next. Either way the delegate is closed and + ``_quit_requested`` is set. A cancellation from anywhere else still propagates. + """ + task = asyncio.current_task() + assert task is not None + terminate = asyncio.ensure_future(self.debug_bridge.wait_for_terminate()) + in_delegate = False + cancelled = False + + def on_terminate(done: "asyncio.Future[None]") -> None: + nonlocal cancelled + if in_delegate and _requested(done): + cancelled = True + task.cancel() + + terminate.add_done_callback(on_terminate) + try: + while not _requested(terminate): + in_delegate = True + try: + event = await anext(events) + except StopAsyncIteration: + return + except asyncio.CancelledError: + if cancelled: + cancelled = False + if task.uncancel() == 0: + break + raise + finally: + in_delegate = False + yield event + self._quit_requested = True + finally: + terminate.remove_done_callback(on_terminate) + if cancelled: + # The delegate swallowed the cancellation this sent it. + task.uncancel() + terminate.cancel() + await events.aclose() + def get_resumable_runtime( self, max_depth: int = 10 ) -> UiPathResumableRuntime | None: @@ -358,3 +415,9 @@ async def _wait_with_quit_check(self) -> None: if term_task in done: raise UiPathDebugQuitError("Debugging terminated during polling.") + + +def _requested(terminate: "asyncio.Future[None]") -> bool: + return ( + terminate.done() and not terminate.cancelled() and terminate.exception() is None + ) diff --git a/tests/test_debugger.py b/tests/test_debugger.py index 0c613a16..0c4150cb 100644 --- a/tests/test_debugger.py +++ b/tests/test_debugger.py @@ -27,6 +27,10 @@ from uipath.runtime.schema import UiPathRuntimeSchema +async def _never() -> None: + await asyncio.Event().wait() + + def make_debug_bridge_mock() -> UiPathDebugProtocol: """Create a debug bridge mock with all methods that UiPathDebugRuntime uses. @@ -43,6 +47,8 @@ def make_debug_bridge_mock() -> UiPathDebugProtocol: bridge_mock.emit_breakpoint_hit = AsyncMock() bridge_mock.emit_state_update = AsyncMock() bridge_mock.wait_for_resume = AsyncMock() + # The bridge returns from wait_for_terminate only once the user quits. + bridge_mock.wait_for_terminate = AsyncMock(side_effect=_never) bridge_mock.get_breakpoints = Mock(return_value=["node-1"]) @@ -475,3 +481,160 @@ async def test_debug_runtime_dispose_suppresses_disconnect_errors(): await debug_runtime.dispose() cast(AsyncMock, bridge.disconnect).assert_awaited_once() + + +class _WorkingRuntime: + """A delegate that emits one state event, then works until it is cancelled.""" + + def __init__(self) -> None: + self.cleaned_up = False + self.emitted: list[str] = [] + + async def dispose(self) -> None: + pass + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + raise AssertionError("streaming is used") + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + try: + self.emitted.append("first") + yield UiPathRuntimeStateEvent(node_name="first", payload={}) + await asyncio.sleep(30) + self.emitted.append("second") + yield UiPathRuntimeStateEvent(node_name="second", payload={}) + yield UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + finally: + self.cleaned_up = True + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + +def _bridge_with_quit() -> tuple[UiPathDebugProtocol, asyncio.Event]: + bridge = make_debug_bridge_mock() + quit_event = asyncio.Event() + cast(Mock, bridge).wait_for_terminate = AsyncMock(side_effect=quit_event.wait) + cast(Mock, bridge.get_breakpoints).return_value = [] + return bridge, quit_event + + +@pytest.mark.asyncio +async def test_debug_runtime_quit_while_running_cancels_the_delegate(): + """A quit that arrives while the delegate works stops it instead of letting it finish.""" + runtime_impl = _WorkingRuntime() + bridge, quit_event = _bridge_with_quit() + debug_runtime = UiPathDebugRuntime(delegate=runtime_impl, debug_bridge=bridge) + + async def quit_soon() -> None: + await asyncio.sleep(0.05) + quit_event.set() + + quitter = asyncio.create_task(quit_soon()) + result = await asyncio.wait_for(debug_runtime.execute({}), timeout=5) + await quitter + + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + assert runtime_impl.emitted == ["first"] + assert runtime_impl.cleaned_up + cast(AsyncMock, bridge.emit_execution_completed).assert_awaited_once_with(result) + + +@pytest.mark.asyncio +async def test_debug_runtime_quit_between_events_ends_before_the_next_one(): + """A quit that arrives while an event is being handled ends the stream before the next.""" + runtime_impl = _WorkingRuntime() + bridge, quit_event = _bridge_with_quit() + + async def quit_on_first_state(_event: Any) -> None: + quit_event.set() + await asyncio.sleep(0) + + cast(AsyncMock, bridge.emit_state_update).side_effect = quit_on_first_state + debug_runtime = UiPathDebugRuntime(delegate=runtime_impl, debug_bridge=bridge) + + result = await asyncio.wait_for(debug_runtime.execute({}), timeout=5) + + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + assert runtime_impl.emitted == ["first"] + assert runtime_impl.cleaned_up + + +@pytest.mark.asyncio +async def test_debug_runtime_cancellation_from_elsewhere_still_propagates(): + """Only a quit is turned into a result; a host's cancellation keeps cancelling.""" + runtime_impl = _WorkingRuntime() + bridge, _quit_event = _bridge_with_quit() + debug_runtime = UiPathDebugRuntime(delegate=runtime_impl, debug_bridge=bridge) + + run = asyncio.create_task(debug_runtime.execute({})) + await asyncio.sleep(0.05) + run.cancel() + + with pytest.raises(asyncio.CancelledError): + await run + assert runtime_impl.cleaned_up + cast(AsyncMock, bridge.emit_execution_completed).assert_not_awaited() + + +class _BreakpointRuntime: + """A delegate that pauses at a breakpoint: its stream ends with the breakpoint result.""" + + async def dispose(self) -> None: + pass + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + raise AssertionError("streaming is used") + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + yield UiPathRuntimeStateEvent(node_name="node-quit", payload={}) + yield UiPathBreakpointResult( + breakpoint_node="node-quit", + breakpoint_type="before", + next_nodes=[], + current_state={}, + ) + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + +@pytest.mark.asyncio +async def test_debug_runtime_quit_at_a_breakpoint_yields_one_result(): + """A quit answered at a breakpoint is not reported a second time by the stream watch.""" + bridge, quit_event = _bridge_with_quit() + + async def resume_then_quit() -> None: + if cast(AsyncMock, bridge.wait_for_resume).await_count > 1: + quit_event.set() + raise UiPathDebugQuitError("quit") + + cast(AsyncMock, bridge.wait_for_resume).side_effect = resume_then_quit + debug_runtime = UiPathDebugRuntime( + delegate=_BreakpointRuntime(), debug_bridge=bridge + ) + + results = [ + event + async for event in debug_runtime.stream({}) + if isinstance(event, UiPathRuntimeResult) + and not isinstance(event, UiPathBreakpointResult) + ] + + assert [r.status for r in results] == [UiPathRuntimeStatus.SUCCESSFUL] From 4477dc9337e12e8a9a6d2c1728db1a93e6ddf6f0 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 6 Oct 2026 14:35:02 +0300 Subject: [PATCH 3/3] refactor(runtime): detect a stop from the task, and watch for a quit in execute() A stopped run is now told apart from a failure by Task.cancelling() on the task that entered the context: a cancel() someone called keeps it above zero through cleanup, and a stray CancelledError leaves it at zero. This replaces the host-installed stop-requested probe and the walk of the exception chain, so uipath.runtime.stop goes away and hosts install nothing. The debug quit is watched in execute(), which runs the stream in its own task and races it against wait_for_terminate(). A quit cancels that task instead of the caller's, so no uncancel() bookkeeping or in-delegate tracking is needed, and the execute() fallback for a delegate that cannot stream is covered too. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_011be41L1VE33oixHNnJRPhw --- src/uipath/runtime/context.py | 34 +++------ src/uipath/runtime/debug/runtime.py | 104 ++++++++++----------------- src/uipath/runtime/stop.py | 28 -------- tests/test_debugger.py | 43 +++++++++++ tests/test_stop.py | 107 +++++++++++----------------- 5 files changed, 130 insertions(+), 186 deletions(-) delete mode 100644 src/uipath/runtime/stop.py diff --git a/src/uipath/runtime/context.py b/src/uipath/runtime/context.py index 9290d942..58e02cb1 100644 --- a/src/uipath/runtime/context.py +++ b/src/uipath/runtime/context.py @@ -22,7 +22,6 @@ from uipath.runtime.logging._interceptor import UiPathRuntimeLogsInterceptor from uipath.runtime.output_sinks import ResultSink, get_log_handler, get_result_sink from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus -from uipath.runtime.stop import StopRequestedProbe, get_stop_requested_probe logger = logging.getLogger(__name__) @@ -129,7 +128,6 @@ class UiPathRuntimeContext(BaseModel): # installed when the context started is the one that receives the result — even if the registry is # changed in between or __exit__ runs on a different event loop. _result_sink: ResultSink | None = PrivateAttr(default=None) - _stop_requested_probe: StopRequestedProbe | None = PrivateAttr(default=None) def _apply_execution_source(self) -> None: """Derive execution_source from the command, if not already set. @@ -250,7 +248,6 @@ def __enter__(self): # the result sink is stashed for __exit__ so the two are read at the same moment (same context). log_handler = get_log_handler() self._result_sink = get_result_sink() - self._stop_requested_probe = get_stop_requested_probe() self.logs_interceptor = UiPathRuntimeLogsInterceptor( min_level=self.logs_min_level, @@ -275,7 +272,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self.result is None: self.result = UiPathRuntimeResult() - if exc_type and exc_val and self._stopped_on_request(exc_val): + if exc_type and exc_val and _task_is_being_cancelled(): self.result.status = UiPathRuntimeStatus.STOPPED self.result.error = None elif exc_type and exc_val: @@ -385,19 +382,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): if hasattr(self, "logs_interceptor"): self.logs_interceptor.teardown() - def _stopped_on_request(self, exc: BaseException) -> bool: - """Whether the run ended by a cancellation the host asked for.""" - probe = self._stop_requested_probe - if probe is None or not _caused_by_cancellation(exc): - return False - try: - return bool(probe()) - except Exception: - logger.exception( - "Stop-requested probe failed; the run is reported as faulted" - ) - return False - def _deliver_result(self, sink: ResultSink, output_payload: Any) -> None: """Spill the output arguments to a file, then hand the result + that path to the sink.""" args_path = self.resolved_output_arguments_file_path @@ -550,12 +534,10 @@ def from_config( return instance -def _caused_by_cancellation(exc: BaseException | None) -> bool: - """Whether ``exc`` is a cancellation, or was raised while handling or because of one.""" - seen: set[int] = set() - while exc is not None and id(exc) not in seen: - if isinstance(exc, asyncio.CancelledError): - return True - seen.add(id(exc)) - exc = exc.__cause__ or exc.__context__ - return False +def _task_is_being_cancelled() -> bool: + """Whether someone called ``cancel()`` on the running task, as opposed to a stray CancelledError.""" + try: + task = asyncio.current_task() + except RuntimeError: + return False + return task is not None and task.cancelling() > 0 diff --git a/src/uipath/runtime/debug/runtime.py b/src/uipath/runtime/debug/runtime.py index 4637533f..b2308e05 100644 --- a/src/uipath/runtime/debug/runtime.py +++ b/src/uipath/runtime/debug/runtime.py @@ -60,16 +60,48 @@ def __init__( if trigger_poll_interval < 0: raise ValueError("trigger_poll_interval must be >= 0") self.trigger_poll_interval = trigger_poll_interval - self._quit_requested = False async def execute( self, input: dict[str, Any] | None = None, options: UiPathExecuteOptions | None = None, ) -> UiPathRuntimeResult: - """Execute the workflow with debug support.""" + """Execute the workflow with debug support; a quit while it runs stops it.""" + run = asyncio.ensure_future( + self._run(input, cast(UiPathStreamOptions, options)) + ) + terminate = asyncio.ensure_future(self.debug_bridge.wait_for_terminate()) + try: + done, _ = await asyncio.wait( + {run, terminate}, return_when=asyncio.FIRST_COMPLETED + ) + if run not in done and not _quit_requested(terminate): + done, _ = await asyncio.wait({run}) + finally: + terminate.cancel() + if not run.done(): + run.cancel() + await asyncio.wait({run}) + + if run in done: + return run.result() + + if not run.cancelled() and run.exception(): + logger.warning( + "Delegate failed while stopping on quit", exc_info=run.exception() + ) + logger.info("Debug session quit by user during execution") + result = UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + await self.debug_bridge.emit_execution_completed(result) + return result + + async def _run( + self, + input: dict[str, Any] | None, + options: UiPathStreamOptions | None, + ) -> UiPathRuntimeResult: final_result = None - async for event in self.stream(input, cast(UiPathStreamOptions, options)): + async for event in self.stream(input, options): if isinstance(event, UiPathRuntimeResult): final_result = event @@ -121,7 +153,6 @@ async def _stream_and_debug( """Stream events from inner runtime and handle debug interactions.""" final_result: UiPathRuntimeResult execution_completed = False - self._quit_requested = False # Starting in paused state - wait for breakpoints and resume try: @@ -174,8 +205,8 @@ async def _stream_and_debug( debug_options.breakpoints = self.debug_bridge.get_breakpoints() # Stream events from inner runtime - async for event in self._until_quit( - self.delegate.stream(current_input, options=debug_options) + async for event in self.delegate.stream( + current_input, options=debug_options ): yield event @@ -262,61 +293,6 @@ async def _stream_and_debug( elif isinstance(event, UiPathRuntimeStateEvent): await self.debug_bridge.emit_state_update(event) - if self._quit_requested and not execution_completed: - logger.info("Debug session quit by user during execution") - final_result = UiPathRuntimeResult( - status=UiPathRuntimeStatus.SUCCESSFUL, - ) - yield final_result - execution_completed = True - - async def _until_quit( - self, events: AsyncGenerator[UiPathRuntimeEvent, None] - ) -> AsyncGenerator[UiPathRuntimeEvent, None]: - """Pass the delegate's events through until the debugger quits. - - A quit while the delegate is working cancels it where it waits, as a stop would; one - between events ends the stream before the next. Either way the delegate is closed and - ``_quit_requested`` is set. A cancellation from anywhere else still propagates. - """ - task = asyncio.current_task() - assert task is not None - terminate = asyncio.ensure_future(self.debug_bridge.wait_for_terminate()) - in_delegate = False - cancelled = False - - def on_terminate(done: "asyncio.Future[None]") -> None: - nonlocal cancelled - if in_delegate and _requested(done): - cancelled = True - task.cancel() - - terminate.add_done_callback(on_terminate) - try: - while not _requested(terminate): - in_delegate = True - try: - event = await anext(events) - except StopAsyncIteration: - return - except asyncio.CancelledError: - if cancelled: - cancelled = False - if task.uncancel() == 0: - break - raise - finally: - in_delegate = False - yield event - self._quit_requested = True - finally: - terminate.remove_done_callback(on_terminate) - if cancelled: - # The delegate swallowed the cancellation this sent it. - task.uncancel() - terminate.cancel() - await events.aclose() - def get_resumable_runtime( self, max_depth: int = 10 ) -> UiPathResumableRuntime | None: @@ -417,7 +393,5 @@ async def _wait_with_quit_check(self) -> None: raise UiPathDebugQuitError("Debugging terminated during polling.") -def _requested(terminate: "asyncio.Future[None]") -> bool: - return ( - terminate.done() and not terminate.cancelled() and terminate.exception() is None - ) +def _quit_requested(terminate: "asyncio.Future[None]") -> bool: + return not terminate.cancelled() and terminate.exception() is None diff --git a/src/uipath/runtime/stop.py b/src/uipath/runtime/stop.py deleted file mode 100644 index 6a820eb4..00000000 --- a/src/uipath/runtime/stop.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Per-execution slot through which a host says whether it asked the execution to stop. - -Unset, a cancelled execution ends like any other failure. Set, a cancellation the host asked for -ends the execution as stopped. A contextvar, so concurrent jobs in one process each see their own -host's answer. -""" - -from __future__ import annotations - -from contextvars import ContextVar -from typing import Callable - -# () -> whether the host has asked this execution to stop -StopRequestedProbe = Callable[[], bool] - -_stop_requested_probe: ContextVar[StopRequestedProbe | None] = ContextVar( - "uipath_stop_requested_probe", default=None -) - - -def set_stop_requested_probe(probe: StopRequestedProbe | None) -> None: - """Install the probe (``None`` clears).""" - _stop_requested_probe.set(probe) - - -def get_stop_requested_probe() -> StopRequestedProbe | None: - """The installed probe, or None.""" - return _stop_requested_probe.get() diff --git a/tests/test_debugger.py b/tests/test_debugger.py index 0c4150cb..a88689fb 100644 --- a/tests/test_debugger.py +++ b/tests/test_debugger.py @@ -638,3 +638,46 @@ async def resume_then_quit() -> None: ] assert [r.status for r in results] == [UiPathRuntimeStatus.SUCCESSFUL] + + +class _NonStreamingWorkingRuntime(_WorkingRuntime): + """A delegate without streaming that works until it is cancelled.""" + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + try: + await asyncio.sleep(30) + return UiPathRuntimeResult(status=UiPathRuntimeStatus.FAULTED) + finally: + self.cleaned_up = True + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + raise UiPathStreamNotSupportedError("no streaming") + yield + + +@pytest.mark.asyncio +async def test_debug_runtime_quit_stops_a_delegate_that_cannot_stream(): + """The quit also reaches a delegate run through the execute() fallback.""" + runtime_impl = _NonStreamingWorkingRuntime() + bridge, quit_event = _bridge_with_quit() + debug_runtime = UiPathDebugRuntime(delegate=runtime_impl, debug_bridge=bridge) + + async def quit_soon() -> None: + await asyncio.sleep(0.05) + quit_event.set() + + quitter = asyncio.create_task(quit_soon()) + result = await asyncio.wait_for(debug_runtime.execute({}), timeout=5) + await quitter + + assert result.status == UiPathRuntimeStatus.SUCCESSFUL + assert runtime_impl.cleaned_up + cast(AsyncMock, bridge.emit_execution_completed).assert_awaited_once_with(result) diff --git a/tests/test_stop.py b/tests/test_stop.py index 51f328d9..43b530ee 100644 --- a/tests/test_stop.py +++ b/tests/test_stop.py @@ -1,4 +1,4 @@ -"""A run the host asked to stop ends as stopped rather than faulted.""" +"""A run whose task is cancelled ends as stopped rather than faulted.""" import asyncio import json @@ -10,7 +10,6 @@ from uipath.runtime.context import UiPathRuntimeContext from uipath.runtime.output_sinks import set_result_sink from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus -from uipath.runtime.stop import get_stop_requested_probe, set_stop_requested_probe class _DummyInterceptor: @@ -29,10 +28,8 @@ def _isolate(monkeypatch: pytest.MonkeyPatch) -> Any: monkeypatch.setattr( "uipath.runtime.context.UiPathRuntimeLogsInterceptor", _DummyInterceptor ) - set_stop_requested_probe(None) set_result_sink(None) yield - set_stop_requested_probe(None) set_result_sink(None) @@ -46,62 +43,72 @@ def _written(ctx: UiPathRuntimeContext) -> dict[str, Any]: return json.loads(Path(ctx.resolved_result_file_path).read_text()) -def _run_cancelled(ctx: UiPathRuntimeContext) -> None: - with pytest.raises(asyncio.CancelledError): - with ctx: - raise asyncio.CancelledError() - - -def test_probe_is_unset_by_default() -> None: - assert get_stop_requested_probe() is None +async def _cancel_while_running(body: Any) -> None: + task = asyncio.create_task(body()) + await asyncio.sleep(0.01) + task.cancel() + await asyncio.wait({task}) -def test_a_cancellation_the_host_asked_for_ends_as_stopped(tmp_path: Path) -> None: - set_stop_requested_probe(lambda: True) +@pytest.mark.asyncio +async def test_a_cancelled_run_ends_as_stopped(tmp_path: Path) -> None: ctx = _ctx(tmp_path) - _run_cancelled(ctx) + async def body() -> None: + with ctx: + await asyncio.sleep(10) + + await _cancel_while_running(body) content = _written(ctx) assert content["status"] == UiPathRuntimeStatus.STOPPED.value assert content.get("error") is None -def test_a_cancellation_without_a_probe_still_faults(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_a_stray_cancelled_error_still_faults(tmp_path: Path) -> None: ctx = _ctx(tmp_path) - _run_cancelled(ctx) + with pytest.raises(asyncio.CancelledError): + with ctx: + child = asyncio.ensure_future(asyncio.sleep(10)) + await asyncio.sleep(0) + child.cancel() + await child content = _written(ctx) assert content["status"] == UiPathRuntimeStatus.FAULTED.value assert content["error"]["code"] == "ERROR_CancelledError" -def test_a_cancellation_the_host_did_not_ask_for_faults(tmp_path: Path) -> None: - set_stop_requested_probe(lambda: False) +def test_a_cancelled_error_outside_an_event_loop_faults(tmp_path: Path) -> None: ctx = _ctx(tmp_path) - _run_cancelled(ctx) + with pytest.raises(asyncio.CancelledError): + with ctx: + raise asyncio.CancelledError() assert _written(ctx)["status"] == UiPathRuntimeStatus.FAULTED.value -def test_an_error_while_stopping_ends_as_stopped(tmp_path: Path) -> None: - set_stop_requested_probe(lambda: True) +@pytest.mark.asyncio +async def test_an_error_while_stopping_ends_as_stopped(tmp_path: Path) -> None: ctx = _ctx(tmp_path) - with pytest.raises(RuntimeError): + async def body() -> None: with ctx: try: - raise asyncio.CancelledError() + await asyncio.sleep(10) except asyncio.CancelledError: raise RuntimeError("cleanup failed while stopping") from None + await _cancel_while_running(body) + assert _written(ctx)["status"] == UiPathRuntimeStatus.STOPPED.value -def test_an_unrelated_error_still_faults_after_a_stop_request(tmp_path: Path) -> None: - set_stop_requested_probe(lambda: True) +@pytest.mark.asyncio +async def test_an_error_in_a_run_nobody_cancelled_faults(tmp_path: Path) -> None: ctx = _ctx(tmp_path) with pytest.raises(ValueError): @@ -113,50 +120,16 @@ def test_an_unrelated_error_still_faults_after_a_stop_request(tmp_path: Path) -> assert content["error"]["code"] == "ERROR_ValueError" -def test_a_run_that_completes_after_a_stop_request_keeps_its_result( - tmp_path: Path, -) -> None: - set_stop_requested_probe(lambda: True) - ctx = _ctx(tmp_path) - - with ctx: - ctx.result = UiPathRuntimeResult(output={"done": True}) - - content = _written(ctx) - assert content["status"] == UiPathRuntimeStatus.SUCCESSFUL.value - assert content["output"] == {"done": True} - - -def test_a_probe_that_raises_reports_the_run_as_faulted(tmp_path: Path) -> None: - def _broken() -> bool: - raise RuntimeError("probe broke") - - set_stop_requested_probe(_broken) - ctx = _ctx(tmp_path) - - _run_cancelled(ctx) - - assert _written(ctx)["status"] == UiPathRuntimeStatus.FAULTED.value - - -def test_the_result_sink_receives_a_stopped_result(tmp_path: Path) -> None: +@pytest.mark.asyncio +async def test_the_result_sink_receives_a_stopped_result(tmp_path: Path) -> None: delivered: list[UiPathRuntimeResult] = [] set_result_sink(lambda result, _path: delivered.append(result)) - set_stop_requested_probe(lambda: True) - ctx = _ctx(tmp_path) - - _run_cancelled(ctx) - - assert [r.status for r in delivered] == [UiPathRuntimeStatus.STOPPED] - - -def test_the_probe_is_read_when_the_context_starts(tmp_path: Path) -> None: - set_stop_requested_probe(lambda: True) ctx = _ctx(tmp_path) - with pytest.raises(asyncio.CancelledError): + async def body() -> None: with ctx: - set_stop_requested_probe(None) - raise asyncio.CancelledError() + await asyncio.sleep(10) - assert _written(ctx)["status"] == UiPathRuntimeStatus.STOPPED.value + await _cancel_while_running(body) + + assert [r.status for r in delivered] == [UiPathRuntimeStatus.STOPPED]