From 8b69e040a726366f634ced86ee0d7dd476fd0e45 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 13:49:26 +0300 Subject: [PATCH 1/3] feat(cli): really stop a job that is already executing Builds on async dispatch. Until now StopJob could only refuse a running job: its body runs on a thread via asyncio.to_thread, which cannot be cancelled. That reasoning missed a fact. run/debug/eval each drive their own event loop with asyncio.run INSIDE the worker thread (cli_run.py:333, cli_debug.py:269, cli_eval.py:531), so a job IS an event loop -- and an event loop can be cancelled. The three call sites now go through `run_job_loop`, which publishes that loop and its root task to a JobControl carried on a ContextVar. asyncio.to_thread propagates contextvars, so no monkeypatching is needed, and outside the server (uipath run on a terminal) it is asyncio.run verbatim. Cancellation is cooperative, not a thread kill: the runtime's context managers still unwind, so UiPathRuntimeContext.__exit__ still writes output.json and the caller's file fallback still works. The stop ladder: 1. cancel the ROOT task only, then wait STOP_GRACE_SECONDS. Cancelling every task would land a second CancelledError inside the cleanup that writes output.json and abort it mid-write. 2. if cleanup itself is stuck, sweep the loop and wait a shorter window. 3. otherwise return False -- a job wedged in a non-cancellable C call (a socket read inside an LLM request) cannot be stopped, and saying so is better than claiming a stop that did not happen. A stopped job reports Stopped, not Faulted. The runtime records it as FAULTED/ERROR_CancelledError because it sees a CancelledError, which is the wrong story for a stop the caller asked for, so the result push carries an explicit `stopped` flag that wins over the document. Stop is also reachable now: POST /jobs/{job_key}/stop on the HTTP transport, which carries all current traffic and previously had no way to reach the registry at all. Notes for review: - _invoke_command discriminates the two CancelledErrors via Task.cancelling(): 0 means the job's own loop was cancelled (an outcome, swallow it); >0 means our awaiting task was cancelled (a shutdown, re-raise). - StopJob takes resume_version as a trailing optional parameter rather than a DTO: uipath-ipc ignores a surplus wire arg and defaults a missing one, so old and new peers interoperate both ways. A suspended job resumes under the same key, so a stop aimed at the previous run must not kill the resumed one. Co-Authored-By: Claude Opus 5 --- .../uipath/src/uipath/_cli/_job_control.py | 113 +++++++++ .../uipath/src/uipath/_cli/_server_core.py | 44 +++- .../uipath/src/uipath/_cli/_server_jobs.py | 108 ++++++++- packages/uipath/src/uipath/_cli/cli_debug.py | 4 +- packages/uipath/src/uipath/_cli/cli_eval.py | 3 +- packages/uipath/src/uipath/_cli/cli_run.py | 4 +- packages/uipath/src/uipath/_cli/cli_server.py | 32 +++ .../uipath/src/uipath/_cli/cli_server_ipc.py | 14 +- .../uipath/tests/cli/test_server_async.py | 185 ++++++++++++++- .../uipath/tests/cli/test_server_callbacks.py | 1 + .../tests/cli/test_server_cancellation.py | 222 ++++++++++++++++++ packages/uipath/tests/cli/test_server_ipc.py | 2 +- 12 files changed, 711 insertions(+), 21 deletions(-) create mode 100644 packages/uipath/src/uipath/_cli/_job_control.py create mode 100644 packages/uipath/tests/cli/test_server_cancellation.py diff --git a/packages/uipath/src/uipath/_cli/_job_control.py b/packages/uipath/src/uipath/_cli/_job_control.py new file mode 100644 index 000000000..f39df2e76 --- /dev/null +++ b/packages/uipath/src/uipath/_cli/_job_control.py @@ -0,0 +1,113 @@ +"""The seam that lets the server cancel the job running on its worker thread. + +Every server command drives its own event loop with ``asyncio.run`` inside the +``asyncio.to_thread`` worker. That loop is the only place a cancellation can actually +land, so the command publishes it here and the server reaches back through +``loop.call_soon_threadsafe``. + +Deliberately import-light so ``_cli/__init__.py``'s lazy-import discipline is untouched. +""" + +import asyncio +import contextvars +import threading +from typing import Any + +CURRENT_JOB_CONTROL: contextvars.ContextVar["JobControl | None"] = ( + contextvars.ContextVar("uipath_current_job_control", default=None) +) + + +class JobControl: + """Handle on one job's event loop, shared between the server loop and the worker.""" + + def __init__(self, job_key: str) -> None: + self.job_key = job_key + self.cancel_requested = False + self._loop: asyncio.AbstractEventLoop | None = None + self._task: "asyncio.Task[Any] | None" = None + # Bound from the job thread, read from the server loop. + self._sync = threading.Lock() + + # ---- job thread --------------------------------------------------------- + + def bind(self, loop: asyncio.AbstractEventLoop, task: "asyncio.Task[Any]") -> None: + with self._sync: + self._loop, self._task = loop, task + pending = self.cancel_requested + if pending: + # A stop that raced the loop's creation: apply it now rather than losing it. + self.cancel() + + def unbind(self) -> None: + with self._sync: + self._loop = self._task = None + + # ---- server loop -------------------------------------------------------- + + @property + def bound(self) -> bool: + with self._sync: + return self._task is not None + + def cancel(self) -> bool: + """Deliver CancelledError to the job's ROOT task only. + + Call this at most once. A second cancel lands inside the runtime's cleanup + ``finally`` blocks — the ones that write ``output.json`` and tear the log + interceptor down — and aborts them, which would destroy the very fallback the + caller relies on. Escalate with ``cancel_all`` instead. + """ + with self._sync: + self.cancel_requested = True + loop, task = self._loop, self._task + if loop is None or task is None: + # Not bound yet; bind() will apply it. + return False + try: + loop.call_soon_threadsafe(task.cancel) + except RuntimeError: + # Loop already closed — the job is finishing anyway. + return False + return True + + def cancel_all(self) -> bool: + """Escalation: cancel every task on the job's loop. + + This includes whatever the cleanup is awaiting, so it gives up on a clean + ``output.json``. Only worth doing once a polite cancel has already failed. + """ + with self._sync: + loop = self._loop + if loop is None: + return False + + def _sweep() -> None: + for task in asyncio.all_tasks(loop): + task.cancel() + + try: + loop.call_soon_threadsafe(_sweep) + except RuntimeError: + return False + return True + + +def run_job_loop(coro: Any) -> Any: + """``asyncio.run`` that publishes its loop and root task to the JobControl in scope. + + Outside the server — ``uipath run`` on a terminal — there is no handle in scope and + this is ``asyncio.run`` verbatim. + """ + control = CURRENT_JOB_CONTROL.get() + if control is None: + return asyncio.run(coro) + + with asyncio.Runner() as runner: + loop = runner.get_loop() + task = loop.create_task(coro, context=contextvars.copy_context()) + control.bind(loop, task) + try: + return loop.run_until_complete(task) + finally: + control.unbind() diff --git a/packages/uipath/src/uipath/_cli/_server_core.py b/packages/uipath/src/uipath/_cli/_server_core.py index 6ff6952f0..750f263cd 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/packages/uipath/src/uipath/_cli/_server_core.py @@ -4,8 +4,10 @@ import json import os import shlex +from collections.abc import Callable from typing import Any +from ._job_control import CURRENT_JOB_CONTROL, JobControl from .cli_debug import debug from .cli_eval import eval from .cli_run import run @@ -55,6 +57,10 @@ def parse_args(args: str | list[str] | None) -> list[str]: DEFAULT_RESULT_FILE = "output.json" DEFAULT_LOGS_FILE = "execution.log" +# Distinct from any click exit code, so the caller can tell "stopped on request" from +# "the job failed" and report Stopped rather than Faulted. +EXIT_CODE_STOPPED = 143 + def _resolve_runtime_file( config_path: str, base_dir: str, key: str, default_name: str @@ -134,8 +140,13 @@ def _read_result_document() -> tuple[str | None, str]: return None, "file" -async def _invoke_command(cmd: Any, args: list[str]) -> dict[str, Any]: +async def _invoke_command( + cmd: Any, args: list[str], control: "JobControl | None" = None +) -> dict[str, Any]: """Invoke one click command and classify how it ended.""" + # asyncio.to_thread propagates contextvars, so the command reads this inside the + # worker and publishes its event loop back through it. + token = CURRENT_JOB_CONTROL.set(control) if control is not None else None try: result_value = await asyncio.to_thread(cmd.main, args, standalone_mode=False) # Under standalone_mode=False click RETURNS ctx.exit(N)'s code instead of @@ -163,8 +174,27 @@ async def _invoke_command(cmd: Any, args: list[str]) -> dict[str, Any]: "Result": None, "Unexpected": False, } + except asyncio.CancelledError: + # Two very different things arrive here. cancelling() > 0 means OUR awaiting task + # was cancelled (server shutdown) — never swallow that, or the shutdown stalls. + # 0 means the CancelledError travelled out of the job's own loop, i.e. our StopJob + # landed: that is an OUTCOME, and swallowing it keeps this task alive so the lock + # unwinds and the result document still gets read. + current = asyncio.current_task() + if current is not None and current.cancelling() > 0: + raise + return { + "ExitCode": EXIT_CODE_STOPPED, + "Error": "Job stopped on request", + "Result": None, + "Unexpected": False, + "Stopped": True, + } except Exception as e: # report any job failure as a result, not a fault return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True} + finally: + if token is not None: + CURRENT_JOB_CONTROL.reset(token) async def _run_command_isolated( @@ -172,12 +202,20 @@ async def _run_command_isolated( args: list[str], env_vars: dict[str, str], working_dir: str | None, + on_started: Callable[[], None] | None = None, + control: "JobControl | None" = None, ) -> dict[str, Any]: - """Run one command with per-job env/cwd isolation (the shared job core).""" + """Run one command with per-job env/cwd isolation (the shared job core). + + ``on_started`` fires once the lock is held, i.e. the moment the job stops being + queued and becomes uncancellable. + """ if _state.lock is None or _state.baseline_env is None: raise RuntimeError("Server state not initialized") async with _state.lock: + if on_started is not None: + on_started() original_cwd = os.getcwd() try: # Start from server baseline + request env vars only, so nothing from @@ -201,7 +239,7 @@ async def _run_command_isolated( "ClientError": True, } - outcome = await _invoke_command(cmd, args) + outcome = await _invoke_command(cmd, args, control) # Must happen before the finally below restores env/cwd: the document's # location comes from this job's UIPATH_CONFIG_PATH and may be relative. document, conveyance = _read_result_document() diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index 65fce17d1..edf393512 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -20,6 +20,7 @@ from aiohttp import ClientSession, ClientTimeout, UnixConnector +from ._job_control import JobControl from ._server_core import _run_command_isolated, resolve_logs_file_path from ._utils._console import ConsoleLogger @@ -179,12 +180,21 @@ def build_result_payload(job_key: str, outcome: dict[str, Any]) -> dict[str, Any "stateConveyance": "file", "jobConveyance": outcome.get("DocumentConveyance", "file"), "job": outcome.get("Document"), + # Explicit so the caller reports Stopped rather than Faulted — the runtime's + # output.json records a cancelled job as FAULTED/ERROR_CancelledError, which is + # the wrong story for a stop the caller itself asked for. + "stopped": bool(outcome.get("Stopped")), } LOG_POLL_SECONDS = 0.25 LOG_BATCH_MAX_LINES = 200 LOG_FLUSH_TIMEOUT_SECONDS = 10 + +# How long a cancelled job gets to unwind before we admit the stop did not take. +STOP_GRACE_SECONDS = 30 +# Extra window after escalating to a full loop sweep. +STOP_ESCALATION_SECONDS = 10 # ``[2026-07-29 17:04:19,123][INFO] message`` — the format the runtime's file handler # emits and that the .NET FileLogsWatcher has always parsed. LOG_LINE_RE = re.compile( @@ -307,6 +317,13 @@ class JobRegistry: def __init__(self) -> None: self._tasks: dict[str, asyncio.Task[None]] = {} + # Jobs past the point of no return: executing on a thread, uncancellable. + self._running: set[str] = set() + # A suspended job resumes under the SAME key, so a stop meant for the previous + # run must not kill the resumed one. + self._resume_versions: dict[str, int | None] = {} + # Handle on each job's own event loop, so a stop can cancel it cooperatively. + self._controls: dict[str, JobControl] = {} def is_active(self, job_key: str) -> bool: task = self._tasks.get(job_key) @@ -320,11 +337,14 @@ def start( env_vars: dict[str, str], working_dir: str | None, callback: JobReporter, + resume_version: int | None = None, ) -> bool: """Register and schedule a job. False if one is already in flight for this key.""" if self.is_active(job_key): return False + self._resume_versions[job_key] = resume_version + task = asyncio.create_task( self._run(job_key, cmd, args, env_vars, working_dir, callback) ) @@ -334,6 +354,9 @@ def start( def _forget(self, job_key: str) -> None: self._tasks.pop(job_key, None) + self._running.discard(job_key) + self._resume_versions.pop(job_key, None) + self._controls.pop(job_key, None) async def _run( self, @@ -371,11 +394,15 @@ async def finish(outcome: dict[str, Any]) -> None: tailer = JobLogTailer(job_key, logs_path, callback) tail_task = asyncio.create_task(tailer.run()) + control = JobControl(job_key) + self._controls[job_key] = control outcome = await _run_command_isolated( cmd, args, env_vars, working_dir, + on_started=lambda: self._running.add(job_key), + control=control, ) except asyncio.CancelledError: # Cancelled before it took the lock — report it rather than going silent, @@ -388,18 +415,89 @@ async def finish(outcome: dict[str, Any]) -> None: await finish(outcome) - async def stop(self, job_key: str) -> bool: + async def stop(self, job_key: str, resume_version: int | None = None) -> bool: """Cancel a job that has not started executing yet. - Once the job is running its body is on a thread via ``asyncio.to_thread``, which - cannot be cancelled — so this only removes work that is still queued. Actually - interrupting a running job is a separate change. + Only queued work can be stopped. Once the job holds the lock its body is on a + thread via ``asyncio.to_thread``, which cannot be cancelled cooperatively — + cancelling the awaiting task would free the lock and report "cancelled before + execution" while the work carried on mutating process globals underneath the + next job. Refusing is the honest answer; real cancellation has to be designed + inside the runtime itself. """ task = self._tasks.get(job_key) if task is None or task.done(): return True - task.cancel() + # A stop is raised against a specific run. A suspended job resumes under the same + # key, so a stop for run N that arrives after N+1 started must not kill N+1. + if resume_version is not None: + registered = self._resume_versions.get(job_key) + if registered is not None and registered != resume_version: + console.warning( + f"StopJob for {job_key} ignored: it targets resume version " + f"{resume_version} but the live run is {registered}." + ) + return False + + if job_key not in self._running: + # Still queued behind the lock: cancelling the task is enough, and the job + # never touched process state. + task.cancel() + return True + + # Executing. Cancel the job's OWN event loop rather than the awaiting task: + # that unwinds the runtime cooperatively, so its context managers run and + # output.json still gets written for the caller to fall back on. + control = self._controls.get(job_key) + if control is None: + _server_log(f"StopJob for {job_key}: no control handle") + return False + + # Rung 1 — cancel the ROOT task only. The runtime's cleanup then unwinds + # normally, which is what writes output.json for the caller to fall back on. + # Cancelling everything here would abort that cleanup mid-write. + if not control.cancel(): + _server_log( + f"StopJob for {job_key}: the job has not started an event loop yet; " + "the request is recorded and applied as soon as it does." + ) + + task = self._tasks.get(job_key) + if task is None: + return True + + try: + await asyncio.wait_for(asyncio.shield(task), STOP_GRACE_SECONDS) + return True + except asyncio.TimeoutError: + pass + except BaseException: + # It ended; how it ended is the result push's business. + return True + + # Rung 2 — cleanup itself is stuck. Sweep the loop, giving up on a clean + # output.json in exchange for releasing the lock. + _server_log( + f"StopJob for {job_key}: cleanup did not finish in {STOP_GRACE_SECONDS}s; " + "cancelling every task on the job's loop." + ) + control.cancel_all() + + try: + await asyncio.wait_for(asyncio.shield(task), STOP_ESCALATION_SECONDS) + return True + except asyncio.TimeoutError: + # Rung 3 would be taking the process down, which costs every queued job. + # Report the truth instead and let the caller decide. + _server_log( + f"StopJob for {job_key}: still running — it is blocked in a call that " + "cannot be interrupted (a socket read inside an LLM request, typically)." + ) + return False + except BaseException: + return True + return True diff --git a/packages/uipath/src/uipath/_cli/cli_debug.py b/packages/uipath/src/uipath/_cli/cli_debug.py index 7f542e871..cff28bc89 100644 --- a/packages/uipath/src/uipath/_cli/cli_debug.py +++ b/packages/uipath/src/uipath/_cli/cli_debug.py @@ -1,4 +1,3 @@ -import asyncio import logging from typing import Any, cast, get_args @@ -28,6 +27,7 @@ from uipath.tracing import LiveTrackingSpanProcessor, LlmOpsHttpExporter from ._governance_bootstrap import GovernanceBootstrap, resolve_governance +from ._job_control import run_job_loop from ._telemetry import track_command from ._utils._console import ConsoleLogger from .middlewares import Middlewares @@ -266,7 +266,7 @@ async def execute_debug_runtime(): finally: trace_manager.shutdown() - asyncio.run(execute_debug_runtime()) + run_job_loop(execute_debug_runtime()) except Exception as e: console.error( f"Error occurred: {e or 'Execution failed'}", include_traceback=True diff --git a/packages/uipath/src/uipath/_cli/cli_eval.py b/packages/uipath/src/uipath/_cli/cli_eval.py index 66bdfad10..4483c3b2f 100644 --- a/packages/uipath/src/uipath/_cli/cli_eval.py +++ b/packages/uipath/src/uipath/_cli/cli_eval.py @@ -39,6 +39,7 @@ LlmOpsHttpExporter, ) +from ._job_control import run_job_loop from ._utils._console import ConsoleLogger logger = logging.getLogger(__name__) @@ -528,7 +529,7 @@ async def execute_eval(): finally: await runtime_factory.dispose() - asyncio.run(execute_eval()) + run_job_loop(execute_eval()) except _EvalDiscoveryError as e: click.echo("\n".join(e.get_usage_help())) diff --git a/packages/uipath/src/uipath/_cli/cli_run.py b/packages/uipath/src/uipath/_cli/cli_run.py index 9d12a86c3..2574c0fac 100644 --- a/packages/uipath/src/uipath/_cli/cli_run.py +++ b/packages/uipath/src/uipath/_cli/cli_run.py @@ -1,4 +1,3 @@ -import asyncio from typing import Any import click @@ -36,6 +35,7 @@ from ._errors import EntrypointDiscoveryException from ._governance_bootstrap import GovernanceBootstrap, resolve_governance +from ._job_control import run_job_loop from ._telemetry import track_command from ._utils._console import ConsoleLogger from .middlewares import Middlewares @@ -330,7 +330,7 @@ async def execute() -> None: finally: trace_manager.shutdown() - asyncio.run(execute()) + run_job_loop(execute()) except _RunDiscoveryError as e: click.echo("\n".join(e.get_usage_help())) diff --git a/packages/uipath/src/uipath/_cli/cli_server.py b/packages/uipath/src/uipath/_cli/cli_server.py index 2b7fd2796..f603280cb 100644 --- a/packages/uipath/src/uipath/_cli/cli_server.py +++ b/packages/uipath/src/uipath/_cli/cli_server.py @@ -211,6 +211,7 @@ async def handle_start(request: web.Request) -> web.Response: env_vars, working_dir, HandlerCallback(result_callback_socket), + resume_version=get_field(message, "resumeVersion", "ResumeVersion"), ) if not accepted: return web.json_response( @@ -267,6 +268,36 @@ async def handle_start(request: web.Request) -> web.Response: ) +async def handle_stop(request: web.Request) -> web.Response: + """Stop a job, reporting whether it actually stopped. + + Always 200 with the outcome in the body — 4xx/5xx stay reserved for request-shaped + failures, matching how /start already reports a failed job. + """ + job_key = request.match_info.get("job_key") + if not job_key: + return web.json_response( + {"success": False, "error": "Missing job_key"}, status=400 + ) + + try: + message: dict[str, Any] = await request.json() + except json.JSONDecodeError: + message = {} + + resume_version = get_field(message, "resumeVersion", "ResumeVersion") + stopped = await get_registry().stop(job_key, resume_version) + + return web.json_response( + { + "success": True, + "job_key": job_key, + "contractVersion": CONTRACT_VERSION, + "stopped": stopped, + } + ) + + ALLOWED_HOSTS = {"127.0.0.1", "localhost", "[::1]"} @@ -301,6 +332,7 @@ def create_app() -> web.Application: app = web.Application(middlewares=[host_validation_middleware]) app.router.add_get("/health", handle_health) app.router.add_post("/jobs/{job_key}/start", handle_start) + app.router.add_post("/jobs/{job_key}/stop", handle_stop) return app diff --git a/packages/uipath/src/uipath/_cli/cli_server_ipc.py b/packages/uipath/src/uipath/_cli/cli_server_ipc.py index 38752f53d..161c917e7 100644 --- a/packages/uipath/src/uipath/_cli/cli_server_ipc.py +++ b/packages/uipath/src/uipath/_cli/cli_server_ipc.py @@ -56,7 +56,12 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: @abstractmethod async def StopJob(self, request: StopJobRequest) -> bool: - """Cancel a running job by key (bool return avoids fire-and-forget).""" + """Cancel a running job (bool return avoids fire-and-forget). + + A nested DTO, matching the .NET peer: a suspended job resumes under the same + JobKey, so ``ResumeVersion`` is what keeps a stop aimed at the previous run from + killing the resumed one. + """ class PythonRuntimeService(IPythonRuntimeServer): @@ -86,6 +91,7 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: request.EnvironmentVariables, request.WorkingDirectory, HandlerCallback(callback_socket), + resume_version=request.ResumeVersion, ) if not accepted: return RunJobResult( @@ -106,11 +112,7 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: return RunJobResult(ExitCode=result["ExitCode"], Error=result["Error"]) async def StopJob(self, request: StopJobRequest) -> bool: - console.info( - f"StopJob requested for {_run_id(request.JobKey, request.ResumeVersion)} " - f"(force={request.ForceStop}) (no-op)" - ) - return True + return await get_registry().stop(request.JobKey, request.ResumeVersion) async def start_ipc_server(pipe_name: str) -> None: diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py index f89ecd4e5..feb605442 100644 --- a/packages/uipath/tests/cli/test_server_async.py +++ b/packages/uipath/tests/cli/test_server_async.py @@ -8,6 +8,7 @@ import asyncio import json import os +import time from typing import Any, cast import click @@ -26,6 +27,7 @@ JobRegistry, build_result_payload, ) +from uipath._cli.cli_server_ipc import StopJobRequest @pytest.fixture(autouse=True) @@ -169,6 +171,7 @@ def test_build_result_payload_shape(): "stateConveyance": "file", "jobConveyance": "inline", "job": '{"status":"successful"}', + "stopped": False, } @@ -251,6 +254,60 @@ async def test_stop_is_true_for_an_unknown_job(): assert await registry.stop("never-seen") is True +async def test_stop_waits_for_a_running_job_to_finish(tmp_path, monkeypatch): + """Stop no longer refuses a running job — it cancels and waits. A job that ends + inside the grace window is a successful stop, whatever ended it.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "STOP_GRACE_SECONDS", 5) + monkeypatch.setattr(jobs, "STOP_ESCALATION_SECONDS", 2) + + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + started = asyncio.Event() + + @click.command() + def _brief_command() -> None: + started.set() + time.sleep(0.3) + + registry.start("job-1", _brief_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(started.wait(), timeout=10) + + assert await registry.stop("job-1") is True + await asyncio.wait_for(callback.done.wait(), timeout=10) + + +async def test_stop_cancels_a_job_that_is_still_queued(tmp_path): + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start("job-1", _holds_the_lock, [], {}, str(tmp_path), first) + await asyncio.wait_for(started.wait(), timeout=10) + + # job-2 is queued behind the lock, so it has not started executing. + registry.start("job-2", _ok_command, [], {}, str(tmp_path), second) + cancelled = await registry.stop("job-2") + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + assert cancelled is True + + # --------------------------------------------------------------------------- # # HTTP dispatch # # --------------------------------------------------------------------------- # @@ -279,9 +336,13 @@ class _RecordingRegistry: def __init__(self, accept: bool = True) -> None: self.accept = accept self.started: list[str] = [] + self.resume_versions: list[int | None] = [] - def start(self, job_key, cmd, args, env_vars, working_dir, callback): + def start( + self, job_key, cmd, args, env_vars, working_dir, callback, resume_version=None + ): self.started.append(job_key) + self.resume_versions.append(resume_version) return self.accept @@ -396,3 +457,125 @@ async def _fake_isolated(cmd, args, env_vars, working_dir): assert result.ExitCode == 7 assert result.Disposition is None + + +async def test_stop_ignores_a_command_for_a_previous_resume_version(tmp_path): + """A suspended job resumes under the SAME key. A stop raised against run N that + arrives after N+1 started must not kill N+1.""" + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start( + "job-1", _holds_the_lock, [], {}, str(tmp_path), first, resume_version=2 + ) + await asyncio.wait_for(started.wait(), timeout=10) + + # A stop for the run that was suspended before this one. + assert await registry.stop("job-1", resume_version=1) is False + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + assert first.results[0]["exitCode"] == 0 + + +async def test_stop_applies_to_the_matching_resume_version(tmp_path): + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start("job-1", _holds_the_lock, [], {}, str(tmp_path), first) + await asyncio.wait_for(started.wait(), timeout=10) + + registry.start( + "job-2", _ok_command, [], {}, str(tmp_path), second, resume_version=3 + ) + assert await registry.stop("job-2", resume_version=3) is True + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + +# --------------------------------------------------------------------------- # +# HTTP stop — the transport production actually uses # +# --------------------------------------------------------------------------- # + + +class _RecordingStopRegistry: + def __init__(self, stopped: bool = True) -> None: + self.stopped = stopped + self.calls: list[tuple[str, int | None]] = [] + + async def stop(self, job_key, resume_version=None): + self.calls.append((job_key, resume_version)) + return self.stopped + + +async def test_http_stop_reaches_the_registry(monkeypatch): + """Stop was previously unreachable over HTTP, which is the transport carrying all + production traffic — an operator stop simply never reached the runtime.""" + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server, "get_registry", lambda: registry) + + response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + + assert response.status == 200 + body = json.loads(response.text) + assert body["stopped"] is True + assert registry.calls == [("job-1", None)] + + +async def test_http_stop_forwards_the_resume_version(monkeypatch): + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server, "get_registry", lambda: registry) + + await cli_server.handle_stop(_FakeRequest("job-1", {"resumeVersion": 2})) + + assert registry.calls == [("job-1", 2)] + + +async def test_http_stop_reports_a_refused_stop(monkeypatch): + """A job wedged in a non-cancellable call must be reported honestly, not as stopped.""" + monkeypatch.setattr( + cli_server, "get_registry", lambda: _RecordingStopRegistry(stopped=False) + ) + + response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + + assert response.status == 200 + assert json.loads(response.text)["stopped"] is False + + +async def test_ipc_stop_forwards_the_dto_fields(monkeypatch): + """The .NET peer sends a StopJobRequest; both identity fields must reach the registry + or the resume-version guard is inert.""" + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server_ipc, "get_registry", lambda: registry) + + await cli_server_ipc.PythonRuntimeService().StopJob( + StopJobRequest(JobKey="job-1", ResumeVersion=2, ForceStop=True) + ) + + assert registry.calls == [("job-1", 2)] diff --git a/packages/uipath/tests/cli/test_server_callbacks.py b/packages/uipath/tests/cli/test_server_callbacks.py index daa699edd..f2015dd73 100644 --- a/packages/uipath/tests/cli/test_server_callbacks.py +++ b/packages/uipath/tests/cli/test_server_callbacks.py @@ -82,6 +82,7 @@ def test_result_payload_carries_the_keys_the_handler_deserializes(): "stateConveyance", "jobConveyance", "job", + "stopped", } diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py new file mode 100644 index 000000000..3e88ee2d7 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -0,0 +1,222 @@ +"""Really stopping a job that is already executing. + +``run``/``debug``/``eval`` each end in ``asyncio.run(...)`` on the worker thread, so a +job *is* an event loop. Capturing that loop turns cancellation from "impossible, the +thread is opaque" into an ordinary ``task.cancel()`` that unwinds the runtime +cooperatively — which is what these pin, using commands shaped like the real ones. +""" + +import asyncio +import time + +import click +import pytest + +from uipath._cli import _server_core +from uipath._cli._job_control import CURRENT_JOB_CONTROL, JobControl, run_job_loop +from uipath._cli._server_core import EXIT_CODE_STOPPED, _ServerState +from uipath._cli._server_jobs import JobRegistry + + +@pytest.fixture(autouse=True) +def _fresh_state(monkeypatch): + monkeypatch.setattr(_server_core, "_state", _ServerState()) + + +class FakeCallback: + def __init__(self) -> None: + self.results: list[dict] = [] + self.done = asyncio.Event() + + async def post_result(self, job_key, payload): + self.results.append(payload) + self.done.set() + return True + + async def post_logs(self, job_key, lines): + return True + + +# Shaped like the real commands: a click command whose body is asyncio.run(...). +_started = None +_cleanup_ran = None + + +@click.command() +def _long_async_command() -> None: + async def body() -> None: + try: + _started.set() + await asyncio.sleep(30) + finally: + # Stands in for UiPathRuntimeContext.__exit__, which writes output.json. + _cleanup_ran.set() + + run_job_loop(body()) + + +@click.command() +def _quick_async_command() -> None: + async def body() -> None: + await asyncio.sleep(0) + + run_job_loop(body()) + + +@click.command() +def _blocking_command() -> None: + # No event loop at all: models a job wedged in a non-cancellable C call. + _started.set() + time.sleep(30) + + +@pytest.fixture(autouse=True) +def _events(): + global _started, _cleanup_ran + _started = __import__("threading").Event() + _cleanup_ran = __import__("threading").Event() + yield + + +async def _wait_for(event, timeout=10.0): + deadline = asyncio.get_running_loop().time() + timeout + while not event.is_set(): + if asyncio.get_running_loop().time() > deadline: + raise AssertionError("timed out waiting for the job to start") + await asyncio.sleep(0.02) + + +# --------------------------------------------------------------------------- # +# the capture mechanism # +# --------------------------------------------------------------------------- # + + +def test_control_starts_unbound(): + control = JobControl("job-1") + + assert control.bound is False + # Nothing to cancel yet, but the request must be remembered for bind(). + assert control.cancel() is False + assert control.cancel_requested is True + + +def test_run_job_loop_is_plain_asyncio_run_outside_the_server(): + """`uipath run` on a terminal has no control in scope and must be unaffected.""" + assert CURRENT_JOB_CONTROL.get() is None + + async def body(): + return 42 + + assert run_job_loop(body()) == 42 + + +async def test_control_binds_to_the_loop_the_job_runs_on(tmp_path): + _server_core._state.init() + control = JobControl("job-1") + + await _server_core._run_command_isolated( + _quick_async_command, [], {}, str(tmp_path), control=control + ) + + # Unbound again once the job finished. + assert control.bound is False + + +async def test_a_stop_that_races_the_loop_is_still_applied(tmp_path): + """A cancel arriving before the job builds its loop must not be dropped.""" + _server_core._state.init() + control = JobControl("job-1") + control.cancel() # before anything is bound + + result = await _server_core._run_command_isolated( + _long_async_command, [], {}, str(tmp_path), control=control + ) + + assert result["Stopped"] is True + + +# --------------------------------------------------------------------------- # +# stopping a running job for real # +# --------------------------------------------------------------------------- # + + +async def test_stop_cancels_a_job_that_is_already_executing(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + stopped = await registry.stop("job-1") + + assert stopped is True, "a running job must actually stop, not just be refused" + await asyncio.wait_for(callback.done.wait(), timeout=10) + assert callback.results[0]["stopped"] is True + assert callback.results[0]["exitCode"] == EXIT_CODE_STOPPED + + +async def test_a_stopped_job_still_unwinds_its_cleanup(tmp_path): + """Cooperative cancellation, not a thread kill: the runtime's context managers must + still run, because that is what writes output.json for the caller to fall back on.""" + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + await registry.stop("job-1") + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert _cleanup_ran.is_set(), "the job's finally block must have run" + + +async def test_stop_releases_the_lock_for_the_next_job(tmp_path): + """A stop that leaves the lock held would wedge the whole server.""" + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), first) + await _wait_for(_started) + await registry.stop("job-1") + await asyncio.wait_for(first.done.wait(), timeout=10) + + registry.start("job-2", _quick_async_command, [], {}, str(tmp_path), second) + await asyncio.wait_for(second.done.wait(), timeout=10) + + assert second.results[0]["exitCode"] == 0 + + +async def test_a_normal_job_is_not_reported_stopped(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _quick_async_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert callback.results[0]["stopped"] is False + assert callback.results[0]["exitCode"] == 0 + + +async def test_stop_reports_failure_when_the_job_cannot_be_interrupted( + tmp_path, monkeypatch +): + """A job with no event loop — wedged in a C call — cannot be cancelled. Say so + rather than claiming a stop that did not happen.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "STOP_GRACE_SECONDS", 1) + monkeypatch.setattr(jobs, "STOP_ESCALATION_SECONDS", 1) + + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _blocking_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + assert await registry.stop("job-1") is False diff --git a/packages/uipath/tests/cli/test_server_ipc.py b/packages/uipath/tests/cli/test_server_ipc.py index 43fe313c9..6110738df 100644 --- a/packages/uipath/tests/cli/test_server_ipc.py +++ b/packages/uipath/tests/cli/test_server_ipc.py @@ -201,7 +201,7 @@ def test_stop_job_accepts_resume_version_and_force_stop(self, pipe): assert result is True def test_stop_job_returns_true(self, pipe): - """StopJob is a no-op stub today, but must ack (bool) so the call is awaitable.""" + """StopJob must ack (bool) so the call is awaitable; an unknown key is a no-op.""" result = asyncio.run( _with_proxy( pipe, lambda p: p.StopJob({"JobKey": "job-1", "ForceStop": True}) From 5468b7d0bfb7db1af5ab9ee80d681175c364921f Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 18:51:04 +0300 Subject: [PATCH 2/3] fix(cli): deliver one cancellation per stop, and only call it Stopped when asked Two defects in the cancellation path, both found in review. JobControl.cancel() documented that it must be called at most once, but JobRegistry.stop() called it unguarded on every request. A stop followed by a force-stop escalation is ordinary, and the second delivery lands inside the runtime's cleanup finally blocks -- the ones that write output.json -- and aborts them, destroying the fallback the caller relies on. Absorbing the repeat belongs in JobControl, not in every caller. _invoke_command treated any CancelledError with cancelling() == 0 as a user-requested stop. cancelling() reports on the awaiting server task, so a CancelledError the job's own code let escape was reported as exit 143 "stopped on request" with no StopJob in sight -- a fault filed as a clean stop. Gate the classification on the control's own cancel_requested. Also corrects JobRegistry.stop()'s docstring, which still described the pre-cancellation behaviour of refusing to stop executing work. --- .../uipath/src/uipath/_cli/_job_control.py | 23 ++++-- .../uipath/src/uipath/_cli/_server_core.py | 28 ++++--- .../uipath/src/uipath/_cli/_server_jobs.py | 14 ++-- .../tests/cli/test_server_cancellation.py | 74 +++++++++++++++++++ 4 files changed, 114 insertions(+), 25 deletions(-) diff --git a/packages/uipath/src/uipath/_cli/_job_control.py b/packages/uipath/src/uipath/_cli/_job_control.py index f39df2e76..1a32b4d64 100644 --- a/packages/uipath/src/uipath/_cli/_job_control.py +++ b/packages/uipath/src/uipath/_cli/_job_control.py @@ -24,6 +24,7 @@ class JobControl: def __init__(self, job_key: str) -> None: self.job_key = job_key self.cancel_requested = False + self._delivered = False self._loop: asyncio.AbstractEventLoop | None = None self._task: "asyncio.Task[Any] | None" = None # Bound from the job thread, read from the server loop. @@ -51,23 +52,29 @@ def bound(self) -> bool: return self._task is not None def cancel(self) -> bool: - """Deliver CancelledError to the job's ROOT task only. + """Deliver CancelledError to the job's ROOT task, at most once. - Call this at most once. A second cancel lands inside the runtime's cleanup - ``finally`` blocks — the ones that write ``output.json`` and tear the log - interceptor down — and aborts them, which would destroy the very fallback the - caller relies on. Escalate with ``cancel_all`` instead. + Delivering twice lands the second one inside the runtime's cleanup ``finally`` + blocks — the ones that write ``output.json`` and tear the log interceptor down — + and aborts them, destroying the very fallback the caller relies on. Repeat stops + are ordinary (a stop followed by a force-stop escalation), so absorbing them is + this object's job rather than every caller's. Escalate with ``cancel_all``. """ with self._sync: self.cancel_requested = True + if self._delivered: + return True loop, task = self._loop, self._task - if loop is None or task is None: - # Not bound yet; bind() will apply it. - return False + if loop is None or task is None: + # Not bound yet; bind() will apply it. + return False + self._delivered = True try: loop.call_soon_threadsafe(task.cancel) except RuntimeError: # Loop already closed — the job is finishing anyway. + with self._sync: + self._delivered = False return False return True diff --git a/packages/uipath/src/uipath/_cli/_server_core.py b/packages/uipath/src/uipath/_cli/_server_core.py index 750f263cd..0bafc2b1a 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/packages/uipath/src/uipath/_cli/_server_core.py @@ -175,20 +175,30 @@ async def _invoke_command( "Unexpected": False, } except asyncio.CancelledError: - # Two very different things arrive here. cancelling() > 0 means OUR awaiting task - # was cancelled (server shutdown) — never swallow that, or the shutdown stalls. - # 0 means the CancelledError travelled out of the job's own loop, i.e. our StopJob - # landed: that is an OUTCOME, and swallowing it keeps this task alive so the lock - # unwinds and the result document still gets read. + # Three very different things arrive here. current = asyncio.current_task() if current is not None and current.cancelling() > 0: + # OUR awaiting task was cancelled (server shutdown). Never swallow that, or + # the shutdown stalls. raise + if control is not None and control.cancel_requested: + # A stop we asked for travelled out of the job's own loop. That is an + # OUTCOME: swallowing it keeps this task alive so the lock unwinds and the + # result document still gets read. + return { + "ExitCode": EXIT_CODE_STOPPED, + "Error": "Job stopped on request", + "Result": None, + "Unexpected": False, + "Stopped": True, + } + # Nobody asked for this: the job's own code let a CancelledError escape. It is a + # failure, and reporting it as Stopped would file a fault as a clean stop. return { - "ExitCode": EXIT_CODE_STOPPED, - "Error": "Job stopped on request", + "ExitCode": 1, + "Error": "Job cancelled", "Result": None, - "Unexpected": False, - "Stopped": True, + "Unexpected": True, } except Exception as e: # report any job failure as a result, not a fault return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True} diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index edf393512..a984c3886 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -416,14 +416,12 @@ async def finish(outcome: dict[str, Any]) -> None: await finish(outcome) async def stop(self, job_key: str, resume_version: int | None = None) -> bool: - """Cancel a job that has not started executing yet. - - Only queued work can be stopped. Once the job holds the lock its body is on a - thread via ``asyncio.to_thread``, which cannot be cancelled cooperatively — - cancelling the awaiting task would free the lock and report "cancelled before - execution" while the work carried on mutating process globals underneath the - next job. Refusing is the honest answer; real cancellation has to be designed - inside the runtime itself. + """Stop a job, whether it is still queued or already executing. + + Queued work is cancelled by dropping the waiting task. Executing work is + cancelled through the job's own event loop so the runtime unwinds cooperatively. + Safe to call repeatedly for the same run: the second request does not deliver a + second cancellation. """ task = self._tasks.get(job_key) if task is None or task.done(): diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py index 3e88ee2d7..8fd0efbac 100644 --- a/packages/uipath/tests/cli/test_server_cancellation.py +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -63,6 +63,31 @@ async def body() -> None: run_job_loop(body()) +@click.command() +def _slow_cleanup_command() -> None: + async def body() -> None: + try: + _started.set() + await asyncio.sleep(30) + finally: + # The real cleanup awaits (flushing traces, closing clients), which is the + # window a second cancellation would land in. + await asyncio.sleep(0.5) + _cleanup_ran.set() + + run_job_loop(body()) + + +@click.command() +def _self_cancelling_command() -> None: + async def body() -> None: + _started.set() + # Nobody asked this job to stop; its own code let a CancelledError escape. + raise asyncio.CancelledError() + + run_job_loop(body()) + + @click.command() def _blocking_command() -> None: # No event loop at all: models a job wedged in a non-cancellable C call. @@ -156,6 +181,55 @@ async def test_stop_cancels_a_job_that_is_already_executing(tmp_path): assert callback.results[0]["exitCode"] == EXIT_CODE_STOPPED +async def test_cancel_delivers_a_single_cancellation_however_often_it_is_called(): + """A second delivery lands in the cleanup and aborts it, so cancel() absorbs it.""" + control = JobControl("job-1") + loop = asyncio.get_running_loop() + task = loop.create_task(asyncio.sleep(30)) + control.bind(loop, task) + + assert control.cancel() is True + assert control.cancel() is True + assert control.cancel() is True + + await asyncio.gather(task, return_exceptions=True) + assert task.cancelling() == 1 + + +async def test_a_repeated_stop_does_not_abort_the_cleanup(tmp_path): + """A stop followed by a force-stop escalation is ordinary; both must be honoured + without the second one interrupting the cleanup that writes output.json.""" + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _slow_cleanup_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + first, second = await asyncio.gather(registry.stop("job-1"), registry.stop("job-1")) + + assert (first, second) == (True, True) + await asyncio.wait_for(callback.done.wait(), timeout=10) + assert _cleanup_ran.is_set(), "the second stop interrupted the job's cleanup" + assert callback.results[0]["stopped"] is True + + +async def test_a_self_inflicted_cancellation_is_a_fault_not_a_stop(tmp_path): + """`cancelling() == 0` alone does not mean the caller asked for a stop — without a + stop request a stray CancelledError is a job failure, and filing it as Stopped + would report a fault as a clean stop.""" + _server_core._state.init() + control = JobControl("job-1") + + result = await _server_core._run_command_isolated( + _self_cancelling_command, [], {}, str(tmp_path), control=control + ) + + assert result.get("Stopped") is not True + assert result["ExitCode"] != EXIT_CODE_STOPPED + assert result["Unexpected"] is True + + async def test_a_stopped_job_still_unwinds_its_cleanup(tmp_path): """Cooperative cancellation, not a thread kill: the runtime's context managers must still run, because that is what writes output.json for the caller to fall back on.""" From 1df44eb75e20d70b63dafdb58ce1c8a47fe2bda3 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 18:58:45 +0300 Subject: [PATCH 3/3] fix(cli): type-check the cancellation tests Same CI gate as the parent commit, applied to the tests this branch adds: give the module-level events a real Event type instead of letting them infer None, annotate the recording callback against the JobReporter protocol, and cast the _FakeRequest stand-ins at the handle_stop call sites. --- .../uipath/tests/cli/test_server_async.py | 10 +++++----- .../tests/cli/test_server_cancellation.py | 19 +++++++++++-------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py index feb605442..ed39f9e69 100644 --- a/packages/uipath/tests/cli/test_server_async.py +++ b/packages/uipath/tests/cli/test_server_async.py @@ -539,10 +539,10 @@ async def test_http_stop_reaches_the_registry(monkeypatch): registry = _RecordingStopRegistry() monkeypatch.setattr(cli_server, "get_registry", lambda: registry) - response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + response = await cli_server.handle_stop(_fake_request("job-1", {})) assert response.status == 200 - body = json.loads(response.text) + body = _body(response) assert body["stopped"] is True assert registry.calls == [("job-1", None)] @@ -551,7 +551,7 @@ async def test_http_stop_forwards_the_resume_version(monkeypatch): registry = _RecordingStopRegistry() monkeypatch.setattr(cli_server, "get_registry", lambda: registry) - await cli_server.handle_stop(_FakeRequest("job-1", {"resumeVersion": 2})) + await cli_server.handle_stop(_fake_request("job-1", {"resumeVersion": 2})) assert registry.calls == [("job-1", 2)] @@ -562,10 +562,10 @@ async def test_http_stop_reports_a_refused_stop(monkeypatch): cli_server, "get_registry", lambda: _RecordingStopRegistry(stopped=False) ) - response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + response = await cli_server.handle_stop(_fake_request("job-1", {})) assert response.status == 200 - assert json.loads(response.text)["stopped"] is False + assert _body(response)["stopped"] is False async def test_ipc_stop_forwards_the_dto_fields(monkeypatch): diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py index 8fd0efbac..79c6787a1 100644 --- a/packages/uipath/tests/cli/test_server_cancellation.py +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -7,7 +7,9 @@ """ import asyncio +import threading import time +from typing import Any import click import pytest @@ -25,21 +27,22 @@ def _fresh_state(monkeypatch): class FakeCallback: def __init__(self) -> None: - self.results: list[dict] = [] + self.results: list[dict[str, Any]] = [] self.done = asyncio.Event() - async def post_result(self, job_key, payload): + async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: self.results.append(payload) self.done.set() return True - async def post_logs(self, job_key, lines): + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: return True # Shaped like the real commands: a click command whose body is asyncio.run(...). -_started = None -_cleanup_ran = None +# Rebound per test by the _events fixture. +_started = threading.Event() +_cleanup_ran = threading.Event() @click.command() @@ -98,12 +101,12 @@ def _blocking_command() -> None: @pytest.fixture(autouse=True) def _events(): global _started, _cleanup_ran - _started = __import__("threading").Event() - _cleanup_ran = __import__("threading").Event() + _started = threading.Event() + _cleanup_ran = threading.Event() yield -async def _wait_for(event, timeout=10.0): +async def _wait_for(event: threading.Event, timeout: float = 10.0) -> None: deadline = asyncio.get_running_loop().time() + timeout while not event.is_set(): if asyncio.get_running_loop().time() > deadline: