feat(cli): really stop a job that is already executing [PC-4873]

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()

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,21 +174,58 @@ 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(
     cmd: Any,
     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()