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..1a32b4d64 --- /dev/null +++ b/packages/uipath/src/uipath/_cli/_job_control.py @@ -0,0 +1,120 @@ +"""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._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. + 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, at most once. + + 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 + 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 + + 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..0bafc2b1a 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,37 @@ async def _invoke_command(cmd: Any, args: list[str]) -> dict[str, Any]: "Result": None, "Unexpected": False, } + except asyncio.CancelledError: + # 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": 1, + "Error": "Job cancelled", + "Result": None, + "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} + finally: + if token is not None: + CURRENT_JOB_CONTROL.reset(token) async def _run_command_isolated( @@ -172,12 +212,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 +249,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..a984c3886 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,87 @@ async def finish(outcome: dict[str, Any]) -> None: await finish(outcome) - async def stop(self, job_key: str) -> bool: - """Cancel a job that has not started executing yet. + async def stop(self, job_key: str, resume_version: int | None = None) -> bool: + """Stop a job, whether it is still queued or already executing. - 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. + 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(): 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..ed39f9e69 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(_fake_request("job-1", {})) + + assert response.status == 200 + body = _body(response) + 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(_fake_request("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(_fake_request("job-1", {})) + + assert response.status == 200 + assert _body(response)["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..79c6787a1 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -0,0 +1,299 @@ +"""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 threading +import time +from typing import Any + +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[str, Any]] = [] + self.done = asyncio.Event() + + 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: str, lines: list[dict[str, Any]]) -> bool: + return True + + +# Shaped like the real commands: a click command whose body is asyncio.run(...). +# Rebound per test by the _events fixture. +_started = threading.Event() +_cleanup_ran = threading.Event() + + +@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 _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. + _started.set() + time.sleep(30) + + +@pytest.fixture(autouse=True) +def _events(): + global _started, _cleanup_ran + _started = threading.Event() + _cleanup_ran = threading.Event() + yield + + +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: + 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_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.""" + _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})