feat(cli): dispatch jobs asynchronously and push results/logs to the caller

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

packages/uipath/src/uipath/_cli/_server_core.py

@@ -1,10 +1,13 @@
 """Transport-agnostic job core shared by the HTTP and uipath-ipc channels."""
 
 import asyncio
+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
@@ -45,17 +48,174 @@
     return []
 
 
+# The document is carried inline on the result push when it fits. uipath_ipc caps a
+# frame at 2 MiB and Orchestrator already spills large outputs to an attachment, so
+# anything bigger stays on disk and the caller reads the file as it always has.
+MAX_INLINE_DOCUMENT_BYTES = 1024 * 1024
+
+DEFAULT_RUNTIME_DIR = "__uipath"
+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
+) -> str | None:
+    """Resolve one ``runtime.*`` file path from a uipath.json.
+
+    Mirrors UiPathRuntimeContext.from_config's ``runtime.dir`` / ``runtime.<key>``
+    mapping. ``base_dir`` anchors relative paths so this works without ever changing
+    the process's cwd.
+    """
+    if not os.path.isabs(config_path):
+        config_path = os.path.join(base_dir, config_path)
+
+    runtime: dict[str, Any] = {}
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            loaded = json.load(f)
+        if isinstance(loaded, dict):
+            runtime = loaded.get("runtime") or {}
+    except (OSError, json.JSONDecodeError, UnicodeDecodeError):
+        # Fall through to the runtime's own defaults rather than giving up. Returning
+        # None here would stop the log tailer from starting at all — and by then the
+        # caller has already been told we forward logs and stopped tailing the file
+        # itself, so the job would produce no logs anywhere.
+        runtime = {}
+
+    directory = runtime.get("dir") or DEFAULT_RUNTIME_DIR
+    name = runtime.get(key) or default_name
+    if not isinstance(directory, str) or not isinstance(name, str):
+        directory, name = DEFAULT_RUNTIME_DIR, default_name
+
+    if not os.path.isabs(directory):
+        directory = os.path.join(base_dir, directory)
+    return os.path.abspath(os.path.join(directory, name))
+
+
+def resolve_result_file_path() -> str | None:
+    """Where this job's terminal document lives. Call inside the job's env and cwd."""
+    return _resolve_runtime_file(
+        os.environ.get("UIPATH_CONFIG_PATH", "uipath.json"),
+        os.getcwd(),
+        "outputFile",
+        DEFAULT_RESULT_FILE,
+    )
+
+
+def resolve_logs_file_path(
+    env_vars: dict[str, str] | None, working_dir: str | None
+) -> str | None:
+    """Where this job will write its log file, derived from the request alone.
+
+    Pure with respect to process state: the log tailer has to know the path *before*
+    the job takes the lock and applies its env/cwd.
+    """
+    env_vars = env_vars or {}
+    base_dir = working_dir or os.getcwd()
+    config_path = env_vars.get("UIPATH_CONFIG_PATH", "uipath.json")
+    return _resolve_runtime_file(config_path, base_dir, "logsFile", DEFAULT_LOGS_FILE)
+
+
+def _read_result_document() -> tuple[str | None, str]:
+    """Return ``(document, conveyance)`` for the terminal result document.
+
+    ``conveyance`` is ``inline`` when the document rides the wire, ``file`` when the
+    caller must read it from disk (too large, unreadable, or never written).
+    """
+    path = resolve_result_file_path()
+    if not path or not os.path.exists(path):
+        return None, "file"
+
+    try:
+        if os.path.getsize(path) > MAX_INLINE_DOCUMENT_BYTES:
+            return None, "file"
+        with open(path, encoding="utf-8") as f:
+            return f.read(), "inline"
+    except (OSError, UnicodeDecodeError):
+        return None, "file"
+
+
+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
+        # raising SystemExit, so a bare int is the exit code, not a result — every
+        # ConsoleLogger.error path lands here via ctx.exit(1). The run/debug/eval
+        # callbacks only ever return a result object or None, so this is unambiguous.
+        if isinstance(result_value, int) and not isinstance(result_value, bool):
+            return {
+                "ExitCode": result_value,
+                "Error": None if result_value == 0 else f"Exit code: {result_value}",
+                "Result": None,
+                "Unexpected": False,
+            }
+        return {
+            "ExitCode": 0,
+            "Error": None,
+            "Result": result_value,
+            "Unexpected": False,
+        }
+    except SystemExit as e:
+        exit_code = e.code if isinstance(e.code, int) else 1
+        return {
+            "ExitCode": exit_code,
+            "Error": None if exit_code == 0 else f"Exit code: {exit_code}",
+            "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(
     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
@@ -79,25 +239,13 @@
                         "ClientError": True,
                     }
 
-            result_value = await asyncio.to_thread(
-                cmd.main, args, standalone_mode=False
-            )
-            return {
-                "ExitCode": 0,
-                "Error": None,
-                "Result": result_value,
-                "Unexpected": False,
-            }
-        except SystemExit as e:
-            exit_code = e.code if isinstance(e.code, int) else 1
-            return {
-                "ExitCode": exit_code,
-                "Error": None if exit_code == 0 else f"Exit code: {exit_code}",
-                "Result": None,
-                "Unexpected": False,
-            }
-        except Exception as e:  # report any job failure as a result, not a fault
-            return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True}
+            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()
+            outcome["Document"] = document
+            outcome["DocumentConveyance"] = conveyance
+            return outcome
         finally:
             # Restore to server baseline.
             try: