From 01c1fb1cd3dfd850caccb6a572e6d9cfe1083ff6 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 13:47:43 +0300 Subject: [PATCH 1/6] feat(cli): dispatch jobs asynchronously and push results/logs to the caller `uipath server` used to block for a whole job and leave its outcome on disk for the caller to find. StartJob now enqueues the work and returns; the server pushes logs while the job runs and the terminal result when it finishes. Behaviour is a pure function of the request: a caller that supplies `resultCallbackSocket` gets async dispatch, one that does not gets the original blocking call, byte for byte. No handshake, no capability gate on the dispatch path -- an older caller cannot send the field and could not serve the callback if it did. POST {callback}/api/python/jobs/{jobKey}/result POST {callback}/api/python/jobs/{jobKey}/logs The readiness ACK advertises protocolVersion + capabilities so the caller knows whether logs will be forwarded before it decides to tail the file itself. Execution stays serialised behind the process-wide lock: a job mutates process globals (logging handlers, OTel providers, env, cwd), so queueing changes who waits, not how many run. Also fixes two pre-existing defects that made the API result meaningless: _run_command_isolated hardcoded ExitCode 0 while click RETURNS ctx.exit(N)'s code under standalone_mode=False, so every ConsoleLogger.error path reported success; and the HTTP body now carries exitCode, which is the field the un-upgraded .NET handler already reads. Notes for review: - Server diagnostics go to a stderr handle bound at import, NOT ConsoleLogger. ConsoleLogger resolves sys.stdout at call time, and the runtime's 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. - A 4xx from the callback is REJECTED, not UNREACHABLE: the caller is up and has moved on, so retrying cannot help and must not trip the shutdown path. - The log tailer holds back an unterminated tail; a handler writes the record and only then flushes, so a poll can otherwise split one line into two entries that cannot be rejoined. - Logs are tailed from the file rather than captured via a handler: the runtime's log interceptor strips every handler but its own, and uipath-runtime ships on its own release train. Stopping a job that is already executing is deliberately NOT in this change -- it follows on top. Co-Authored-By: Claude Opus 5 --- .../uipath/src/uipath/_cli/_server_core.py | 148 ++++++- .../uipath/src/uipath/_cli/_server_jobs.py | 398 ++++++++++++++++++ packages/uipath/src/uipath/_cli/cli_server.py | 86 +++- .../uipath/src/uipath/_cli/cli_server_ipc.py | 31 ++ packages/uipath/tests/cli/test_server.py | 61 +++ .../uipath/tests/cli/test_server_async.py | 384 +++++++++++++++++ .../uipath/tests/cli/test_server_callbacks.py | 226 ++++++++++ packages/uipath/tests/cli/test_server_logs.py | 296 +++++++++++++ .../uipath/tests/cli/test_server_result.py | 187 ++++++++ 9 files changed, 1794 insertions(+), 23 deletions(-) create mode 100644 packages/uipath/src/uipath/_cli/_server_jobs.py create mode 100644 packages/uipath/tests/cli/test_server_async.py create mode 100644 packages/uipath/tests/cli/test_server_callbacks.py create mode 100644 packages/uipath/tests/cli/test_server_logs.py create mode 100644 packages/uipath/tests/cli/test_server_result.py diff --git a/packages/uipath/src/uipath/_cli/_server_core.py b/packages/uipath/src/uipath/_cli/_server_core.py index c426126ff..6ff6952f0 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/packages/uipath/src/uipath/_cli/_server_core.py @@ -1,6 +1,7 @@ """Transport-agnostic job core shared by the HTTP and uipath-ipc channels.""" import asyncio +import json import os import shlex from typing import Any @@ -45,6 +46,127 @@ 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" + + +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]) -> dict[str, Any]: + """Invoke one click command and classify how it ended.""" + 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 Exception as e: # report any job failure as a result, not a fault + return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True} + + async def _run_command_isolated( cmd: Any, args: list[str], @@ -79,25 +201,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) + # 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..810b6741e --- /dev/null +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -0,0 +1,398 @@ +"""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 ._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"), + } + + +LOG_POLL_SECONDS = 0.25 +LOG_BATCH_MAX_LINES = 200 +LOG_FLUSH_TIMEOUT_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]] = {} + + 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, + ) -> bool: + """Register and schedule a job. False if one is already in flight for this key.""" + if self.is_active(job_key): + return False + + 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) + + 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()) + + outcome = await _run_command_isolated( + cmd, + args, + env_vars, + working_dir, + ) + 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) -> bool: + """Cancel a job that has not started executing yet. + + Once the job is running its body is on a thread via ``asyncio.to_thread``, which + cannot be cancelled — so this only removes work that is still queued. Actually + interrupting a running job is a separate change. + """ + task = self._tasks.get(job_key) + if task is None or task.done(): + return True + + task.cancel() + return True + + +_registry = JobRegistry() + + +def get_registry() -> JobRegistry: + return _registry diff --git a/packages/uipath/src/uipath/_cli/cli_server.py b/packages/uipath/src/uipath/_cli/cli_server.py index d822bdd8c..2b7fd2796 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 ( @@ -112,9 +119,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) @@ -186,6 +198,40 @@ 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), + ) + 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"]: @@ -199,12 +245,25 @@ 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, "error": result["Error"]} + { + "success": False, + "job_key": job_key, + "exitCode": result["ExitCode"], + "error": result["Error"], + } ) @@ -373,7 +432,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 11c461ae6..38752f53d 100644 --- a/packages/uipath/src/uipath/_cli/cli_server_ipc.py +++ b/packages/uipath/src/uipath/_cli/cli_server_ipc.py @@ -2,6 +2,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() @@ -23,6 +24,10 @@ class RunJobRequest: 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 @dataclass @@ -36,6 +41,10 @@ class StopJobRequest: class RunJobResult: 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): @@ -68,6 +77,28 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: f"Running job {_run_id(request.JobKey, request.ResumeVersion)}: {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), + ) + if not accepted: + return RunJobResult( + ExitCode=1, + Error=f"Job {request.JobKey} is already in flight", + ContractVersion=CONTRACT_VERSION, + ) + return RunJobResult( + ExitCode=0, + Disposition="accepted", + ContractVersion=CONTRACT_VERSION, + ) + result = await _run_command_isolated( cmd, args, request.EnvironmentVariables, request.WorkingDirectory ) 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..9f33d18af --- /dev/null +++ b/packages/uipath/tests/cli/test_server_async.py @@ -0,0 +1,384 @@ +"""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 +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"}', + } + + +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 + + +# --------------------------------------------------------------------------- # +# 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] = [] + + def start(self, job_key, cmd, args, env_vars, working_dir, callback): + self.started.append(job_key) + 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().RunJob( + cli_server_ipc.RunJobRequest( + 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().RunJob( + cli_server_ipc.RunJobRequest( + 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().RunJob( + cli_server_ipc.RunJobRequest(JobKey="job-1", Command="run") + ) + + assert result.ExitCode == 7 + assert result.Disposition is None 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..daa699edd --- /dev/null +++ b/packages/uipath/tests/cli/test_server_callbacks.py @@ -0,0 +1,226 @@ +"""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", + } + + +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_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 From e5099410243618bc71cbee795e4aa00786eea213 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 18:55:04 +0300 Subject: [PATCH 2/6] fix(cli): type the job reporter as a protocol so the suite type-checks CI runs mypy over tests as well as src, and these tests were red: the recording fakes they pass to JobRegistry.start and JobLogTailer are not HandlerCallbacks. Neither collaborator wants the socket transport -- they want somewhere to send a result and somewhere to send logs -- so say that: JobReporter is a Protocol with those two methods, and HandlerCallback satisfies it structurally. The fakes then type-check as themselves rather than needing a cast at every call site. Worth noting the protocol immediately found a real gap: test_server_async's FakeCallback had no post_logs at all, so the tailer's calls into it were only ever working by accident of it never being exercised there. The rest is local: cast the _FakeRequest stand-ins to web.Request at the two handle_start call sites, and narrow web.Response.text (str | None) before json.loads and `in`. --- .../uipath/src/uipath/_cli/_server_jobs.py | 20 ++++++++++--- .../uipath/tests/cli/test_server_async.py | 30 ++++++++++++++----- .../uipath/tests/cli/test_server_result.py | 7 +++-- 3 files changed, 42 insertions(+), 15 deletions(-) diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index 810b6741e..65fce17d1 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -16,7 +16,7 @@ import os import re import sys -from typing import Any +from typing import Any, Protocol from aiohttp import ClientSession, ClientTimeout, UnixConnector @@ -90,6 +90,18 @@ class _PostOutcome(enum.Enum): UNREACHABLE = "unreachable" # transport failure or 5xx — worth retrying +class JobReporter(Protocol): + """Where a job's logs and outcome go. + + All the registry and the tailer need. Kept separate from HandlerCallback so a + reporter can be substituted without inheriting the socket transport. + """ + + async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: ... + + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: ... + + class HandlerCallback: """Posts job lifecycle events to the caller's Unix-socket HTTP API.""" @@ -200,7 +212,7 @@ class JobLogTailer: still produced exactly as before, we just also forward it. """ - def __init__(self, job_key: str, path: str, callback: "HandlerCallback") -> None: + def __init__(self, job_key: str, path: str, callback: "JobReporter") -> None: self.job_key = job_key self.path = path self.callback = callback @@ -307,7 +319,7 @@ def start( args: list[str], env_vars: dict[str, str], working_dir: str | None, - callback: HandlerCallback, + callback: JobReporter, ) -> bool: """Register and schedule a job. False if one is already in flight for this key.""" if self.is_active(job_key): @@ -330,7 +342,7 @@ async def _run( args: list[str], env_vars: dict[str, str], working_dir: str | None, - callback: HandlerCallback, + callback: JobReporter, ) -> None: tailer: JobLogTailer | None = None tail_task: "asyncio.Task[None] | None" = None diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py index 9f33d18af..f89ecd4e5 100644 --- a/packages/uipath/tests/cli/test_server_async.py +++ b/packages/uipath/tests/cli/test_server_async.py @@ -8,10 +8,11 @@ import asyncio import json import os -from typing import Any +from typing import Any, cast import click import pytest +from aiohttp import web from uipath._cli import _server_core, cli_server, cli_server_ipc from uipath._cli._server_core import ( @@ -44,6 +45,9 @@ async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: self.done.set() return True + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: + return True + def _write_config(tmp_path, runtime: dict[str, Any] | None) -> str: config = {"runtime": runtime} if runtime is not None else {} @@ -261,6 +265,16 @@ async def json(self) -> dict[str, Any]: return self._payload +def _fake_request(job_key: str, payload: dict[str, Any]) -> web.Request: + """handle_start only touches match_info and json(); the cast keeps mypy honest.""" + return cast(web.Request, _FakeRequest(job_key, payload)) + + +def _body(response: web.Response) -> dict[str, Any]: + assert response.text is not None + return cast("dict[str, Any]", json.loads(response.text)) + + class _RecordingRegistry: def __init__(self, accept: bool = True) -> None: self.accept = accept @@ -276,14 +290,14 @@ async def test_http_start_with_callback_returns_accepted(monkeypatch): monkeypatch.setattr(cli_server, "get_registry", lambda: registry) response = await cli_server.handle_start( - _FakeRequest( + _fake_request( "job-1", {"command": "run", "resultCallbackSocket": "/tmp/ack.sock"}, ) ) assert response.status == 202 - body = json.loads(response.text) + body = _body(response) assert body["disposition"] == "accepted" assert body["contractVersion"] == CONTRACT_VERSION assert registry.started == ["job-1"] @@ -295,13 +309,13 @@ async def test_http_start_duplicate_is_409(monkeypatch): ) response = await cli_server.handle_start( - _FakeRequest( + _fake_request( "job-1", {"command": "run", "resultCallbackSocket": "/tmp/ack.sock"} ) ) assert response.status == 409 - assert json.loads(response.text)["success"] is False + assert _body(response)["success"] is False async def test_http_start_without_callback_stays_synchronous(monkeypatch): @@ -319,11 +333,11 @@ async def _fake_isolated(cmd, args, env_vars, working_dir): lambda: pytest.fail("registry must not be used without a callback socket"), ) - response = await cli_server.handle_start(_FakeRequest("job-1", {"command": "run"})) + response = await cli_server.handle_start(_fake_request("job-1", {"command": "run"})) assert called.get("ran") is True assert response.status == 200 - body = json.loads(response.text) + body = _body(response) assert body["success"] is True assert "disposition" not in body @@ -362,7 +376,7 @@ async def test_ipc_start_duplicate_reports_failure(monkeypatch): assert result.ExitCode == 1 assert result.Disposition is None - assert "already in flight" in result.Error + assert "already in flight" in (result.Error or "") async def test_ipc_start_without_callback_stays_synchronous(monkeypatch): diff --git a/packages/uipath/tests/cli/test_server_result.py b/packages/uipath/tests/cli/test_server_result.py index 0e6f81bb8..f92f9bad8 100644 --- a/packages/uipath/tests/cli/test_server_result.py +++ b/packages/uipath/tests/cli/test_server_result.py @@ -7,7 +7,7 @@ rather than a stub, because it is the whole reason the exit code can be wrong. """ -from typing import Any +from typing import Any, cast import click import pytest @@ -132,7 +132,8 @@ 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"})) + request = cast(web.Request, _FakeRequest("job-1", {"command": "run"})) + return await cli_server.handle_start(request) async def test_success_body_carries_exit_code(monkeypatch): @@ -157,7 +158,7 @@ async def test_failure_body_is_200_but_says_so_and_carries_exit_code(monkeypatch ) assert response.status == 200 - body = response.text + body = response.text or "" assert '"success": false' in body assert '"exitCode": 1' in body assert "Exit code: 1" in body From e7d0b151b7f7a6ab0d1fc5c46486d8a1c723f09d Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 13:49:26 +0300 Subject: [PATCH 3/6] feat(cli): really stop a job that is already executing Builds on async dispatch. Until now StopJob could only refuse a running job: its body runs on a thread via asyncio.to_thread, which cannot be cancelled. That reasoning missed a fact. run/debug/eval each drive their own event loop with asyncio.run INSIDE the worker thread (cli_run.py:333, cli_debug.py:269, cli_eval.py:531), so a job IS an event loop -- and an event loop can be cancelled. The three call sites now go through `run_job_loop`, which publishes that loop and its root task to a JobControl carried on a ContextVar. asyncio.to_thread propagates contextvars, so no monkeypatching is needed, and outside the server (uipath run on a terminal) it is asyncio.run verbatim. Cancellation is cooperative, not a thread kill: the runtime's context managers still unwind, so UiPathRuntimeContext.__exit__ still writes output.json and the caller's file fallback still works. The stop ladder: 1. cancel the ROOT task only, then wait STOP_GRACE_SECONDS. Cancelling every task would land a second CancelledError inside the cleanup that writes output.json and abort it mid-write. 2. if cleanup itself is stuck, sweep the loop and wait a shorter window. 3. otherwise return False -- a job wedged in a non-cancellable C call (a socket read inside an LLM request) cannot be stopped, and saying so is better than claiming a stop that did not happen. A stopped job reports Stopped, not Faulted. The runtime records it as FAULTED/ERROR_CancelledError because it sees a CancelledError, which is the wrong story for a stop the caller asked for, so the result push carries an explicit `stopped` flag that wins over the document. Stop is also reachable now: POST /jobs/{job_key}/stop on the HTTP transport, which carries all current traffic and previously had no way to reach the registry at all. Notes for review: - _invoke_command discriminates the two CancelledErrors via Task.cancelling(): 0 means the job's own loop was cancelled (an outcome, swallow it); >0 means our awaiting task was cancelled (a shutdown, re-raise). - StopJob takes resume_version as a trailing optional parameter rather than a DTO: uipath-ipc ignores a surplus wire arg and defaults a missing one, so old and new peers interoperate both ways. A suspended job resumes under the same key, so a stop aimed at the previous run must not kill the resumed one. Co-Authored-By: Claude Opus 5 --- .../uipath/src/uipath/_cli/_job_control.py | 113 +++++++++ .../uipath/src/uipath/_cli/_server_core.py | 44 +++- .../uipath/src/uipath/_cli/_server_jobs.py | 108 ++++++++- packages/uipath/src/uipath/_cli/cli_debug.py | 4 +- packages/uipath/src/uipath/_cli/cli_eval.py | 3 +- packages/uipath/src/uipath/_cli/cli_run.py | 4 +- packages/uipath/src/uipath/_cli/cli_server.py | 32 +++ .../uipath/src/uipath/_cli/cli_server_ipc.py | 14 +- .../uipath/tests/cli/test_server_async.py | 185 ++++++++++++++- .../uipath/tests/cli/test_server_callbacks.py | 1 + .../tests/cli/test_server_cancellation.py | 222 ++++++++++++++++++ packages/uipath/tests/cli/test_server_ipc.py | 2 +- 12 files changed, 711 insertions(+), 21 deletions(-) create mode 100644 packages/uipath/src/uipath/_cli/_job_control.py create mode 100644 packages/uipath/tests/cli/test_server_cancellation.py diff --git a/packages/uipath/src/uipath/_cli/_job_control.py b/packages/uipath/src/uipath/_cli/_job_control.py new file mode 100644 index 000000000..f39df2e76 --- /dev/null +++ b/packages/uipath/src/uipath/_cli/_job_control.py @@ -0,0 +1,113 @@ +"""The seam that lets the server cancel the job running on its worker thread. + +Every server command drives its own event loop with ``asyncio.run`` inside the +``asyncio.to_thread`` worker. That loop is the only place a cancellation can actually +land, so the command publishes it here and the server reaches back through +``loop.call_soon_threadsafe``. + +Deliberately import-light so ``_cli/__init__.py``'s lazy-import discipline is untouched. +""" + +import asyncio +import contextvars +import threading +from typing import Any + +CURRENT_JOB_CONTROL: contextvars.ContextVar["JobControl | None"] = ( + contextvars.ContextVar("uipath_current_job_control", default=None) +) + + +class JobControl: + """Handle on one job's event loop, shared between the server loop and the worker.""" + + def __init__(self, job_key: str) -> None: + self.job_key = job_key + self.cancel_requested = False + self._loop: asyncio.AbstractEventLoop | None = None + self._task: "asyncio.Task[Any] | None" = None + # Bound from the job thread, read from the server loop. + self._sync = threading.Lock() + + # ---- job thread --------------------------------------------------------- + + def bind(self, loop: asyncio.AbstractEventLoop, task: "asyncio.Task[Any]") -> None: + with self._sync: + self._loop, self._task = loop, task + pending = self.cancel_requested + if pending: + # A stop that raced the loop's creation: apply it now rather than losing it. + self.cancel() + + def unbind(self) -> None: + with self._sync: + self._loop = self._task = None + + # ---- server loop -------------------------------------------------------- + + @property + def bound(self) -> bool: + with self._sync: + return self._task is not None + + def cancel(self) -> bool: + """Deliver CancelledError to the job's ROOT task only. + + Call this at most once. A second cancel lands inside the runtime's cleanup + ``finally`` blocks — the ones that write ``output.json`` and tear the log + interceptor down — and aborts them, which would destroy the very fallback the + caller relies on. Escalate with ``cancel_all`` instead. + """ + with self._sync: + self.cancel_requested = True + loop, task = self._loop, self._task + if loop is None or task is None: + # Not bound yet; bind() will apply it. + return False + try: + loop.call_soon_threadsafe(task.cancel) + except RuntimeError: + # Loop already closed — the job is finishing anyway. + return False + return True + + def cancel_all(self) -> bool: + """Escalation: cancel every task on the job's loop. + + This includes whatever the cleanup is awaiting, so it gives up on a clean + ``output.json``. Only worth doing once a polite cancel has already failed. + """ + with self._sync: + loop = self._loop + if loop is None: + return False + + def _sweep() -> None: + for task in asyncio.all_tasks(loop): + task.cancel() + + try: + loop.call_soon_threadsafe(_sweep) + except RuntimeError: + return False + return True + + +def run_job_loop(coro: Any) -> Any: + """``asyncio.run`` that publishes its loop and root task to the JobControl in scope. + + Outside the server — ``uipath run`` on a terminal — there is no handle in scope and + this is ``asyncio.run`` verbatim. + """ + control = CURRENT_JOB_CONTROL.get() + if control is None: + return asyncio.run(coro) + + with asyncio.Runner() as runner: + loop = runner.get_loop() + task = loop.create_task(coro, context=contextvars.copy_context()) + control.bind(loop, task) + try: + return loop.run_until_complete(task) + finally: + control.unbind() diff --git a/packages/uipath/src/uipath/_cli/_server_core.py b/packages/uipath/src/uipath/_cli/_server_core.py index 6ff6952f0..750f263cd 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/packages/uipath/src/uipath/_cli/_server_core.py @@ -4,8 +4,10 @@ import json import os import shlex +from collections.abc import Callable from typing import Any +from ._job_control import CURRENT_JOB_CONTROL, JobControl from .cli_debug import debug from .cli_eval import eval from .cli_run import run @@ -55,6 +57,10 @@ def parse_args(args: str | list[str] | None) -> list[str]: DEFAULT_RESULT_FILE = "output.json" DEFAULT_LOGS_FILE = "execution.log" +# Distinct from any click exit code, so the caller can tell "stopped on request" from +# "the job failed" and report Stopped rather than Faulted. +EXIT_CODE_STOPPED = 143 + def _resolve_runtime_file( config_path: str, base_dir: str, key: str, default_name: str @@ -134,8 +140,13 @@ def _read_result_document() -> tuple[str | None, str]: return None, "file" -async def _invoke_command(cmd: Any, args: list[str]) -> dict[str, Any]: +async def _invoke_command( + cmd: Any, args: list[str], control: "JobControl | None" = None +) -> dict[str, Any]: """Invoke one click command and classify how it ended.""" + # asyncio.to_thread propagates contextvars, so the command reads this inside the + # worker and publishes its event loop back through it. + token = CURRENT_JOB_CONTROL.set(control) if control is not None else None try: result_value = await asyncio.to_thread(cmd.main, args, standalone_mode=False) # Under standalone_mode=False click RETURNS ctx.exit(N)'s code instead of @@ -163,8 +174,27 @@ async def _invoke_command(cmd: Any, args: list[str]) -> dict[str, Any]: "Result": None, "Unexpected": False, } + except asyncio.CancelledError: + # Two very different things arrive here. cancelling() > 0 means OUR awaiting task + # was cancelled (server shutdown) — never swallow that, or the shutdown stalls. + # 0 means the CancelledError travelled out of the job's own loop, i.e. our StopJob + # landed: that is an OUTCOME, and swallowing it keeps this task alive so the lock + # unwinds and the result document still gets read. + current = asyncio.current_task() + if current is not None and current.cancelling() > 0: + raise + return { + "ExitCode": EXIT_CODE_STOPPED, + "Error": "Job stopped on request", + "Result": None, + "Unexpected": False, + "Stopped": True, + } except Exception as e: # report any job failure as a result, not a fault return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True} + finally: + if token is not None: + CURRENT_JOB_CONTROL.reset(token) async def _run_command_isolated( @@ -172,12 +202,20 @@ async def _run_command_isolated( args: list[str], env_vars: dict[str, str], working_dir: str | None, + on_started: Callable[[], None] | None = None, + control: "JobControl | None" = None, ) -> dict[str, Any]: - """Run one command with per-job env/cwd isolation (the shared job core).""" + """Run one command with per-job env/cwd isolation (the shared job core). + + ``on_started`` fires once the lock is held, i.e. the moment the job stops being + queued and becomes uncancellable. + """ if _state.lock is None or _state.baseline_env is None: raise RuntimeError("Server state not initialized") async with _state.lock: + if on_started is not None: + on_started() original_cwd = os.getcwd() try: # Start from server baseline + request env vars only, so nothing from @@ -201,7 +239,7 @@ async def _run_command_isolated( "ClientError": True, } - outcome = await _invoke_command(cmd, args) + outcome = await _invoke_command(cmd, args, control) # Must happen before the finally below restores env/cwd: the document's # location comes from this job's UIPATH_CONFIG_PATH and may be relative. document, conveyance = _read_result_document() diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index 65fce17d1..edf393512 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -20,6 +20,7 @@ from aiohttp import ClientSession, ClientTimeout, UnixConnector +from ._job_control import JobControl from ._server_core import _run_command_isolated, resolve_logs_file_path from ._utils._console import ConsoleLogger @@ -179,12 +180,21 @@ def build_result_payload(job_key: str, outcome: dict[str, Any]) -> dict[str, Any "stateConveyance": "file", "jobConveyance": outcome.get("DocumentConveyance", "file"), "job": outcome.get("Document"), + # Explicit so the caller reports Stopped rather than Faulted — the runtime's + # output.json records a cancelled job as FAULTED/ERROR_CancelledError, which is + # the wrong story for a stop the caller itself asked for. + "stopped": bool(outcome.get("Stopped")), } LOG_POLL_SECONDS = 0.25 LOG_BATCH_MAX_LINES = 200 LOG_FLUSH_TIMEOUT_SECONDS = 10 + +# How long a cancelled job gets to unwind before we admit the stop did not take. +STOP_GRACE_SECONDS = 30 +# Extra window after escalating to a full loop sweep. +STOP_ESCALATION_SECONDS = 10 # ``[2026-07-29 17:04:19,123][INFO] message`` — the format the runtime's file handler # emits and that the .NET FileLogsWatcher has always parsed. LOG_LINE_RE = re.compile( @@ -307,6 +317,13 @@ class JobRegistry: def __init__(self) -> None: self._tasks: dict[str, asyncio.Task[None]] = {} + # Jobs past the point of no return: executing on a thread, uncancellable. + self._running: set[str] = set() + # A suspended job resumes under the SAME key, so a stop meant for the previous + # run must not kill the resumed one. + self._resume_versions: dict[str, int | None] = {} + # Handle on each job's own event loop, so a stop can cancel it cooperatively. + self._controls: dict[str, JobControl] = {} def is_active(self, job_key: str) -> bool: task = self._tasks.get(job_key) @@ -320,11 +337,14 @@ def start( env_vars: dict[str, str], working_dir: str | None, callback: JobReporter, + resume_version: int | None = None, ) -> bool: """Register and schedule a job. False if one is already in flight for this key.""" if self.is_active(job_key): return False + self._resume_versions[job_key] = resume_version + task = asyncio.create_task( self._run(job_key, cmd, args, env_vars, working_dir, callback) ) @@ -334,6 +354,9 @@ def start( def _forget(self, job_key: str) -> None: self._tasks.pop(job_key, None) + self._running.discard(job_key) + self._resume_versions.pop(job_key, None) + self._controls.pop(job_key, None) async def _run( self, @@ -371,11 +394,15 @@ async def finish(outcome: dict[str, Any]) -> None: tailer = JobLogTailer(job_key, logs_path, callback) tail_task = asyncio.create_task(tailer.run()) + control = JobControl(job_key) + self._controls[job_key] = control outcome = await _run_command_isolated( cmd, args, env_vars, working_dir, + on_started=lambda: self._running.add(job_key), + control=control, ) except asyncio.CancelledError: # Cancelled before it took the lock — report it rather than going silent, @@ -388,18 +415,89 @@ async def finish(outcome: dict[str, Any]) -> None: await finish(outcome) - async def stop(self, job_key: str) -> bool: + async def stop(self, job_key: str, resume_version: int | None = None) -> bool: """Cancel a job that has not started executing yet. - Once the job is running its body is on a thread via ``asyncio.to_thread``, which - cannot be cancelled — so this only removes work that is still queued. Actually - interrupting a running job is a separate change. + Only queued work can be stopped. Once the job holds the lock its body is on a + thread via ``asyncio.to_thread``, which cannot be cancelled cooperatively — + cancelling the awaiting task would free the lock and report "cancelled before + execution" while the work carried on mutating process globals underneath the + next job. Refusing is the honest answer; real cancellation has to be designed + inside the runtime itself. """ task = self._tasks.get(job_key) if task is None or task.done(): return True - task.cancel() + # A stop is raised against a specific run. A suspended job resumes under the same + # key, so a stop for run N that arrives after N+1 started must not kill N+1. + if resume_version is not None: + registered = self._resume_versions.get(job_key) + if registered is not None and registered != resume_version: + console.warning( + f"StopJob for {job_key} ignored: it targets resume version " + f"{resume_version} but the live run is {registered}." + ) + return False + + if job_key not in self._running: + # Still queued behind the lock: cancelling the task is enough, and the job + # never touched process state. + task.cancel() + return True + + # Executing. Cancel the job's OWN event loop rather than the awaiting task: + # that unwinds the runtime cooperatively, so its context managers run and + # output.json still gets written for the caller to fall back on. + control = self._controls.get(job_key) + if control is None: + _server_log(f"StopJob for {job_key}: no control handle") + return False + + # Rung 1 — cancel the ROOT task only. The runtime's cleanup then unwinds + # normally, which is what writes output.json for the caller to fall back on. + # Cancelling everything here would abort that cleanup mid-write. + if not control.cancel(): + _server_log( + f"StopJob for {job_key}: the job has not started an event loop yet; " + "the request is recorded and applied as soon as it does." + ) + + task = self._tasks.get(job_key) + if task is None: + return True + + try: + await asyncio.wait_for(asyncio.shield(task), STOP_GRACE_SECONDS) + return True + except asyncio.TimeoutError: + pass + except BaseException: + # It ended; how it ended is the result push's business. + return True + + # Rung 2 — cleanup itself is stuck. Sweep the loop, giving up on a clean + # output.json in exchange for releasing the lock. + _server_log( + f"StopJob for {job_key}: cleanup did not finish in {STOP_GRACE_SECONDS}s; " + "cancelling every task on the job's loop." + ) + control.cancel_all() + + try: + await asyncio.wait_for(asyncio.shield(task), STOP_ESCALATION_SECONDS) + return True + except asyncio.TimeoutError: + # Rung 3 would be taking the process down, which costs every queued job. + # Report the truth instead and let the caller decide. + _server_log( + f"StopJob for {job_key}: still running — it is blocked in a call that " + "cannot be interrupted (a socket read inside an LLM request, typically)." + ) + return False + except BaseException: + return True + return True diff --git a/packages/uipath/src/uipath/_cli/cli_debug.py b/packages/uipath/src/uipath/_cli/cli_debug.py index 7d7eceba2..12351900d 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 @@ -29,6 +28,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 @@ -295,7 +295,7 @@ async def execute_debug_runtime(): finally: trace_manager.shutdown() - asyncio.run(execute_debug_runtime()) + run_job_loop(execute_debug_runtime()) except Exception as e: console.error( f"Error occurred: {e or 'Execution failed'}", include_traceback=True diff --git a/packages/uipath/src/uipath/_cli/cli_eval.py b/packages/uipath/src/uipath/_cli/cli_eval.py index 66bdfad10..4483c3b2f 100644 --- a/packages/uipath/src/uipath/_cli/cli_eval.py +++ b/packages/uipath/src/uipath/_cli/cli_eval.py @@ -39,6 +39,7 @@ LlmOpsHttpExporter, ) +from ._job_control import run_job_loop from ._utils._console import ConsoleLogger logger = logging.getLogger(__name__) @@ -528,7 +529,7 @@ async def execute_eval(): finally: await runtime_factory.dispose() - asyncio.run(execute_eval()) + run_job_loop(execute_eval()) except _EvalDiscoveryError as e: click.echo("\n".join(e.get_usage_help())) diff --git a/packages/uipath/src/uipath/_cli/cli_run.py b/packages/uipath/src/uipath/_cli/cli_run.py index 9d12a86c3..2574c0fac 100644 --- a/packages/uipath/src/uipath/_cli/cli_run.py +++ b/packages/uipath/src/uipath/_cli/cli_run.py @@ -1,4 +1,3 @@ -import asyncio from typing import Any import click @@ -36,6 +35,7 @@ from ._errors import EntrypointDiscoveryException from ._governance_bootstrap import GovernanceBootstrap, resolve_governance +from ._job_control import run_job_loop from ._telemetry import track_command from ._utils._console import ConsoleLogger from .middlewares import Middlewares @@ -330,7 +330,7 @@ async def execute() -> None: finally: trace_manager.shutdown() - asyncio.run(execute()) + run_job_loop(execute()) except _RunDiscoveryError as e: click.echo("\n".join(e.get_usage_help())) diff --git a/packages/uipath/src/uipath/_cli/cli_server.py b/packages/uipath/src/uipath/_cli/cli_server.py index 2b7fd2796..f603280cb 100644 --- a/packages/uipath/src/uipath/_cli/cli_server.py +++ b/packages/uipath/src/uipath/_cli/cli_server.py @@ -211,6 +211,7 @@ async def handle_start(request: web.Request) -> web.Response: env_vars, working_dir, HandlerCallback(result_callback_socket), + resume_version=get_field(message, "resumeVersion", "ResumeVersion"), ) if not accepted: return web.json_response( @@ -267,6 +268,36 @@ async def handle_start(request: web.Request) -> web.Response: ) +async def handle_stop(request: web.Request) -> web.Response: + """Stop a job, reporting whether it actually stopped. + + Always 200 with the outcome in the body — 4xx/5xx stay reserved for request-shaped + failures, matching how /start already reports a failed job. + """ + job_key = request.match_info.get("job_key") + if not job_key: + return web.json_response( + {"success": False, "error": "Missing job_key"}, status=400 + ) + + try: + message: dict[str, Any] = await request.json() + except json.JSONDecodeError: + message = {} + + resume_version = get_field(message, "resumeVersion", "ResumeVersion") + stopped = await get_registry().stop(job_key, resume_version) + + return web.json_response( + { + "success": True, + "job_key": job_key, + "contractVersion": CONTRACT_VERSION, + "stopped": stopped, + } + ) + + ALLOWED_HOSTS = {"127.0.0.1", "localhost", "[::1]"} @@ -301,6 +332,7 @@ def create_app() -> web.Application: app = web.Application(middlewares=[host_validation_middleware]) app.router.add_get("/health", handle_health) app.router.add_post("/jobs/{job_key}/start", handle_start) + app.router.add_post("/jobs/{job_key}/stop", handle_stop) return app diff --git a/packages/uipath/src/uipath/_cli/cli_server_ipc.py b/packages/uipath/src/uipath/_cli/cli_server_ipc.py index 38752f53d..161c917e7 100644 --- a/packages/uipath/src/uipath/_cli/cli_server_ipc.py +++ b/packages/uipath/src/uipath/_cli/cli_server_ipc.py @@ -56,7 +56,12 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: @abstractmethod async def StopJob(self, request: StopJobRequest) -> bool: - """Cancel a running job by key (bool return avoids fire-and-forget).""" + """Cancel a running job (bool return avoids fire-and-forget). + + A nested DTO, matching the .NET peer: a suspended job resumes under the same + JobKey, so ``ResumeVersion`` is what keeps a stop aimed at the previous run from + killing the resumed one. + """ class PythonRuntimeService(IPythonRuntimeServer): @@ -86,6 +91,7 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: request.EnvironmentVariables, request.WorkingDirectory, HandlerCallback(callback_socket), + resume_version=request.ResumeVersion, ) if not accepted: return RunJobResult( @@ -106,11 +112,7 @@ async def RunJob(self, request: RunJobRequest) -> RunJobResult: return RunJobResult(ExitCode=result["ExitCode"], Error=result["Error"]) async def StopJob(self, request: StopJobRequest) -> bool: - console.info( - f"StopJob requested for {_run_id(request.JobKey, request.ResumeVersion)} " - f"(force={request.ForceStop}) (no-op)" - ) - return True + return await get_registry().stop(request.JobKey, request.ResumeVersion) async def start_ipc_server(pipe_name: str) -> None: diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py index f89ecd4e5..feb605442 100644 --- a/packages/uipath/tests/cli/test_server_async.py +++ b/packages/uipath/tests/cli/test_server_async.py @@ -8,6 +8,7 @@ import asyncio import json import os +import time from typing import Any, cast import click @@ -26,6 +27,7 @@ JobRegistry, build_result_payload, ) +from uipath._cli.cli_server_ipc import StopJobRequest @pytest.fixture(autouse=True) @@ -169,6 +171,7 @@ def test_build_result_payload_shape(): "stateConveyance": "file", "jobConveyance": "inline", "job": '{"status":"successful"}', + "stopped": False, } @@ -251,6 +254,60 @@ async def test_stop_is_true_for_an_unknown_job(): assert await registry.stop("never-seen") is True +async def test_stop_waits_for_a_running_job_to_finish(tmp_path, monkeypatch): + """Stop no longer refuses a running job — it cancels and waits. A job that ends + inside the grace window is a successful stop, whatever ended it.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "STOP_GRACE_SECONDS", 5) + monkeypatch.setattr(jobs, "STOP_ESCALATION_SECONDS", 2) + + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + started = asyncio.Event() + + @click.command() + def _brief_command() -> None: + started.set() + time.sleep(0.3) + + registry.start("job-1", _brief_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(started.wait(), timeout=10) + + assert await registry.stop("job-1") is True + await asyncio.wait_for(callback.done.wait(), timeout=10) + + +async def test_stop_cancels_a_job_that_is_still_queued(tmp_path): + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start("job-1", _holds_the_lock, [], {}, str(tmp_path), first) + await asyncio.wait_for(started.wait(), timeout=10) + + # job-2 is queued behind the lock, so it has not started executing. + registry.start("job-2", _ok_command, [], {}, str(tmp_path), second) + cancelled = await registry.stop("job-2") + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + assert cancelled is True + + # --------------------------------------------------------------------------- # # HTTP dispatch # # --------------------------------------------------------------------------- # @@ -279,9 +336,13 @@ class _RecordingRegistry: def __init__(self, accept: bool = True) -> None: self.accept = accept self.started: list[str] = [] + self.resume_versions: list[int | None] = [] - def start(self, job_key, cmd, args, env_vars, working_dir, callback): + def start( + self, job_key, cmd, args, env_vars, working_dir, callback, resume_version=None + ): self.started.append(job_key) + self.resume_versions.append(resume_version) return self.accept @@ -396,3 +457,125 @@ async def _fake_isolated(cmd, args, env_vars, working_dir): assert result.ExitCode == 7 assert result.Disposition is None + + +async def test_stop_ignores_a_command_for_a_previous_resume_version(tmp_path): + """A suspended job resumes under the SAME key. A stop raised against run N that + arrives after N+1 started must not kill N+1.""" + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start( + "job-1", _holds_the_lock, [], {}, str(tmp_path), first, resume_version=2 + ) + await asyncio.wait_for(started.wait(), timeout=10) + + # A stop for the run that was suspended before this one. + assert await registry.stop("job-1", resume_version=1) is False + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + assert first.results[0]["exitCode"] == 0 + + +async def test_stop_applies_to_the_matching_resume_version(tmp_path): + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + started = asyncio.Event() + release = asyncio.Event() + + @click.command() + def _holds_the_lock() -> None: + started.set() + waited = 0 + while not release.is_set() and waited < 200: + time.sleep(0.02) + waited += 1 + + registry.start("job-1", _holds_the_lock, [], {}, str(tmp_path), first) + await asyncio.wait_for(started.wait(), timeout=10) + + registry.start( + "job-2", _ok_command, [], {}, str(tmp_path), second, resume_version=3 + ) + assert await registry.stop("job-2", resume_version=3) is True + + release.set() + await asyncio.wait_for(first.done.wait(), timeout=10) + + +# --------------------------------------------------------------------------- # +# HTTP stop — the transport production actually uses # +# --------------------------------------------------------------------------- # + + +class _RecordingStopRegistry: + def __init__(self, stopped: bool = True) -> None: + self.stopped = stopped + self.calls: list[tuple[str, int | None]] = [] + + async def stop(self, job_key, resume_version=None): + self.calls.append((job_key, resume_version)) + return self.stopped + + +async def test_http_stop_reaches_the_registry(monkeypatch): + """Stop was previously unreachable over HTTP, which is the transport carrying all + production traffic — an operator stop simply never reached the runtime.""" + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server, "get_registry", lambda: registry) + + response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + + assert response.status == 200 + body = json.loads(response.text) + assert body["stopped"] is True + assert registry.calls == [("job-1", None)] + + +async def test_http_stop_forwards_the_resume_version(monkeypatch): + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server, "get_registry", lambda: registry) + + await cli_server.handle_stop(_FakeRequest("job-1", {"resumeVersion": 2})) + + assert registry.calls == [("job-1", 2)] + + +async def test_http_stop_reports_a_refused_stop(monkeypatch): + """A job wedged in a non-cancellable call must be reported honestly, not as stopped.""" + monkeypatch.setattr( + cli_server, "get_registry", lambda: _RecordingStopRegistry(stopped=False) + ) + + response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + + assert response.status == 200 + assert json.loads(response.text)["stopped"] is False + + +async def test_ipc_stop_forwards_the_dto_fields(monkeypatch): + """The .NET peer sends a StopJobRequest; both identity fields must reach the registry + or the resume-version guard is inert.""" + registry = _RecordingStopRegistry() + monkeypatch.setattr(cli_server_ipc, "get_registry", lambda: registry) + + await cli_server_ipc.PythonRuntimeService().StopJob( + StopJobRequest(JobKey="job-1", ResumeVersion=2, ForceStop=True) + ) + + assert registry.calls == [("job-1", 2)] diff --git a/packages/uipath/tests/cli/test_server_callbacks.py b/packages/uipath/tests/cli/test_server_callbacks.py index daa699edd..f2015dd73 100644 --- a/packages/uipath/tests/cli/test_server_callbacks.py +++ b/packages/uipath/tests/cli/test_server_callbacks.py @@ -82,6 +82,7 @@ def test_result_payload_carries_the_keys_the_handler_deserializes(): "stateConveyance", "jobConveyance", "job", + "stopped", } diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py new file mode 100644 index 000000000..3e88ee2d7 --- /dev/null +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -0,0 +1,222 @@ +"""Really stopping a job that is already executing. + +``run``/``debug``/``eval`` each end in ``asyncio.run(...)`` on the worker thread, so a +job *is* an event loop. Capturing that loop turns cancellation from "impossible, the +thread is opaque" into an ordinary ``task.cancel()`` that unwinds the runtime +cooperatively — which is what these pin, using commands shaped like the real ones. +""" + +import asyncio +import time + +import click +import pytest + +from uipath._cli import _server_core +from uipath._cli._job_control import CURRENT_JOB_CONTROL, JobControl, run_job_loop +from uipath._cli._server_core import EXIT_CODE_STOPPED, _ServerState +from uipath._cli._server_jobs import JobRegistry + + +@pytest.fixture(autouse=True) +def _fresh_state(monkeypatch): + monkeypatch.setattr(_server_core, "_state", _ServerState()) + + +class FakeCallback: + def __init__(self) -> None: + self.results: list[dict] = [] + self.done = asyncio.Event() + + async def post_result(self, job_key, payload): + self.results.append(payload) + self.done.set() + return True + + async def post_logs(self, job_key, lines): + return True + + +# Shaped like the real commands: a click command whose body is asyncio.run(...). +_started = None +_cleanup_ran = None + + +@click.command() +def _long_async_command() -> None: + async def body() -> None: + try: + _started.set() + await asyncio.sleep(30) + finally: + # Stands in for UiPathRuntimeContext.__exit__, which writes output.json. + _cleanup_ran.set() + + run_job_loop(body()) + + +@click.command() +def _quick_async_command() -> None: + async def body() -> None: + await asyncio.sleep(0) + + run_job_loop(body()) + + +@click.command() +def _blocking_command() -> None: + # No event loop at all: models a job wedged in a non-cancellable C call. + _started.set() + time.sleep(30) + + +@pytest.fixture(autouse=True) +def _events(): + global _started, _cleanup_ran + _started = __import__("threading").Event() + _cleanup_ran = __import__("threading").Event() + yield + + +async def _wait_for(event, timeout=10.0): + deadline = asyncio.get_running_loop().time() + timeout + while not event.is_set(): + if asyncio.get_running_loop().time() > deadline: + raise AssertionError("timed out waiting for the job to start") + await asyncio.sleep(0.02) + + +# --------------------------------------------------------------------------- # +# the capture mechanism # +# --------------------------------------------------------------------------- # + + +def test_control_starts_unbound(): + control = JobControl("job-1") + + assert control.bound is False + # Nothing to cancel yet, but the request must be remembered for bind(). + assert control.cancel() is False + assert control.cancel_requested is True + + +def test_run_job_loop_is_plain_asyncio_run_outside_the_server(): + """`uipath run` on a terminal has no control in scope and must be unaffected.""" + assert CURRENT_JOB_CONTROL.get() is None + + async def body(): + return 42 + + assert run_job_loop(body()) == 42 + + +async def test_control_binds_to_the_loop_the_job_runs_on(tmp_path): + _server_core._state.init() + control = JobControl("job-1") + + await _server_core._run_command_isolated( + _quick_async_command, [], {}, str(tmp_path), control=control + ) + + # Unbound again once the job finished. + assert control.bound is False + + +async def test_a_stop_that_races_the_loop_is_still_applied(tmp_path): + """A cancel arriving before the job builds its loop must not be dropped.""" + _server_core._state.init() + control = JobControl("job-1") + control.cancel() # before anything is bound + + result = await _server_core._run_command_isolated( + _long_async_command, [], {}, str(tmp_path), control=control + ) + + assert result["Stopped"] is True + + +# --------------------------------------------------------------------------- # +# stopping a running job for real # +# --------------------------------------------------------------------------- # + + +async def test_stop_cancels_a_job_that_is_already_executing(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + stopped = await registry.stop("job-1") + + assert stopped is True, "a running job must actually stop, not just be refused" + await asyncio.wait_for(callback.done.wait(), timeout=10) + assert callback.results[0]["stopped"] is True + assert callback.results[0]["exitCode"] == EXIT_CODE_STOPPED + + +async def test_a_stopped_job_still_unwinds_its_cleanup(tmp_path): + """Cooperative cancellation, not a thread kill: the runtime's context managers must + still run, because that is what writes output.json for the caller to fall back on.""" + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + await registry.stop("job-1") + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert _cleanup_ran.is_set(), "the job's finally block must have run" + + +async def test_stop_releases_the_lock_for_the_next_job(tmp_path): + """A stop that leaves the lock held would wedge the whole server.""" + _server_core._state.init() + registry = JobRegistry() + first = FakeCallback() + second = FakeCallback() + + registry.start("job-1", _long_async_command, [], {}, str(tmp_path), first) + await _wait_for(_started) + await registry.stop("job-1") + await asyncio.wait_for(first.done.wait(), timeout=10) + + registry.start("job-2", _quick_async_command, [], {}, str(tmp_path), second) + await asyncio.wait_for(second.done.wait(), timeout=10) + + assert second.results[0]["exitCode"] == 0 + + +async def test_a_normal_job_is_not_reported_stopped(tmp_path): + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _quick_async_command, [], {}, str(tmp_path), callback) + await asyncio.wait_for(callback.done.wait(), timeout=10) + + assert callback.results[0]["stopped"] is False + assert callback.results[0]["exitCode"] == 0 + + +async def test_stop_reports_failure_when_the_job_cannot_be_interrupted( + tmp_path, monkeypatch +): + """A job with no event loop — wedged in a C call — cannot be cancelled. Say so + rather than claiming a stop that did not happen.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "STOP_GRACE_SECONDS", 1) + monkeypatch.setattr(jobs, "STOP_ESCALATION_SECONDS", 1) + + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _blocking_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + assert await registry.stop("job-1") is False diff --git a/packages/uipath/tests/cli/test_server_ipc.py b/packages/uipath/tests/cli/test_server_ipc.py index 43fe313c9..6110738df 100644 --- a/packages/uipath/tests/cli/test_server_ipc.py +++ b/packages/uipath/tests/cli/test_server_ipc.py @@ -201,7 +201,7 @@ def test_stop_job_accepts_resume_version_and_force_stop(self, pipe): assert result is True def test_stop_job_returns_true(self, pipe): - """StopJob is a no-op stub today, but must ack (bool) so the call is awaitable.""" + """StopJob must ack (bool) so the call is awaitable; an unknown key is a no-op.""" result = asyncio.run( _with_proxy( pipe, lambda p: p.StopJob({"JobKey": "job-1", "ForceStop": True}) From f08578b39f37c46f057f218f264fc0a03a42bb19 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 18:51:04 +0300 Subject: [PATCH 4/6] fix(cli): deliver one cancellation per stop, and only call it Stopped when asked Two defects in the cancellation path, both found in review. JobControl.cancel() documented that it must be called at most once, but JobRegistry.stop() called it unguarded on every request. A stop followed by a force-stop escalation is ordinary, and the second delivery lands inside the runtime's cleanup finally blocks -- the ones that write output.json -- and aborts them, destroying the fallback the caller relies on. Absorbing the repeat belongs in JobControl, not in every caller. _invoke_command treated any CancelledError with cancelling() == 0 as a user-requested stop. cancelling() reports on the awaiting server task, so a CancelledError the job's own code let escape was reported as exit 143 "stopped on request" with no StopJob in sight -- a fault filed as a clean stop. Gate the classification on the control's own cancel_requested. Also corrects JobRegistry.stop()'s docstring, which still described the pre-cancellation behaviour of refusing to stop executing work. --- .../uipath/src/uipath/_cli/_job_control.py | 23 ++++-- .../uipath/src/uipath/_cli/_server_core.py | 28 ++++--- .../uipath/src/uipath/_cli/_server_jobs.py | 14 ++-- .../tests/cli/test_server_cancellation.py | 74 +++++++++++++++++++ 4 files changed, 114 insertions(+), 25 deletions(-) diff --git a/packages/uipath/src/uipath/_cli/_job_control.py b/packages/uipath/src/uipath/_cli/_job_control.py index f39df2e76..1a32b4d64 100644 --- a/packages/uipath/src/uipath/_cli/_job_control.py +++ b/packages/uipath/src/uipath/_cli/_job_control.py @@ -24,6 +24,7 @@ class JobControl: def __init__(self, job_key: str) -> None: self.job_key = job_key self.cancel_requested = False + self._delivered = False self._loop: asyncio.AbstractEventLoop | None = None self._task: "asyncio.Task[Any] | None" = None # Bound from the job thread, read from the server loop. @@ -51,23 +52,29 @@ def bound(self) -> bool: return self._task is not None def cancel(self) -> bool: - """Deliver CancelledError to the job's ROOT task only. + """Deliver CancelledError to the job's ROOT task, at most once. - Call this at most once. A second cancel lands inside the runtime's cleanup - ``finally`` blocks — the ones that write ``output.json`` and tear the log - interceptor down — and aborts them, which would destroy the very fallback the - caller relies on. Escalate with ``cancel_all`` instead. + Delivering twice lands the second one inside the runtime's cleanup ``finally`` + blocks — the ones that write ``output.json`` and tear the log interceptor down — + and aborts them, destroying the very fallback the caller relies on. Repeat stops + are ordinary (a stop followed by a force-stop escalation), so absorbing them is + this object's job rather than every caller's. Escalate with ``cancel_all``. """ with self._sync: self.cancel_requested = True + if self._delivered: + return True loop, task = self._loop, self._task - if loop is None or task is None: - # Not bound yet; bind() will apply it. - return False + if loop is None or task is None: + # Not bound yet; bind() will apply it. + return False + self._delivered = True try: loop.call_soon_threadsafe(task.cancel) except RuntimeError: # Loop already closed — the job is finishing anyway. + with self._sync: + self._delivered = False return False return True diff --git a/packages/uipath/src/uipath/_cli/_server_core.py b/packages/uipath/src/uipath/_cli/_server_core.py index 750f263cd..0bafc2b1a 100644 --- a/packages/uipath/src/uipath/_cli/_server_core.py +++ b/packages/uipath/src/uipath/_cli/_server_core.py @@ -175,20 +175,30 @@ async def _invoke_command( "Unexpected": False, } except asyncio.CancelledError: - # Two very different things arrive here. cancelling() > 0 means OUR awaiting task - # was cancelled (server shutdown) — never swallow that, or the shutdown stalls. - # 0 means the CancelledError travelled out of the job's own loop, i.e. our StopJob - # landed: that is an OUTCOME, and swallowing it keeps this task alive so the lock - # unwinds and the result document still gets read. + # Three very different things arrive here. current = asyncio.current_task() if current is not None and current.cancelling() > 0: + # OUR awaiting task was cancelled (server shutdown). Never swallow that, or + # the shutdown stalls. raise + if control is not None and control.cancel_requested: + # A stop we asked for travelled out of the job's own loop. That is an + # OUTCOME: swallowing it keeps this task alive so the lock unwinds and the + # result document still gets read. + return { + "ExitCode": EXIT_CODE_STOPPED, + "Error": "Job stopped on request", + "Result": None, + "Unexpected": False, + "Stopped": True, + } + # Nobody asked for this: the job's own code let a CancelledError escape. It is a + # failure, and reporting it as Stopped would file a fault as a clean stop. return { - "ExitCode": EXIT_CODE_STOPPED, - "Error": "Job stopped on request", + "ExitCode": 1, + "Error": "Job cancelled", "Result": None, - "Unexpected": False, - "Stopped": True, + "Unexpected": True, } except Exception as e: # report any job failure as a result, not a fault return {"ExitCode": 1, "Error": str(e), "Result": None, "Unexpected": True} diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index edf393512..a984c3886 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -416,14 +416,12 @@ async def finish(outcome: dict[str, Any]) -> None: await finish(outcome) async def stop(self, job_key: str, resume_version: int | None = None) -> bool: - """Cancel a job that has not started executing yet. - - Only queued work can be stopped. Once the job holds the lock its body is on a - thread via ``asyncio.to_thread``, which cannot be cancelled cooperatively — - cancelling the awaiting task would free the lock and report "cancelled before - execution" while the work carried on mutating process globals underneath the - next job. Refusing is the honest answer; real cancellation has to be designed - inside the runtime itself. + """Stop a job, whether it is still queued or already executing. + + Queued work is cancelled by dropping the waiting task. Executing work is + cancelled through the job's own event loop so the runtime unwinds cooperatively. + Safe to call repeatedly for the same run: the second request does not deliver a + second cancellation. """ task = self._tasks.get(job_key) if task is None or task.done(): diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py index 3e88ee2d7..8fd0efbac 100644 --- a/packages/uipath/tests/cli/test_server_cancellation.py +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -63,6 +63,31 @@ async def body() -> None: run_job_loop(body()) +@click.command() +def _slow_cleanup_command() -> None: + async def body() -> None: + try: + _started.set() + await asyncio.sleep(30) + finally: + # The real cleanup awaits (flushing traces, closing clients), which is the + # window a second cancellation would land in. + await asyncio.sleep(0.5) + _cleanup_ran.set() + + run_job_loop(body()) + + +@click.command() +def _self_cancelling_command() -> None: + async def body() -> None: + _started.set() + # Nobody asked this job to stop; its own code let a CancelledError escape. + raise asyncio.CancelledError() + + run_job_loop(body()) + + @click.command() def _blocking_command() -> None: # No event loop at all: models a job wedged in a non-cancellable C call. @@ -156,6 +181,55 @@ async def test_stop_cancels_a_job_that_is_already_executing(tmp_path): assert callback.results[0]["exitCode"] == EXIT_CODE_STOPPED +async def test_cancel_delivers_a_single_cancellation_however_often_it_is_called(): + """A second delivery lands in the cleanup and aborts it, so cancel() absorbs it.""" + control = JobControl("job-1") + loop = asyncio.get_running_loop() + task = loop.create_task(asyncio.sleep(30)) + control.bind(loop, task) + + assert control.cancel() is True + assert control.cancel() is True + assert control.cancel() is True + + await asyncio.gather(task, return_exceptions=True) + assert task.cancelling() == 1 + + +async def test_a_repeated_stop_does_not_abort_the_cleanup(tmp_path): + """A stop followed by a force-stop escalation is ordinary; both must be honoured + without the second one interrupting the cleanup that writes output.json.""" + _server_core._state.init() + registry = JobRegistry() + callback = FakeCallback() + + registry.start("job-1", _slow_cleanup_command, [], {}, str(tmp_path), callback) + await _wait_for(_started) + + first, second = await asyncio.gather(registry.stop("job-1"), registry.stop("job-1")) + + assert (first, second) == (True, True) + await asyncio.wait_for(callback.done.wait(), timeout=10) + assert _cleanup_ran.is_set(), "the second stop interrupted the job's cleanup" + assert callback.results[0]["stopped"] is True + + +async def test_a_self_inflicted_cancellation_is_a_fault_not_a_stop(tmp_path): + """`cancelling() == 0` alone does not mean the caller asked for a stop — without a + stop request a stray CancelledError is a job failure, and filing it as Stopped + would report a fault as a clean stop.""" + _server_core._state.init() + control = JobControl("job-1") + + result = await _server_core._run_command_isolated( + _self_cancelling_command, [], {}, str(tmp_path), control=control + ) + + assert result.get("Stopped") is not True + assert result["ExitCode"] != EXIT_CODE_STOPPED + assert result["Unexpected"] is True + + async def test_a_stopped_job_still_unwinds_its_cleanup(tmp_path): """Cooperative cancellation, not a thread kill: the runtime's context managers must still run, because that is what writes output.json for the caller to fall back on.""" From e14f4b29d26780f990ec952bdaf69a80b8ff9f1a Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 18:58:45 +0300 Subject: [PATCH 5/6] fix(cli): type-check the cancellation tests Same CI gate as the parent commit, applied to the tests this branch adds: give the module-level events a real Event type instead of letting them infer None, annotate the recording callback against the JobReporter protocol, and cast the _FakeRequest stand-ins at the handle_stop call sites. --- .../uipath/tests/cli/test_server_async.py | 10 +++++----- .../tests/cli/test_server_cancellation.py | 19 +++++++++++-------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/packages/uipath/tests/cli/test_server_async.py b/packages/uipath/tests/cli/test_server_async.py index feb605442..ed39f9e69 100644 --- a/packages/uipath/tests/cli/test_server_async.py +++ b/packages/uipath/tests/cli/test_server_async.py @@ -539,10 +539,10 @@ async def test_http_stop_reaches_the_registry(monkeypatch): registry = _RecordingStopRegistry() monkeypatch.setattr(cli_server, "get_registry", lambda: registry) - response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + response = await cli_server.handle_stop(_fake_request("job-1", {})) assert response.status == 200 - body = json.loads(response.text) + body = _body(response) assert body["stopped"] is True assert registry.calls == [("job-1", None)] @@ -551,7 +551,7 @@ async def test_http_stop_forwards_the_resume_version(monkeypatch): registry = _RecordingStopRegistry() monkeypatch.setattr(cli_server, "get_registry", lambda: registry) - await cli_server.handle_stop(_FakeRequest("job-1", {"resumeVersion": 2})) + await cli_server.handle_stop(_fake_request("job-1", {"resumeVersion": 2})) assert registry.calls == [("job-1", 2)] @@ -562,10 +562,10 @@ async def test_http_stop_reports_a_refused_stop(monkeypatch): cli_server, "get_registry", lambda: _RecordingStopRegistry(stopped=False) ) - response = await cli_server.handle_stop(_FakeRequest("job-1", {})) + response = await cli_server.handle_stop(_fake_request("job-1", {})) assert response.status == 200 - assert json.loads(response.text)["stopped"] is False + assert _body(response)["stopped"] is False async def test_ipc_stop_forwards_the_dto_fields(monkeypatch): diff --git a/packages/uipath/tests/cli/test_server_cancellation.py b/packages/uipath/tests/cli/test_server_cancellation.py index 8fd0efbac..79c6787a1 100644 --- a/packages/uipath/tests/cli/test_server_cancellation.py +++ b/packages/uipath/tests/cli/test_server_cancellation.py @@ -7,7 +7,9 @@ """ import asyncio +import threading import time +from typing import Any import click import pytest @@ -25,21 +27,22 @@ def _fresh_state(monkeypatch): class FakeCallback: def __init__(self) -> None: - self.results: list[dict] = [] + self.results: list[dict[str, Any]] = [] self.done = asyncio.Event() - async def post_result(self, job_key, payload): + async def post_result(self, job_key: str, payload: dict[str, Any]) -> bool: self.results.append(payload) self.done.set() return True - async def post_logs(self, job_key, lines): + async def post_logs(self, job_key: str, lines: list[dict[str, Any]]) -> bool: return True # Shaped like the real commands: a click command whose body is asyncio.run(...). -_started = None -_cleanup_ran = None +# Rebound per test by the _events fixture. +_started = threading.Event() +_cleanup_ran = threading.Event() @click.command() @@ -98,12 +101,12 @@ def _blocking_command() -> None: @pytest.fixture(autouse=True) def _events(): global _started, _cleanup_ran - _started = __import__("threading").Event() - _cleanup_ran = __import__("threading").Event() + _started = threading.Event() + _cleanup_ran = threading.Event() yield -async def _wait_for(event, timeout=10.0): +async def _wait_for(event: threading.Event, timeout: float = 10.0) -> None: deadline = asyncio.get_running_loop().time() + timeout while not event.is_set(): if asyncio.get_running_loop().time() > deadline: From 31710dd851fcf9d6640998e3d52fd30aaf973b15 Mon Sep 17 00:00:00 2001 From: Robert Ursu Date: Tue, 4 Aug 2026 15:41:10 +0300 Subject: [PATCH 6/6] fix(cli): bound how much log a single drain reads into memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tailer did f.read() with no limit, so one drain pulled everything written since the last poll — and the decoded str costs more again. Normally the 250 ms poll keeps that tiny, but a job that logs heavily between polls, a starved tailer, or the final drain after a long run could all pull a large log in at once. The same read-it-all-to-forward-it shape has OOMed pods on the handler side. Each pass now reads at most LOG_READ_CHUNK_BYTES and loops until caught up, so peak memory is the window rather than the backlog. Two cases only reachable once the read is bounded: - A chunk boundary landing mid-line: trimmed back to the last newline, which is also what keeps the decode safe, since \n never appears inside a multi-byte UTF-8 sequence. - A single line longer than the window: emitted as a fragment rather than held. Buffering it without bound would defeat the point, and stalling on it would wedge the tailer permanently. Co-Authored-By: Claude Opus 5 --- .../uipath/src/uipath/_cli/_server_jobs.py | 61 +++++++++++++------ packages/uipath/tests/cli/test_server_logs.py | 49 +++++++++++++++ 2 files changed, 90 insertions(+), 20 deletions(-) diff --git a/packages/uipath/src/uipath/_cli/_server_jobs.py b/packages/uipath/src/uipath/_cli/_server_jobs.py index a984c3886..a9662974a 100644 --- a/packages/uipath/src/uipath/_cli/_server_jobs.py +++ b/packages/uipath/src/uipath/_cli/_server_jobs.py @@ -189,6 +189,10 @@ def build_result_payload(job_key: str, outcome: dict[str, Any]) -> dict[str, Any LOG_POLL_SECONDS = 0.25 LOG_BATCH_MAX_LINES = 200 +# Cap on how much log a single drain pulls into memory. Without it a job that logs +# heavily between polls -- or the final drain after a long run -- reads the whole +# remainder at once, and the decoded str costs more again. +LOG_READ_CHUNK_BYTES = 256 * 1024 LOG_FLUSH_TIMEOUT_SECONDS = 10 # How long a cancelled job gets to unwind before we admit the stop did not take. @@ -249,31 +253,48 @@ def stop(self) -> None: self._stop.set() async def _drain(self, final: bool = False) -> None: - try: - if not os.path.exists(self.path): + """Forward whatever is new, a bounded window at a time until caught up.""" + while True: + 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(LOG_READ_CHUNK_BYTES) + except OSError: 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: + if not chunk: return - chunk = chunk[: cut + 1] - self._offset += len(chunk) + # A short read means we reached the end of the file; a full one means there + # is more waiting and we should come back round. + more_pending = len(chunk) == LOG_READ_CHUNK_BYTES + + if not chunk.endswith(b"\n"): + # A logging handler writes the record and only then flushes, so the tail + # can be half a line. + cut = chunk.rfind(b"\n") + if cut != -1: + # Trim to the last complete line. Safe to decode: \n never appears + # inside a multi-byte UTF-8 sequence. + chunk = chunk[: cut + 1] + elif not (more_pending or final): + # Partial line still being written — wait for its newline. + return + # Otherwise: a single line longer than the read window, or the last drain + # with nothing more coming. Emit the fragment rather than buffering + # without bound. + + self._offset += len(chunk) + await self._emit(chunk) + + if not more_pending: + return + async def _emit(self, chunk: bytes) -> None: batch: list[dict[str, Any]] = [] for raw in chunk.decode("utf-8", errors="replace").splitlines(): line = raw.rstrip("\r") diff --git a/packages/uipath/tests/cli/test_server_logs.py b/packages/uipath/tests/cli/test_server_logs.py index 3b8094765..69ec80625 100644 --- a/packages/uipath/tests/cli/test_server_logs.py +++ b/packages/uipath/tests/cli/test_server_logs.py @@ -294,3 +294,52 @@ async def test_tailer_tolerates_a_missing_file(tmp_path): await asyncio.wait_for(tailer.run(), timeout=5) assert callback.batches == [] + + +# --------------------------------------------------------------------------- # +# bounded reads # +# --------------------------------------------------------------------------- # + + +async def test_tailer_catches_up_across_multiple_read_windows(tmp_path, monkeypatch): + """A job that logs more than one window between polls must still be forwarded in + full — bounding the read must not mean dropping the remainder.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "LOG_READ_CHUNK_BYTES", 512) + + log_file = tmp_path / "execution.log" + total = 200 # ~40 bytes each, comfortably several windows + log_file.write_bytes( + b"".join(b"[2026-07-29 17:00:00,000][INFO] line-%d\n" % i for i in range(total)) + ) + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + tailer.stop() + await asyncio.wait_for(tailer.run(), timeout=10) + + messages = [line["message"] for line in callback.lines] + assert len(messages) == total + assert messages[0] == "line-0" + assert messages[-1] == f"line-{total - 1}" + + +async def test_tailer_emits_a_line_longer_than_the_read_window(tmp_path, monkeypatch): + """A single pathologically long line must not stall the tailer forever: bounded + memory beats perfect line integrity here.""" + import uipath._cli._server_jobs as jobs + + monkeypatch.setattr(jobs, "LOG_READ_CHUNK_BYTES", 64) + + log_file = tmp_path / "execution.log" + log_file.write_bytes(b"x" * 500 + b"\n[2026-07-29 17:00:00,000][INFO] after\n") + + callback = RecordingCallback() + tailer = JobLogTailer("job-1", str(log_file), callback) + tailer.stop() + await asyncio.wait_for(tailer.run(), timeout=10) + + joined = "".join(line["message"] for line in callback.lines) + assert "x" * 500 in joined, "the long line must be forwarded, even if fragmented" + assert any(line["message"] == "after" for line in callback.lines)