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 c426126ff..750f263cd 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/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 @@ def parse_args(args: str | list[str] | None) -> list[str]: 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.`` + 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 @@ async def _run_command_isolated( "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: diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py new file mode 100644 index 000000000..4467a438b --- /dev/null +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -0,0 +1,496 @@ +"""Async job dispatch: enqueue a job, run it, push the outcome back to the caller. + +``StartJob`` used to block for the whole job, which capped a run at the caller's +request timeout and left no channel for anything that happens *during* a job. Here +the call returns as soon as the job is registered, and the outcome travels back over +the caller-owned HTTP socket it already told us about when we ACKed readiness. + +Execution is still serialised by ``_server_core``'s process-wide lock — a job mutates +process globals (logging handlers, OTel providers, env, cwd), so only one may run at a +time. Queueing changes who waits, not how many run. +""" + +import asyncio +import contextlib +import enum +import os +import re +import sys +from typing import Any + +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 + +console = ConsoleLogger() + +# Server-side diagnostics must NOT go through ConsoleLogger while a job is running. +# ConsoleLogger resolves sys.stdout at call time, and the runtime's log interceptor has +# replaced it with a writer feeding the job's execution.log — which the tailer then reads +# and posts back. With the callback down that is a self-feeding loop: warn -> written to +# execution.log -> tailed -> posted -> fails -> warn. Bind the real stream once, before +# any job can redirect it. +_SERVER_STDERR = sys.stderr + + +def _server_log(message: str) -> None: + """Diagnostics that must never be captured as a job's own logs.""" + try: + print(message, file=_SERVER_STDERR, flush=True) + except Exception: + pass + + +CONTRACT_VERSION = 1 + +# A lost terminal push strands the job: StartJob no longer blocks, so no request timeout +# will complete it and nothing else on the caller's side is watching. The realistic cause +# is the callback socket being briefly unavailable (handler restart), so retry over +# minutes rather than seconds — roughly 4.5 min of wall clock across these attempts. +RESULT_PUSH_ATTEMPTS = 8 +RESULT_PUSH_BACKOFF_SECONDS = 2.0 +RESULT_PUSH_BACKOFF_CAP_SECONDS = 60.0 +CALLBACK_TIMEOUT_SECONDS = 30 + +# Exit code used when the caller's callback socket is unreachable. Deliberately NOT 137: +# the caller skips result processing entirely on SIGKILL, which would throw away the very +# files we are exiting in order to let it read. +EXIT_CALLBACK_UNREACHABLE = 70 + +_shutdown_event: asyncio.Event | None = None +_shutdown_loop: asyncio.AbstractEventLoop | None = None + + +def shutdown_event() -> asyncio.Event: + """Set when the server must stop because it can no longer reach the caller. + + Rebound whenever the running loop changes. The server has exactly one loop for its + whole life, so in production this is created once; the rebinding keeps the flag from + leaking between the short-lived loops that tests create. + """ + global _shutdown_event, _shutdown_loop + loop = asyncio.get_running_loop() + if _shutdown_event is None or _shutdown_loop is not loop: + _shutdown_event = asyncio.Event() + _shutdown_loop = loop + return _shutdown_event + + +def request_shutdown(reason: str) -> None: + _server_log(f"Requesting server shutdown: {reason}") + shutdown_event().set() + + +class _PostOutcome(enum.Enum): + """Why a callback POST ended, because the three cases need different handling.""" + + DELIVERED = "delivered" + REJECTED = "rejected" # 4xx — the caller is up and has moved on; do not retry + UNREACHABLE = "unreachable" # transport failure or 5xx — worth retrying + + +class HandlerCallback: + """Posts job lifecycle events to the caller's Unix-socket HTTP API.""" + + def __init__(self, socket_path: str) -> None: + self.socket_path = socket_path + + async def _post(self, path: str, payload: dict[str, Any]) -> _PostOutcome: + conn = UnixConnector(path=self.socket_path) + try: + async with ClientSession( + connector=conn, timeout=ClientTimeout(total=CALLBACK_TIMEOUT_SECONDS) + ) as session: + async with session.post( + f"http://localhost{path}", json=payload + ) as response: + if response.status < 300: + return _PostOutcome.DELIVERED + body = await response.text() + _server_log( + f"Callback {path} rejected with {response.status}: {body}" + ) + # 4xx is the caller telling us this job is unknown or already + # finished. Retrying cannot change that, and treating it as a lost + # push would take the whole server down over one stale report. + if 400 <= response.status < 500: + return _PostOutcome.REJECTED + return _PostOutcome.UNREACHABLE + except Exception as e: + _server_log(f"Callback {path} failed: {e}") + return _PostOutcome.UNREACHABLE + + async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: + path = f"/api/python/jobs/{job_key}/result" + for attempt in range(1, RESULT_PUSH_ATTEMPTS + 1): + outcome = await self._post(path, payload) + if outcome is _PostOutcome.DELIVERED: + return True + if outcome is _PostOutcome.REJECTED: + # Refused, not lost: the caller is reachable and has moved on. Report + # success so no one tries to recover a job that is already resolved. + return True + if attempt < RESULT_PUSH_ATTEMPTS: + await asyncio.sleep( + min( + RESULT_PUSH_BACKOFF_SECONDS * (2 ** (attempt - 1)), + RESULT_PUSH_BACKOFF_CAP_SECONDS, + ) + ) + # Deliberately warning, not error: ConsoleLogger.error is NoReturn (it calls + # ctx.exit), which would raise outside a click context and kill this task. + _server_log( + f"Giving up on the result push for job {job_key} after " + f"{RESULT_PUSH_ATTEMPTS} attempts; the caller will fall back to the file." + ) + return False + + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: + # Logs are best-effort: a dropped batch must never delay or fail the job. + outcome = await self._post( + f"/api/python/jobs/{job_key}/logs", + {"contractVersion": CONTRACT_VERSION, "jobKey": job_key, "lines": lines}, + ) + return outcome is _PostOutcome.DELIVERED + + +def build_result_payload(job_key: str, outcome: dict[str, Any]) -> dict[str, Any]: + """Shape ``_run_command_isolated``'s outcome into the terminal result push.""" + return { + "contractVersion": CONTRACT_VERSION, + "jobKey": job_key, + "exitCode": outcome.get("ExitCode", 1), + "error": outcome.get("Error"), + "unexpected": bool(outcome.get("Unexpected")), + # state.db is never carried: agent code opens that path directly. + "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( + r"^\[(?P[^\]]+)\]\[(?P\w+)\] (?P.*)$" +) + + +def parse_log_line(line: str) -> dict[str, Any]: + match = LOG_LINE_RE.match(line) + if not match: + return {"timestamp": None, "level": None, "message": line} + return { + "timestamp": match.group("timestamp"), + "level": match.group("level"), + "message": match.group("message"), + } + + +class JobLogTailer: + """Follows the job's log file and pushes batches to the caller. + + The runtime's log interceptor strips every handler but its own, so there is no + supported seam to attach a second one — and ``uipath-runtime`` ships on its own + release train. Tailing the file it already writes keeps this additive: the file is + still produced exactly as before, we just also forward it. + """ + + def __init__(self, job_key: str, path: str, callback: "HandlerCallback") -> None: + self.job_key = job_key + self.path = path + self.callback = callback + self._offset = 0 + self._stop = asyncio.Event() + + async def run(self) -> None: + try: + while not self._stop.is_set(): + await self._drain() + try: + await asyncio.wait_for(self._stop.wait(), timeout=LOG_POLL_SECONDS) + except asyncio.TimeoutError: + pass + # The job has finished; pick up whatever it wrote on the way out, including + # any unterminated final line. + await self._drain(final=True) + except asyncio.CancelledError: + raise + except Exception as e: # logs must never take the job down with them + _server_log(f"Log tailing stopped for job {self.job_key}: {e}") + + def stop(self) -> None: + self._stop.set() + + async def _drain(self, final: bool = False) -> None: + try: + if not os.path.exists(self.path): + return + # Binary with explicit offsets: a text-mode tell() cookie is opaque, and we + # need to rewind past an unterminated tail. + with open(self.path, "rb") as f: + f.seek(self._offset) + chunk = f.read() + except OSError: + return + + if not chunk: + return + + # A logging handler writes the record and only then flushes, so the tail can be + # half a line. Leave it for the next poll rather than reporting a fragment as a + # complete entry — except on the final drain, where nothing more is coming. + if not final and not chunk.endswith(b"\n"): + cut = chunk.rfind(b"\n") + if cut == -1: + return + chunk = chunk[: cut + 1] + + self._offset += len(chunk) + + batch: list[dict[str, Any]] = [] + for raw in chunk.decode("utf-8", errors="replace").splitlines(): + line = raw.rstrip("\r") + if not line: + continue + batch.append(parse_log_line(line)) + if len(batch) >= LOG_BATCH_MAX_LINES: + await self.callback.post_logs(self.job_key, batch) + batch = [] + + if batch: + await self.callback.post_logs(self.job_key, batch) + + +async def _stop_tailer( + tailer: JobLogTailer | None, task: "asyncio.Task[None] | None" +) -> None: + """Let the tailer drain what the job just wrote, then make sure it is gone.""" + if task is None: + return + if tailer is not None: + tailer.stop() + try: + # Shielded so a cancellation of the job task still allows the final drain. + await asyncio.wait_for(asyncio.shield(task), LOG_FLUSH_TIMEOUT_SECONDS) + except asyncio.CancelledError: + # Aimed at US, not the tailer. Tidy up, then let it propagate — swallowing it + # would carry on into the multi-minute result retry and ignore the shutdown. + task.cancel() + with contextlib.suppress(BaseException): + await task + raise + except (asyncio.TimeoutError, Exception): + task.cancel() + with contextlib.suppress(BaseException): + await task + + +class JobRegistry: + """Tracks in-flight jobs so they can be deduplicated and cancelled.""" + + 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) + return task is not None and not task.done() + + def start( + self, + job_key: str, + cmd: Any, + args: list[str], + env_vars: dict[str, str], + working_dir: str | None, + callback: HandlerCallback, + 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) + ) + self._tasks[job_key] = task + task.add_done_callback(lambda _: self._forget(job_key)) + return True + + 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, + job_key: str, + cmd: Any, + args: list[str], + env_vars: dict[str, str], + working_dir: str | None, + callback: HandlerCallback, + ) -> None: + tailer: JobLogTailer | None = None + tail_task: "asyncio.Task[None] | None" = None + + async def finish(outcome: dict[str, Any]) -> None: + # Flush logs BEFORE the terminal result: the result is what completes the + # job on the caller's side, and a log line arriving after that has nowhere + # left to go. + await _stop_tailer(tailer, tail_task) + delivered = await callback.post_result( + job_key, build_result_payload(job_key, outcome) + ) + if not delivered: + # StartJob no longer blocks, so nothing on the caller's side will ever + # time this job out — it would hang forever, and so would every other + # job we are holding, since none of them can be reported either. + # Exiting hands them all to the caller's service-exit path, which + # completes them from the result files the runtime already wrote. + request_shutdown(f"result callback unreachable for job {job_key}") + + try: + # Inside the try: anything that throws here must still produce a terminal + # report, or the caller waits forever on a job it believes was accepted. + logs_path = resolve_logs_file_path(env_vars, working_dir) + if logs_path: + 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, + # otherwise the caller waits forever for a job that will never run. + await finish({"ExitCode": 1, "Error": "Job cancelled before execution"}) + raise + except Exception as e: + await finish({"ExitCode": 1, "Error": str(e), "Unexpected": True}) + return + + 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. + """ + task = self._tasks.get(job_key) + if task is None or task.done(): + return True + + # 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 + + +_registry = JobRegistry() + + +def get_registry() -> JobRegistry: + return _registry 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 6c020dca7..89c9baa8b 100644 --- a/packages/uipath/src/uipath/_cli/cli_server.py +++ b/packages/uipath/src/uipath/_cli/cli_server.py @@ -18,6 +18,13 @@ _state, parse_args, ) +from ._server_jobs import ( + CONTRACT_VERSION, + EXIT_CALLBACK_UNREACHABLE, + HandlerCallback, + get_registry, + shutdown_event, +) from ._telemetry import track_command from ._utils._console import ConsoleLogger from .cli_server_ipc import ( @@ -110,9 +117,14 @@ def get_field(message: dict[str, Any], *keys: str) -> Any: async def send_ack(ack_socket_path: str, server_socket_path: str) -> None: """Send acknowledgment via HTTP POST to the ack socket.""" - ack_message: dict[str, str] = { + # capabilities lets the caller know what this runtime can do before it dispatches + # anything — notably whether it will push logs, which decides whether the caller + # tails the log file itself. Unknown keys are ignored by older callers. + ack_message: dict[str, Any] = { "status": "ready", "socket": server_socket_path, + "protocolVersion": CONTRACT_VERSION, + "capabilities": ["asyncStart", "resultPush", "logPush"], } conn = UnixConnector(path=ack_socket_path) @@ -184,6 +196,41 @@ async def handle_start(request: web.Request) -> web.Response: console.info(f"Starting job {job_key}: {command_name} {args}") + # A caller that gave us somewhere to report to gets async dispatch; one that did + # not (an older handler) gets the original blocking call, unchanged. + result_callback_socket = get_field( + message, "resultCallbackSocket", "ResultCallbackSocket" + ) + if isinstance(result_callback_socket, str) and result_callback_socket: + accepted = get_registry().start( + job_key, + cmd, + args, + env_vars, + working_dir, + HandlerCallback(result_callback_socket), + resume_version=get_field(message, "resumeVersion", "ResumeVersion"), + ) + if not accepted: + return web.json_response( + { + "success": False, + "job_key": job_key, + "contractVersion": CONTRACT_VERSION, + "error": f"Job {job_key} is already in flight", + }, + status=409, + ) + return web.json_response( + { + "success": True, + "job_key": job_key, + "contractVersion": CONTRACT_VERSION, + "disposition": "accepted", + }, + status=202, + ) + result = await _run_command_isolated(cmd, args, env_vars, working_dir) if result["Unexpected"]: @@ -197,12 +244,55 @@ async def handle_start(request: web.Request) -> web.Response: {"success": False, "job_key": job_key, "error": result["Error"]}, status=400, ) + # exitCode is additive and redundant with success, but it is what handlers built + # against the original response shape actually read — emitting it makes them + # correct without an upgrade. if result["ExitCode"] == 0: return web.json_response( - {"success": True, "job_key": job_key, "result": result["Result"]} + { + "success": True, + "job_key": job_key, + "exitCode": 0, + "result": result["Result"], + } + ) + return web.json_response( + { + "success": False, + "job_key": job_key, + "exitCode": result["ExitCode"], + "error": result["Error"], + } + ) + + +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": False, "job_key": job_key, "error": result["Error"]} + { + "success": True, + "job_key": job_key, + "contractVersion": CONTRACT_VERSION, + "stopped": stopped, + } ) @@ -240,6 +330,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 @@ -371,7 +462,26 @@ async def _serve( if ipc_pipe: tasks.append(start_ipc_server(ipc_pipe)) - await asyncio.gather(*tasks) + channels = [asyncio.ensure_future(t) for t in tasks] + stopper = asyncio.ensure_future(shutdown_event().wait()) + + done, pending = await asyncio.wait( + {*channels, stopper}, return_when=asyncio.FIRST_COMPLETED + ) + + for task in pending: + task.cancel() + + if stopper in done: + console.warning( + "Stopping the server: the result callback is unreachable, so no job can be " + "reported. The caller completes pending jobs from their result files." + ) + raise SystemExit(EXIT_CALLBACK_UNREACHABLE) + + # A channel returning or throwing on its own is still fatal; surface it. + for task in done: + task.result() def _run_server( diff --git a/packages/uipath/src/uipath/_cli/cli_server_ipc.py b/packages/uipath/src/uipath/_cli/cli_server_ipc.py index 24ebd100c..63d6f73d4 100644 --- a/packages/uipath/src/uipath/_cli/cli_server_ipc.py +++ b/packages/uipath/src/uipath/_cli/cli_server_ipc.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from ._server_core import COMMANDS, _run_command_isolated, _state, parse_args +from ._server_jobs import CONTRACT_VERSION, HandlerCallback, get_registry from ._utils._console import ConsoleLogger console = ConsoleLogger() @@ -30,6 +31,11 @@ class PythonRunRequest: Args: str | list[str] | None = None WorkingDirectory: str | None = None EnvironmentVariables: dict[str, str] = field(default_factory=dict) + # Where to POST the terminal result. Absent ⇒ the caller predates async dispatch, + # so run inline and return the outcome on this call, exactly as before. + ResultCallbackSocket: str | None = None + ContractVersion: int = 0 + ResumeVersion: int | None = None @dataclass @@ -38,6 +44,10 @@ class PythonRunResult: ExitCode: int = 0 Error: str | None = None + # "accepted" ⇒ the job was queued and its result will be pushed. Absent/"completed" + # ⇒ ExitCode is terminal (the legacy contract). + Disposition: str | None = None + ContractVersion: int = 0 class IPythonRuntimeServer(ABC): @@ -48,8 +58,13 @@ async def StartJob(self, request: PythonRunRequest) -> PythonRunResult: """Run a job → PythonRunResult(ExitCode, Error).""" @abstractmethod - async def StopJob(self, job_key: str) -> bool: - """Cancel a running job by key (bool return avoids fire-and-forget).""" + async def StopJob(self, job_key: str, resume_version: int | None = None) -> bool: + """Cancel a running job by key (bool return avoids fire-and-forget). + + ``resume_version`` is a TRAILING OPTIONAL parameter rather than a DTO: uipath-ipc + ignores a surplus wire argument and defaults a missing one, so an old handler and + a new venv interoperate in both directions. A DTO would break the old pairing. + """ class PythonRuntimeService(IPythonRuntimeServer): @@ -70,15 +85,37 @@ async def StartJob(self, request: PythonRunRequest) -> PythonRunResult: console.info(f"Starting job {request.JobKey}: {command_name} {args}") + callback_socket = request.ResultCallbackSocket + if callback_socket: + accepted = get_registry().start( + request.JobKey, + cmd, + args, + request.EnvironmentVariables, + request.WorkingDirectory, + HandlerCallback(callback_socket), + resume_version=request.ResumeVersion, + ) + if not accepted: + return PythonRunResult( + ExitCode=1, + Error=f"Job {request.JobKey} is already in flight", + ContractVersion=CONTRACT_VERSION, + ) + return PythonRunResult( + ExitCode=0, + Disposition="accepted", + ContractVersion=CONTRACT_VERSION, + ) + result = await _run_command_isolated( cmd, args, request.EnvironmentVariables, request.WorkingDirectory ) - # IPC contract (PythonRunResult) carries only ExitCode + Error. + # Legacy inline contract: ExitCode + Error only. return PythonRunResult(ExitCode=result["ExitCode"], Error=result["Error"]) - async def StopJob(self, job_key: str) -> bool: - console.info(f"StopJob requested for {job_key} (no-op)") - return True + async def StopJob(self, job_key: str, resume_version: int | None = None) -> bool: + return await get_registry().stop(job_key, resume_version) async def start_ipc_server(pipe_name: str) -> None: diff --git a/packages/uipath/tests/cli/test_server.py b/packages/uipath/tests/cli/test_server.py index 70d29cb38..86e88102c 100644 --- a/packages/uipath/tests/cli/test_server.py +++ b/packages/uipath/tests/cli/test_server.py @@ -36,6 +36,20 @@ def main(input: Input) -> str: """ +@pytest.fixture +def failing_script() -> str: + return """ +from dataclasses import dataclass + +@dataclass +class Input: + message: str = "" + +def main(input: Input) -> str: + raise RuntimeError("agent blew up") +""" + + def create_uipath_json(script_path: str, entrypoint_name: str = "main"): """Helper to create uipath.json with functions.""" return {"functions": {entrypoint_name: f"{script_path}:main"}} @@ -131,6 +145,7 @@ def test_start_job_success(self, server, temp_dir, simple_script): ) assert response["success"] is True + assert response["exitCode"] == 0 assert response["job_key"] == job_key assert os.path.exists(output_file) @@ -138,6 +153,52 @@ def test_start_job_success(self, server, temp_dir, simple_script): output = f.read() assert "Hello" in output + def test_start_job_that_faults_is_not_reported_successful( + self, server, temp_dir, failing_script + ): + """A job that runs and faults must say so in the body. + + The response is still 200 — the request succeeded, the job did not — so the + body is the only signal. Click returns ctx.exit(1) rather than raising under + standalone_mode=False, so an exit code that is not threaded through lands here + as a successful job. + """ + port = server + job_key = "test-job-faulted" + + with pytest.MonkeyPatch().context() as mp: + mp.chdir(temp_dir) + + script_file = "entrypoint.py" + with open(os.path.join(temp_dir, script_file), "w") as f: + f.write(failing_script) + + with open(os.path.join(temp_dir, "uipath.json"), "w") as f: + json.dump(create_uipath_json(script_file), f) + + input_file = os.path.join(temp_dir, "input.json") + with open(input_file, "w") as f: + json.dump({"message": "Hello"}, f) + + response = asyncio.run( + start_job( + port, + job_key, + "run", + [ + "main", + "--input-file", + input_file, + "--output-file", + os.path.join(temp_dir, "output.json"), + ], + ) + ) + + assert response["success"] is False + assert response["exitCode"] != 0 + assert response["job_key"] == job_key + def test_start_job_unknown_command(self, server): """Test starting a job with unknown command.""" port = server diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py new file mode 100644 index 000000000..5c20dce08 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_async.py @@ -0,0 +1,553 @@ +"""Async dispatch: enqueue, run, push the outcome back. + +The contract is opt-in per request: a caller that supplies a result-callback socket +gets async dispatch, one that does not gets the original blocking call. These tests +pin both halves, because the blocking half is what every un-upgraded caller still uses. +""" + +import asyncio +import json +import os +import time +from typing import Any + +import click +import pytest + +from uipath._cli import _server_core, cli_server, cli_server_ipc +from uipath._cli._server_core import ( + MAX_INLINE_DOCUMENT_BYTES, + _read_result_document, + _ServerState, + resolve_result_file_path, +) +from uipath._cli._server_jobs import ( + CONTRACT_VERSION, + JobRegistry, + build_result_payload, +) + + +@pytest.fixture(autouse=True) +def _fresh_state(monkeypatch): + monkeypatch.setattr(_server_core, "_state", _ServerState()) + + +class FakeCallback: + """Records result pushes instead of posting them.""" + + 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 + + +def _write_config(tmp_path, runtime: dict[str, Any] | None) -> str: + config = {"runtime": runtime} if runtime is not None else {} + path = tmp_path / "uipath.json" + path.write_text(json.dumps(config), encoding="utf-8") + return str(path) + + +# --------------------------------------------------------------------------- # +# result document resolution # +# --------------------------------------------------------------------------- # + + +def test_resolve_result_file_path_uses_runtime_config(tmp_path, monkeypatch): + monkeypatch.setenv( + "UIPATH_CONFIG_PATH", + _write_config( + tmp_path, {"dir": str(tmp_path / "__uipath"), "outputFile": "output.json"} + ), + ) + + resolved = resolve_result_file_path() + + assert resolved == os.path.abspath(str(tmp_path / "__uipath" / "output.json")) + + +def test_resolve_result_file_path_defaults_when_runtime_block_absent( + tmp_path, monkeypatch +): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("UIPATH_CONFIG_PATH", _write_config(tmp_path, None)) + + resolved = resolve_result_file_path() + + assert resolved == os.path.abspath(os.path.join("__uipath", "output.json")) + + +def test_resolve_result_file_path_falls_back_to_defaults_when_config_missing( + tmp_path, monkeypatch +): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("UIPATH_CONFIG_PATH", str(tmp_path / "nope.json")) + + resolved = resolve_result_file_path() + + assert resolved == os.path.abspath(os.path.join("__uipath", "output.json")) + + +def test_read_result_document_inline(tmp_path, monkeypatch): + runtime_dir = tmp_path / "__uipath" + runtime_dir.mkdir() + (runtime_dir / "output.json").write_text( + '{"status":"successful"}', encoding="utf-8" + ) + monkeypatch.setenv( + "UIPATH_CONFIG_PATH", + _write_config(tmp_path, {"dir": str(runtime_dir), "outputFile": "output.json"}), + ) + + document, conveyance = _read_result_document() + + assert conveyance == "inline" + assert document == '{"status":"successful"}' + + +def test_read_result_document_falls_back_to_file_when_absent(tmp_path, monkeypatch): + monkeypatch.setenv( + "UIPATH_CONFIG_PATH", + _write_config( + tmp_path, {"dir": str(tmp_path / "__uipath"), "outputFile": "output.json"} + ), + ) + + document, conveyance = _read_result_document() + + assert (document, conveyance) == (None, "file") + + +def test_read_result_document_falls_back_to_file_when_oversized(tmp_path, monkeypatch): + """Large outputs already work via the file; they must not start failing on the wire.""" + runtime_dir = tmp_path / "__uipath" + runtime_dir.mkdir() + (runtime_dir / "output.json").write_text( + "x" * (MAX_INLINE_DOCUMENT_BYTES + 1), encoding="utf-8" + ) + monkeypatch.setenv( + "UIPATH_CONFIG_PATH", + _write_config(tmp_path, {"dir": str(runtime_dir), "outputFile": "output.json"}), + ) + + document, conveyance = _read_result_document() + + assert (document, conveyance) == (None, "file") + + +# --------------------------------------------------------------------------- # +# payload shape # +# --------------------------------------------------------------------------- # + + +def test_build_result_payload_shape(): + payload = build_result_payload( + "job-1", + { + "ExitCode": 0, + "Error": None, + "Unexpected": False, + "Document": '{"status":"successful"}', + "DocumentConveyance": "inline", + }, + ) + + assert payload == { + "contractVersion": CONTRACT_VERSION, + "jobKey": "job-1", + "exitCode": 0, + "error": None, + "unexpected": False, + "stateConveyance": "file", + "jobConveyance": "inline", + "job": '{"status":"successful"}', + "stopped": False, + } + + +def test_build_result_payload_defaults_to_failure_when_outcome_is_bare(): + payload = build_result_payload("job-1", {}) + + assert payload["exitCode"] == 1 + assert payload["jobConveyance"] == "file" + assert payload["job"] is None + + +# --------------------------------------------------------------------------- # +# registry # +# --------------------------------------------------------------------------- # + + +@click.command() +def _ok_command() -> None: + return None + + +@click.command() +def _failing_command() -> None: + click.get_current_context().exit(3) + + +async def test_registry_runs_job_and_pushes_result(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + assert registry.start("job-1", _ok_command, [], {}, str(tmp_path), callback) + + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert len(callback.results) == 1 + assert callback.results[0]["jobKey"] == "job-1" + assert callback.results[0]["exitCode"] == 0 + + +async def test_registry_pushes_failure_exit_code(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _failing_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert callback.results[0]["exitCode"] == 3 + + +async def test_registry_rejects_a_duplicate_job_key(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + assert registry.start("job-1", _ok_command, [], {}, str(tmp_path), callback) is True + # Second start while the first is in flight must not silently replace it. + second = registry.start("job-1", _ok_command, [], {}, str(tmp_path), callback) + + await asyncio.wait_for(callback.done.wait(), timeout=10) + assert second is False + + +async def test_registry_frees_the_key_after_completion(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _ok_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + await asyncio.sleep(0) # let the done-callback run + + assert registry.is_active("job-1") is False + + +async def test_stop_is_true_for_an_unknown_job(): + registry = JobRegistry() + + 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 # +# --------------------------------------------------------------------------- # + + +class _FakeRequest: + def __init__(self, job_key: str, payload: dict[str, Any]) -> None: + self.match_info = {"job_key": job_key} + self._payload = payload + + async def json(self) -> dict[str, Any]: + return self._payload + + +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, resume_version=None + ): + self.started.append(job_key) + self.resume_versions.append(resume_version) + return self.accept + + +async def test_http_start_with_callback_returns_accepted(monkeypatch): + registry = _RecordingRegistry() + monkeypatch.setattr(cli_server, "get_registry", lambda: registry) + + response = await cli_server.handle_start( + _FakeRequest( + "job-1", + {"command": "run", "resultCallbackSocket": "/tmp/ack.sock"}, + ) + ) + + assert response.status == 202 + body = json.loads(response.text) + assert body["disposition"] == "accepted" + assert body["contractVersion"] == CONTRACT_VERSION + assert registry.started == ["job-1"] + + +async def test_http_start_duplicate_is_409(monkeypatch): + monkeypatch.setattr( + cli_server, "get_registry", lambda: _RecordingRegistry(accept=False) + ) + + response = await cli_server.handle_start( + _FakeRequest( + "job-1", {"command": "run", "resultCallbackSocket": "/tmp/ack.sock"} + ) + ) + + assert response.status == 409 + assert json.loads(response.text)["success"] is False + + +async def test_http_start_without_callback_stays_synchronous(monkeypatch): + """An un-upgraded caller must get the original blocking behaviour, byte for byte.""" + called: dict[str, Any] = {} + + async def _fake_isolated(cmd, args, env_vars, working_dir): + called["ran"] = True + return {"ExitCode": 0, "Error": None, "Result": {"out": 1}, "Unexpected": False} + + monkeypatch.setattr(cli_server, "_run_command_isolated", _fake_isolated) + monkeypatch.setattr( + cli_server, + "get_registry", + lambda: pytest.fail("registry must not be used without a callback socket"), + ) + + response = await cli_server.handle_start(_FakeRequest("job-1", {"command": "run"})) + + assert called.get("ran") is True + assert response.status == 200 + body = json.loads(response.text) + assert body["success"] is True + assert "disposition" not in body + + +# --------------------------------------------------------------------------- # +# IPC dispatch # +# --------------------------------------------------------------------------- # + + +async def test_ipc_start_with_callback_returns_accepted(monkeypatch): + registry = _RecordingRegistry() + monkeypatch.setattr(cli_server_ipc, "get_registry", lambda: registry) + + result = await cli_server_ipc.PythonRuntimeService().StartJob( + cli_server_ipc.PythonRunRequest( + JobKey="job-1", Command="run", ResultCallbackSocket="/tmp/ack.sock" + ) + ) + + assert result.Disposition == "accepted" + assert result.ExitCode == 0 + assert result.ContractVersion == CONTRACT_VERSION + assert registry.started == ["job-1"] + + +async def test_ipc_start_duplicate_reports_failure(monkeypatch): + monkeypatch.setattr( + cli_server_ipc, "get_registry", lambda: _RecordingRegistry(accept=False) + ) + + result = await cli_server_ipc.PythonRuntimeService().StartJob( + cli_server_ipc.PythonRunRequest( + JobKey="job-1", Command="run", ResultCallbackSocket="/tmp/ack.sock" + ) + ) + + assert result.ExitCode == 1 + assert result.Disposition is None + assert "already in flight" in result.Error + + +async def test_ipc_start_without_callback_stays_synchronous(monkeypatch): + async def _fake_isolated(cmd, args, env_vars, working_dir): + return { + "ExitCode": 7, + "Error": "Exit code: 7", + "Result": None, + "Unexpected": False, + } + + monkeypatch.setattr(cli_server_ipc, "_run_command_isolated", _fake_isolated) + + result = await cli_server_ipc.PythonRuntimeService().StartJob( + cli_server_ipc.PythonRunRequest(JobKey="job-1", Command="run") + ) + + 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 diff --git a/packages/uipath/tests/cli/test_server_callbacks.py b/packages/uipath/tests/cli/test_server_callbacks.py new file mode 100644 index 000000000..f2015dd73 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_callbacks.py @@ -0,0 +1,227 @@ +"""The callback URLs and payload keys the .NET handler parses. + +The handler declares these routes as its own string constants +(``Constants.PythonRuntimeContract`` in GenericExecutors.PythonCoded, pinned by +``PythonRuntimeContractRoutesTests``). Nothing at build time ties the two sides together, +so a rename on either side is a 404 on every push — which this side treats as retryable, +leaving the job to never complete. These assertions make such a rename a deliberate, +visible change with a failing test on both sides. +""" + +from typing import Any + +import pytest + +from uipath._cli._server_jobs import ( + CONTRACT_VERSION, + HandlerCallback, + build_result_payload, +) + + +class _CapturingCallback(HandlerCallback): + """Captures what would be POSTed instead of opening a socket.""" + + def __init__(self) -> None: + super().__init__("/tmp/never-connected.sock") + self.posts: list[tuple[str, dict[str, Any]]] = [] + self.response = True + + async def _post(self, path: str, payload: dict[str, Any]): + from uipath._cli._server_jobs import _PostOutcome + + self.posts.append((path, payload)) + return _PostOutcome.DELIVERED if self.response else _PostOutcome.UNREACHABLE + + +async def test_result_is_posted_to_the_route_the_handler_serves(): + callback = _CapturingCallback() + + await callback.post_result("abc-123", {"exitCode": 0}) + + path, _ = callback.posts[0] + assert path == "/api/python/jobs/abc-123/result" + + +async def test_logs_are_posted_to_the_route_the_handler_serves(): + callback = _CapturingCallback() + + await callback.post_logs("abc-123", [{"message": "hello"}]) + + path, _ = callback.posts[0] + assert path == "/api/python/jobs/abc-123/logs" + + +async def test_log_payload_carries_the_keys_the_handler_deserializes(): + """RuntimeJobLogsDto maps contractVersion / jobKey / lines[].{timestamp,level,message}.""" + callback = _CapturingCallback() + + await callback.post_logs( + "abc-123", + [{"timestamp": "2026-07-29 17:00:00,000", "level": "INFO", "message": "hi"}], + ) + + _, payload = callback.posts[0] + assert payload["contractVersion"] == CONTRACT_VERSION + assert payload["jobKey"] == "abc-123" + assert payload["lines"] == [ + {"timestamp": "2026-07-29 17:00:00,000", "level": "INFO", "message": "hi"} + ] + + +def test_result_payload_carries_the_keys_the_handler_deserializes(): + """RuntimeJobResultDto maps exactly these camelCase keys.""" + payload = build_result_payload("abc-123", {"ExitCode": 0, "Document": "{}"}) + + assert set(payload) == { + "contractVersion", + "jobKey", + "exitCode", + "error", + "unexpected", + "stateConveyance", + "jobConveyance", + "job", + "stopped", + } + + +async def test_result_push_retries_until_it_lands(): + """A briefly unavailable callback socket must not strand the job.""" + callback = _CapturingCallback() + callback.response = False + + attempts: list[int] = [] + + from uipath._cli._server_jobs import _PostOutcome + + async def _flaky(path: str, payload: dict[str, Any]) -> _PostOutcome: + attempts.append(1) + return ( + _PostOutcome.DELIVERED if len(attempts) >= 3 else _PostOutcome.UNREACHABLE + ) + + callback._post = _flaky # type: ignore[method-assign] + + # Collapse the backoff so the test does not actually wait minutes. + import uipath._cli._server_jobs as jobs + + original = jobs.RESULT_PUSH_BACKOFF_SECONDS + jobs.RESULT_PUSH_BACKOFF_SECONDS = 0.0 + try: + assert await callback.post_result("abc-123", {"exitCode": 0}) is True + finally: + jobs.RESULT_PUSH_BACKOFF_SECONDS = original + + assert len(attempts) == 3 + + +async def test_result_push_eventually_gives_up_without_raising(): + """Giving up must be a warning, not an exception: this runs in a background task and + ConsoleLogger.error is NoReturn (it calls ctx.exit).""" + callback = _CapturingCallback() + callback.response = False + + import uipath._cli._server_jobs as jobs + + original_backoff = jobs.RESULT_PUSH_BACKOFF_SECONDS + original_attempts = jobs.RESULT_PUSH_ATTEMPTS + jobs.RESULT_PUSH_BACKOFF_SECONDS = 0.0 + jobs.RESULT_PUSH_ATTEMPTS = 3 + try: + result = await callback.post_result("abc-123", {"exitCode": 0}) + finally: + jobs.RESULT_PUSH_BACKOFF_SECONDS = original_backoff + jobs.RESULT_PUSH_ATTEMPTS = original_attempts + + assert result is False + assert len(callback.posts) == 3 + + +async def test_undeliverable_result_asks_the_server_to_stop(tmp_path): + """If the callback socket is unreachable, no job we hold can ever be reported and + StartJob no longer blocks — so every one of them would hang. Exiting hands them to + the caller's service-exit path, which completes them from their result files.""" + import click + + from uipath._cli import _server_core + from uipath._cli._server_jobs import JobRegistry, shutdown_event + + @click.command() + def _ok() -> None: + return None + + class _DeadCallback(HandlerCallback): + def __init__(self) -> None: + super().__init__("/tmp/unreachable.sock") + self.done = __import__("asyncio").Event() + + async def post_result(self, job_key, payload): + self.done.set() + return False + + async def post_logs(self, job_key, lines): + return False + + import asyncio + + _server_core._state.init() + assert not shutdown_event().is_set() + + callback = _DeadCallback() + JobRegistry().start("job-1", _ok, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + await asyncio.sleep(0) + + assert shutdown_event().is_set() + + +async def test_a_delivered_result_does_not_stop_the_server(tmp_path): + import asyncio + + import click + + from uipath._cli import _server_core + from uipath._cli._server_jobs import JobRegistry, shutdown_event + + @click.command() + def _ok() -> None: + return None + + class _LiveCallback(HandlerCallback): + def __init__(self) -> None: + super().__init__("/tmp/live.sock") + self.done = asyncio.Event() + + async def post_result(self, job_key, payload): + self.done.set() + return True + + async def post_logs(self, job_key, lines): + return True + + _server_core._state.init() + callback = _LiveCallback() + JobRegistry().start("job-1", _ok, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + await asyncio.sleep(0) + + assert not shutdown_event().is_set() + + +async def test_logs_are_not_retried(): + """A dropped log batch must never delay the job.""" + callback = _CapturingCallback() + callback.response = False + + assert await callback.post_logs("abc-123", [{"message": "x"}]) is False + assert len(callback.posts) == 1 + + +@pytest.mark.parametrize("job_key", ["abc-123", "00000000-0000-0000-0000-000000000000"]) +async def test_route_is_built_from_the_job_key_verbatim(job_key): + callback = _CapturingCallback() + + await callback.post_result(job_key, {}) + + assert callback.posts[0][0] == f"/api/python/jobs/{job_key}/result" 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_logs.py b/packages/uipath/tests/cli/test_server_logs.py new file mode 100644 index 000000000..3b8094765 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_logs.py @@ -0,0 +1,296 @@ +"""Forwarding job logs to the caller while the job runs. + +The runtime's log interceptor strips every handler but its own, so the log file it +already writes is the only supported source. These pin the tailing, the parsing of the +runtime's line format, and the ordering guarantee that matters: the last log batch is +flushed before the terminal result, because the result is what ends the job. +""" + +import asyncio +import json +import os +from typing import Any + +import pytest + +from uipath._cli import _server_core +from uipath._cli._server_core import _ServerState, resolve_logs_file_path +from uipath._cli._server_jobs import ( + LOG_BATCH_MAX_LINES, + JobLogTailer, + parse_log_line, +) + + +@pytest.fixture(autouse=True) +def _fresh_state(monkeypatch): + monkeypatch.setattr(_server_core, "_state", _ServerState()) + + +class RecordingCallback: + def __init__(self) -> None: + self.batches: list[list[dict[str, Any]]] = [] + self.results: list[dict[str, Any]] = [] + self.order: list[str] = [] + + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: + self.batches.append(lines) + self.order.append("logs") + return True + + async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: + self.results.append(payload) + self.order.append("result") + return True + + @property + def lines(self) -> list[dict[str, Any]]: + return [line for batch in self.batches for line in batch] + + +# --------------------------------------------------------------------------- # +# line parsing # +# --------------------------------------------------------------------------- # + + +def test_parse_log_line_splits_the_runtime_format(): + parsed = parse_log_line("[2026-07-29 17:04:19,123][INFO] hello world") + + assert parsed == { + "timestamp": "2026-07-29 17:04:19,123", + "level": "INFO", + "message": "hello world", + } + + +def test_parse_log_line_passes_through_an_unstructured_line(): + """Tracebacks and bare prints have no prefix; they must survive verbatim.""" + parsed = parse_log_line(' File "agent.py", line 3, in main') + + assert parsed["timestamp"] is None + assert parsed["level"] is None + assert parsed["message"] == ' File "agent.py", line 3, in main' + + +def test_parse_log_line_keeps_brackets_inside_the_message(): + parsed = parse_log_line("[2026-07-29 17:04:19,123][ERROR] failed [attempt 2]") + + assert parsed["level"] == "ERROR" + assert parsed["message"] == "failed [attempt 2]" + + +# --------------------------------------------------------------------------- # +# path resolution # +# --------------------------------------------------------------------------- # + + +def test_resolve_logs_file_path_is_pure_wrt_process_state(tmp_path): + """The tailer needs the path before the job applies its env, so this must not + depend on os.environ or the cwd.""" + config = tmp_path / "uipath.json" + config.write_text( + json.dumps({"runtime": {"dir": "__uipath", "logsFile": "execution.log"}}), + encoding="utf-8", + ) + + resolved = resolve_logs_file_path( + {"UIPATH_CONFIG_PATH": "uipath.json"}, str(tmp_path) + ) + + assert resolved == os.path.abspath(str(tmp_path / "__uipath" / "execution.log")) + + +def test_resolve_logs_file_path_handles_an_absolute_runtime_dir(tmp_path): + runtime_dir = tmp_path / "elsewhere" + config = tmp_path / "uipath.json" + config.write_text( + json.dumps({"runtime": {"dir": str(runtime_dir), "logsFile": "execution.log"}}), + encoding="utf-8", + ) + + resolved = resolve_logs_file_path( + {"UIPATH_CONFIG_PATH": str(config)}, str(tmp_path) + ) + + assert resolved == os.path.abspath(str(runtime_dir / "execution.log")) + + +def test_resolve_logs_file_path_falls_back_to_the_default_without_a_config(tmp_path): + """Must never return None: the caller has already stopped tailing the file itself, + so a job with an unreadable config would otherwise produce no logs anywhere.""" + resolved = resolve_logs_file_path({}, str(tmp_path)) + + assert resolved == os.path.abspath(str(tmp_path / "__uipath" / "execution.log")) + + +# --------------------------------------------------------------------------- # +# tailing # +# --------------------------------------------------------------------------- # + + +async def test_tailer_forwards_lines_written_while_running(tmp_path): + log_file = tmp_path / "execution.log" + log_file.write_text("[2026-07-29 17:00:00,000][INFO] first\n", encoding="utf-8") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + task = asyncio.create_task(tailer.run()) + + await asyncio.sleep(0.4) + with open(log_file, "a", encoding="utf-8") as f: + f.write("[2026-07-29 17:00:01,000][ERROR] second\n") + await asyncio.sleep(0.4) + + tailer.stop() + await asyncio.wait_for(task, timeout=5) + + messages = [line["message"] for line in callback.lines] + assert messages == ["first", "second"] + + +async def test_tailer_never_replays_a_line(tmp_path): + log_file = tmp_path / "execution.log" + log_file.write_text("[2026-07-29 17:00:00,000][INFO] once\n", encoding="utf-8") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + task = asyncio.create_task(tailer.run()) + + await asyncio.sleep(0.8) # several poll cycles over an unchanged file + tailer.stop() + await asyncio.wait_for(task, timeout=5) + + assert [line["message"] for line in callback.lines] == ["once"] + + +async def test_tailer_drains_what_was_written_after_stop_was_requested(tmp_path): + """The last lines a job writes land after it finishes; they must not be lost.""" + log_file = tmp_path / "execution.log" + log_file.write_text("", encoding="utf-8") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + task = asyncio.create_task(tailer.run()) + await asyncio.sleep(0.1) + + with open(log_file, "a", encoding="utf-8") as f: + f.write("[2026-07-29 17:00:02,000][INFO] final\n") + tailer.stop() + + await asyncio.wait_for(task, timeout=5) + + assert [line["message"] for line in callback.lines] == ["final"] + + +async def test_tailer_batches_long_output(tmp_path): + log_file = tmp_path / "execution.log" + total = LOG_BATCH_MAX_LINES + 25 + log_file.write_text( + "".join(f"[2026-07-29 17:00:00,000][INFO] line-{i}\n" for i in range(total)), + encoding="utf-8", + ) + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + tailer.stop() + await asyncio.wait_for(tailer.run(), timeout=5) + + assert len(callback.batches) >= 2 + assert len(callback.batches[0]) == LOG_BATCH_MAX_LINES + assert len(callback.lines) == total + + +async def test_logs_are_flushed_before_the_terminal_result(tmp_path): + """The result is what ends the job on the caller's side. Any log line pushed after + it is dropped by design, so the last batch MUST precede it.""" + import click + + from uipath._cli._server_jobs import JobRegistry + + runtime_dir = tmp_path / "__uipath" + runtime_dir.mkdir() + (tmp_path / "uipath.json").write_text( + json.dumps({"runtime": {"dir": "__uipath", "logsFile": "execution.log"}}), + encoding="utf-8", + ) + + log_path = runtime_dir / "execution.log" + + @click.command() + def _logging_command() -> None: + with open(log_path, "a", encoding="utf-8") as f: + f.write("[2026-07-29 17:00:00,000][INFO] from inside the job\n") + + _server_core._state.init() + callback = RecordingCallback() + registry = JobRegistry() + + registry.start( + "job-1", + _logging_command, + [], + {"UIPATH_CONFIG_PATH": str(tmp_path / "uipath.json")}, + str(tmp_path), + callback, + ) + + for _ in range(200): + if callback.results: + break + await asyncio.sleep(0.05) + + assert callback.results, "job never reported a result" + assert "logs" in callback.order, "the job's log line was never forwarded" + assert callback.order.index("logs") < callback.order.index("result") + assert callback.order[-1] == "result" + + +async def test_tailer_does_not_split_a_partially_written_line(tmp_path): + """A logging handler writes the record and only then flushes, so a poll can catch + half a line. Emitting the fragment produces two bogus entries that cannot be + rejoined — the tail has to wait for its newline.""" + log_file = tmp_path / "execution.log" + complete = b"[2026-07-29 17:00:00,000][INFO] complete" + b"\n" + log_file.write_bytes(complete + b"[2026-07-29 17:00:01,000][INFO] half") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + task = asyncio.create_task(tailer.run()) + await asyncio.sleep(0.4) + + assert [line["message"] for line in callback.lines] == ["complete"] + + with open(log_file, "ab") as f: + f.write(b"-and-half" + b"\n") + await asyncio.sleep(0.4) + + tailer.stop() + await asyncio.wait_for(task, timeout=5) + + assert [line["message"] for line in callback.lines] == ["complete", "half-and-half"] + + +async def test_tailer_emits_an_unterminated_final_line(tmp_path): + """On the last drain nothing more is coming, so a trailing fragment is real output + and must not be dropped.""" + log_file = tmp_path / "execution.log" + log_file.write_bytes(b"no trailing newline") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + tailer.stop() + + await asyncio.wait_for(tailer.run(), timeout=5) + + assert [line["message"] for line in callback.lines] == ["no trailing newline"] + + +async def test_tailer_tolerates_a_missing_file(tmp_path): + """A job that never logs must not produce an error or a stuck tailer.""" + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(tmp_path / "never-created.log"), callback) + tailer.stop() + + await asyncio.wait_for(tailer.run(), timeout=5) + + assert callback.batches == [] diff --git a/packages/uipath/tests/cli/test_server_result.py b/packages/uipath/tests/cli/test_server_result.py new file mode 100644 index 000000000..0e6f81bb8 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_result.py @@ -0,0 +1,187 @@ +"""The outcome a job reports back over both server channels. + +``_run_command_isolated`` is the single job core behind the HTTP and uipath-ipc +channels, so the exit code it produces is what BOTH wires report. Click's +``standalone_mode=False`` returns ``ctx.exit(N)``'s code instead of raising +``SystemExit`` — these tests pin that behaviour against real click commands +rather than a stub, because it is the whole reason the exit code can be wrong. +""" + +from typing import Any + +import click +import pytest +from aiohttp import web + +from uipath._cli import _server_core, cli_server +from uipath._cli._server_core import _run_command_isolated, _ServerState +from uipath._cli._utils._console import ConsoleLogger + + +@pytest.fixture(autouse=True) +def _fresh_state(monkeypatch): + """Give each test a state object whose lock binds to that test's event loop.""" + monkeypatch.setattr(_server_core, "_state", _ServerState()) + + +@click.command() +def _returns_object() -> dict[str, Any]: + return {"ok": True} + + +@click.command() +def _returns_none() -> None: + return None + + +@click.command() +def _exits_one() -> None: + click.get_current_context().exit(1) + + +@click.command() +def _exits_zero() -> None: + click.get_current_context().exit(0) + + +@click.command() +def _console_errors() -> None: + ConsoleLogger().error("boom") + + +@click.command() +def _raises() -> None: + raise RuntimeError("kaboom") + + +async def _run(cmd: Any) -> dict[str, Any]: + _server_core._state.init() + return await _run_command_isolated(cmd, [], {}, None) + + +# --------------------------------------------------------------------------- # +# exit-code normalisation # +# --------------------------------------------------------------------------- # + + +async def test_console_error_reports_a_failing_exit_code(): + """The regression that mattered: ConsoleLogger.error ends in ctx.exit(1), which + click RETURNS under standalone_mode=False — it must not read as success.""" + result = await _run(_console_errors) + + assert result["ExitCode"] == 1 + assert result["Error"] == "Exit code: 1" + assert result["Unexpected"] is False + + +async def test_ctx_exit_nonzero_becomes_the_exit_code(): + result = await _run(_exits_one) + + assert result["ExitCode"] == 1 + assert result["Result"] is None + + +async def test_ctx_exit_zero_is_success(): + result = await _run(_exits_zero) + + assert result["ExitCode"] == 0 + assert result["Error"] is None + + +async def test_object_return_is_a_result_not_an_exit_code(): + result = await _run(_returns_object) + + assert result["ExitCode"] == 0 + assert result["Error"] is None + assert result["Result"] == {"ok": True} + + +async def test_none_return_is_success(): + result = await _run(_returns_none) + + assert result["ExitCode"] == 0 + assert result["Result"] is None + + +async def test_unhandled_exception_is_flagged_unexpected(): + result = await _run(_raises) + + assert result["ExitCode"] == 1 + assert result["Unexpected"] is True + assert "kaboom" in result["Error"] + + +# --------------------------------------------------------------------------- # +# HTTP response envelope # +# --------------------------------------------------------------------------- # + + +class _FakeRequest: + """Minimal stand-in for web.Request: handle_start only uses match_info + json().""" + + def __init__(self, job_key: str, payload: dict[str, Any]) -> None: + self.match_info = {"job_key": job_key} + self._payload = payload + + async def json(self) -> dict[str, Any]: + return self._payload + + +async def _post_start(monkeypatch, command_result: dict[str, Any]) -> web.Response: + async def _fake_isolated(cmd, args, env_vars, working_dir): + return command_result + + monkeypatch.setattr(cli_server, "_run_command_isolated", _fake_isolated) + return await cli_server.handle_start(_FakeRequest("job-1", {"command": "run"})) + + +async def test_success_body_carries_exit_code(monkeypatch): + response = await _post_start( + monkeypatch, + {"ExitCode": 0, "Error": None, "Result": {"out": 1}, "Unexpected": False}, + ) + + assert response.status == 200 + assert response.text is not None + body = response.text + assert '"success": true' in body + assert '"exitCode": 0' in body + + +async def test_failure_body_is_200_but_says_so_and_carries_exit_code(monkeypatch): + """A failed job answers 200 with the outcome in the body — so the body must be + unambiguous. A handler reading only exitCode has to see a non-zero value.""" + response = await _post_start( + monkeypatch, + {"ExitCode": 1, "Error": "Exit code: 1", "Result": None, "Unexpected": False}, + ) + + assert response.status == 200 + body = response.text + assert '"success": false' in body + assert '"exitCode": 1' in body + assert "Exit code: 1" in body + + +async def test_unexpected_failure_is_500(monkeypatch): + response = await _post_start( + monkeypatch, + {"ExitCode": 1, "Error": "kaboom", "Result": None, "Unexpected": True}, + ) + + assert response.status == 500 + + +async def test_client_error_is_400(monkeypatch): + response = await _post_start( + monkeypatch, + { + "ExitCode": 1, + "Error": "Cannot change to working directory", + "Result": None, + "Unexpected": False, + "ClientError": True, + }, + ) + + assert response.status == 400