From 723e5dcf933773966db9d304daa170513040bee0 Mon Sep 17 00:00:00 2001 From: Naman Tyagi Date: Thu, 13 Aug 2026 11:17:04 +0530 Subject: [PATCH] [agentserver] Add process-isolated hard execution cap for task timeouts (spawn + fork) Enforce a hard cap on per-turn task timeouts by running the handler in an isolated child process the timeout watchdog can SIGKILL after a bounded grace, without disrupting the main container or co-located tasks, and without the recovery system resurrecting the force-stopped turn. - Opt-in, default off (AGENTSERVER_TASK_ISOLATION); zero overhead when off. - Tiered watchdog: cooperative cancel -> grace -> SIGKILL -> existing finalization (one-shot delete / multi-turn drain-or-suspend). - Spawn backend (create_subprocess_exec, portable) + optional per-chain reuse with idle-TTL (AGENTSERVER_TASK_WORKER_REUSE). - Fork backend (multiprocessing fork, Linux-only, AGENTSERVER_TASK_WORKER_FORK): inherits the imported app (COW, ~ms start, no re-import), reuses all shared worker code via a _ForkProcAdapter; child sanitized via _after_fork_child. - Durable terminal-outcome IPC contract; marker + backstop recovery detection. - Design doc: docs/task-timeout-isolation-design.md. - Tests: spawn + fork one-shot/multi-turn/reuse (green on Linux; fork skips on non-Linux). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6a7cdfb6-6f7c-483b-adef-432cc13f4149 --- .../ai/agentserver/core/tasks/_decorator.py | 30 + .../ai/agentserver/core/tasks/_isolation.py | 1172 +++++++++++++++++ .../ai/agentserver/core/tasks/_manager.py | 666 +++++++++- .../docs/task-timeout-isolation-design.md | 414 ++++++ .../tests/tasks/_isolation_handlers.py | 102 ++ .../tests/tasks/test_isolation.py | 82 ++ .../tasks/test_isolation_fork_multiturn.py | 183 +++ .../tasks/test_isolation_fork_oneshot.py | 109 ++ .../tests/tasks/test_isolation_multiturn.py | 99 ++ .../tests/tasks/test_isolation_oneshot.py | 161 +++ .../tests/tasks/test_isolation_reuse.py | 149 +++ 11 files changed, 3144 insertions(+), 23 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_isolation.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/docs/task-timeout-isolation-design.md create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/_isolation_handlers.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_multiturn.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_oneshot.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_multiturn.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_oneshot.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_reuse.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py index f842d67515f3..903ae20cd66d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_decorator.py @@ -64,6 +64,36 @@ async def my_task(ctx: TaskContext[MyInput]) -> MyOutput: _MAX_TASK_TIMEOUT = timedelta(days=7) +def _resolve_hard_stop_grace() -> timedelta: + """Spec 037 #8 (hard execution cap) — grace after the cooperative timeout + cancel before the framework FORCE-stops the handler. + + Defaults to 1 hour; overridable via + ``AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS`` (used to shorten the + window for tests / the hosted POC). Non-positive or unparseable values fall + back to the 1-hour default. + + :return: The hard-stop grace window. + :rtype: ~datetime.timedelta + """ + import os # pylint: disable=import-outside-toplevel + + raw = os.environ.get("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "").strip() + if raw: + try: + secs = float(raw) + if secs > 0: + return timedelta(seconds=secs) + except ValueError: + pass + return timedelta(hours=1) + + +#: Hard execution cap grace (see :func:`_resolve_hard_stop_grace`). Resolved at +# import; the watchdog reads it per turn so an env override takes effect. +_TIMEOUT_HARD_STOP_GRACE = _resolve_hard_stop_grace() + + def _validate_task_name(name: str | None) -> None: """Spec 037 #7 — ``name`` is a required, explicit, stable identity anchor. diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_isolation.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_isolation.py new file mode 100644 index 000000000000..e34bd202d045 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_isolation.py @@ -0,0 +1,1172 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Process-isolated handler execution for resilient tasks. + +The user's task handler ``fn(ctx)`` is run in a **separate child Python +process** launched via :func:`asyncio.create_subprocess_exec`. The main +process keeps ownership of the task store, lease, watchdog, and drain +logic; the child is a pure compute unit. This makes a per-turn ``timeout`` +enforceable: if a handler ignores the cooperative cancel, the main process +can ``kill()`` the child — a guaranteed stop that works even on a CPU-bound +loop, without corrupting the main process or co-located tasks. + +Cross-OS: uses ``create_subprocess_exec`` (spawn) so it runs on Linux, +macOS, and Windows. The handler is resolved **by name** in the child (never +pickled); ``ctx`` is proxied over a small length-prefixed-JSON protocol. + +Channels (from the child's point of view): + +* **stdin** — control messages from main (cancel / shutdown / round-trip + responses). +* **stdout** — protocol messages to main (stream emits, metadata flush / + pending-count requests, terminal outcome). The child redirects its + ``sys.stdout`` to stderr on startup so user ``print()`` cannot corrupt the + protocol; the true fd 1 is dup'd to a private protocol fd. +* **stderr** — user logs, forwarded to the main process logger. + +This module is BOTH the parent-side runner (imported by the task manager) +AND the child entrypoint (``python -m +azure.ai.agentserver.core.tasks._isolation``). +""" + +from __future__ import annotations + +import asyncio # pylint: disable=do-not-import-asyncio +import json +import logging +import multiprocessing +import os +import socket +import struct +import sys +import threading +import traceback +from typing import Any, Awaitable, Callable, Optional + +logger = logging.getLogger("azure.ai.agentserver.tasks.isolation") + +# --------------------------------------------------------------------------- +# Wire protocol — length-prefixed (4-byte big-endian) UTF-8 JSON. +# --------------------------------------------------------------------------- + +_LEN = struct.Struct(">I") + +# Parent -> child +MSG_RUN = "run" +MSG_CANCEL = "cancel" +MSG_SHUTDOWN = "shutdown" +MSG_RESP = "resp" # response to a child->parent round-trip request + +# Child -> parent +MSG_EMIT = "emit" +MSG_STREAM_CLOSE = "stream_close" +MSG_REQ = "req" # round-trip request (flush / last_cursor / pending_count) +MSG_RESULT = "result" +MSG_ERROR = "error" +MSG_SUSPEND = "suspend" +MSG_EXIT_FOR_RECOVERY = "exit_for_recovery" + +# Round-trip request kinds (payload of MSG_REQ) +REQ_FLUSH = "flush" +REQ_LAST_CURSOR = "last_cursor" +REQ_PENDING_COUNT = "pending_count" + + +def _pack(obj: dict) -> bytes: + body = json.dumps(obj, default=_json_default).encode("utf-8") + return _LEN.pack(len(body)) + body + + +def _json_default(o: Any) -> Any: + # Best-effort fallback for non-JSON values crossing the boundary. + return repr(o) + + +def _blocking_read_exact(fd: int, n: int) -> Optional[bytes]: + """Read exactly ``n`` bytes from a raw fd; None on EOF.""" + buf = bytearray() + while len(buf) < n: + try: + chunk = os.read(fd, n - len(buf)) + except OSError: + return None + if not chunk: + return None + buf.extend(chunk) + return bytes(buf) + + +def _blocking_read_msg(fd: int) -> Optional[dict]: + header = _blocking_read_exact(fd, _LEN.size) + if header is None: + return None + (length,) = _LEN.unpack(header) + body = _blocking_read_exact(fd, length) + if body is None: + return None + return json.loads(body.decode("utf-8")) + + +# =========================================================================== +# CHILD SIDE +# =========================================================================== + + +class _HybridCancel: + """A cancellation signal usable both by a CPU-bound poller and an awaiter. + + ``is_set()`` reads a plain ``bool`` flipped directly by the control + thread, so a CPU-bound handler that polls it sees the cancel *immediately* + without the event loop needing to run. ``wait()`` awaits an + :class:`asyncio.Event` woken via ``call_soon_threadsafe`` so a + cooperating handler can ``await ctx.cancel.wait()`` as well. + + Mirrors the ``asyncio.Event`` surface the handler uses (``is_set`` / + ``wait`` / ``set``). + """ + + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop = loop + self._flag = False + self._event = asyncio.Event() + + def set_from_thread(self) -> None: + self._flag = True + try: + self._loop.call_soon_threadsafe(self._event.set) + except RuntimeError: + # Loop already closed — the flag is still observable via is_set(). + pass + + def set(self) -> None: + self._flag = True + self._event.set() + + def is_set(self) -> bool: + return self._flag + + async def wait(self) -> bool: + if self._flag: + return True + await self._event.wait() + return True + + +class _ChildProtocol: + """Child-side protocol endpoint: reads control on a thread, writes to main.""" + + def __init__(self, in_fd: int, out_fd: int, loop: asyncio.AbstractEventLoop) -> None: + self._in_fd = in_fd + self._out_fd = out_fd + self._loop = loop + self._write_lock = threading.Lock() + self._req_seq = 0 + self._pending: dict[Any, "asyncio.Future[Any]"] = {} + self.cancel = _HybridCancel(loop) + self.shutdown = _HybridCancel(loop) + self.timeout_exceeded = False + self.cancel_requested = False + self.pending = 0 + self._ctx: Any = None + self._reader_thread: Optional[threading.Thread] = None + self._closed = False + # --- reuse (persistent worker) support --------------------------- + # When enabled, subsequent RUN messages for later turns arrive on the + # control channel and are pushed onto this queue for the turn loop. + self._turn_mode = False + self._run_queue: "Optional[asyncio.Queue[Optional[dict]]]" = None + + def enable_turn_mode(self) -> None: + """Switch the protocol into persistent multi-turn mode. + + Creates the run queue the reader thread feeds with later-turn RUN + snapshots. Must be called from the event-loop thread. + """ + self._turn_mode = True + self._run_queue = asyncio.Queue() + + def begin_turn(self) -> None: + """Reset per-turn cancel state before a new turn's ctx is built. + + ``shutdown`` is process-scoped and intentionally NOT reset. ``cancel`` + and the cause booleans are per-turn — a fresh :class:`_HybridCancel` + ensures a cooperatively-honored cancel from a prior turn does not leak + into the next turn's ``ctx.cancel``. + """ + self.cancel = _HybridCancel(self._loop) + self.timeout_exceeded = False + self.cancel_requested = False + self.pending = 0 + self._ctx = None + + async def next_run(self) -> Optional[dict]: + """Await the next turn's RUN snapshot (or ``None`` on shutdown/EOF).""" + assert self._run_queue is not None + return await self._run_queue.get() + + def bind_context(self, ctx: Any) -> None: + self._ctx = ctx + + # ---- low-level send (thread-safe, blocking os.write for backpressure) --- + def _send(self, obj: dict) -> None: + data = _pack(obj) + with self._write_lock: + try: + os.write(self._out_fd, data) + except OSError: + pass + + # ---- control reader thread ------------------------------------------ + def start_reader(self) -> None: + self._reader_thread = threading.Thread( + target=self._reader_loop, name="agentserver-worker-control", daemon=True + ) + self._reader_thread.start() + + def _reader_loop(self) -> None: + while not self._closed: + msg = _blocking_read_msg(self._in_fd) + if msg is None: + # Parent closed the control channel — treat as shutdown. + self.shutdown.set_from_thread() + if self._turn_mode and self._run_queue is not None: + self._loop.call_soon_threadsafe(self._run_queue.put_nowait, None) + return + t = msg.get("t") + if t == MSG_RUN and self._turn_mode: + # Later-turn RUN (turns 2..N of a reused worker). Hand the + # snapshot to the turn loop via the run queue. + if self._run_queue is not None: + self._loop.call_soon_threadsafe(self._run_queue.put_nowait, msg.get("snapshot")) + continue + if t == MSG_CANCEL: + if msg.get("timeout_exceeded"): + self.timeout_exceeded = True + if msg.get("cancel_requested"): + self.cancel_requested = True + if "pending" in msg: + self.pending = msg["pending"] + # Set the cause booleans on ctx BEFORE cancel (C-CAN-2 ordering) + # directly from this thread so a CPU-bound handler that polls + # ctx.cancel.is_set() then reads the causes sees them set. + if self._ctx is not None: + self._ctx.timeout_exceeded = self.timeout_exceeded + self._ctx.cancel_requested = self.cancel_requested + self.cancel.set_from_thread() + elif t == MSG_SHUTDOWN: + self.shutdown.set_from_thread() + if self._turn_mode and self._run_queue is not None: + # Unblock a turn loop waiting in next_run() so it can exit. + self._loop.call_soon_threadsafe(self._run_queue.put_nowait, None) + elif t == MSG_RESP: + rid = msg.get("id") + fut = self._pending.pop(rid, None) + if fut is not None: + self._loop.call_soon_threadsafe(_safe_set_result, fut, msg.get("value")) + + # ---- round-trip request from the handler coroutine ------------------ + async def request(self, kind: str, payload: dict) -> Any: + self._req_seq += 1 + rid = self._req_seq + fut: "asyncio.Future[Any]" = self._loop.create_future() + self._pending[rid] = fut + self._send({"t": MSG_REQ, "id": rid, "kind": kind, **payload}) + return await fut + + # ---- one-way notifications ------------------------------------------ + def emit(self, stream_id: str, payload: Any, close: bool) -> None: + self._send({"t": MSG_EMIT, "stream_id": stream_id, "payload": payload, "close": close}) + + def stream_close(self, stream_id: str) -> None: + self._send({"t": MSG_STREAM_CLOSE, "stream_id": stream_id}) + + def terminal(self, obj: dict) -> None: + self._send(obj) + + def close(self) -> None: + self._closed = True + + +def _safe_set_result(fut: "asyncio.Future[Any]", value: Any) -> None: + if not fut.done(): + fut.set_result(value) + + +class _ProxyStream: + """Child-side EventStream proxy: forwards emit/close/last_cursor to main.""" + + def __init__(self, stream_id: str, proto: _ChildProtocol) -> None: + self._id = stream_id + self._proto = proto + + async def emit(self, payload: Any, *, close: bool = False) -> None: + self._proto.emit(self._id, payload, close) + + async def close(self) -> None: + self._proto.stream_close(self._id) + + async def last_cursor(self) -> Optional[int]: + return await self._proto.request(REQ_LAST_CURSOR, {"stream_id": self._id}) + + def subscribe(self, *, after: Optional[int] = None): # pragma: no cover + raise RuntimeError("subscribe() is not available inside an isolated task worker") + + +def _install_stream_proxy(proto: _ChildProtocol) -> None: + """Point the streams singleton's factory at the IPC proxy (Gap 1).""" + try: + from azure.ai.agentserver.core.streaming import streams # pylint: disable=import-outside-toplevel + + streams._factory = lambda _id: _ProxyStream(_id, proto) # noqa: SLF001 + except Exception: # pylint: disable=broad-except + logger.debug("streams proxy not installed (streaming module unavailable)", exc_info=True) + + +class _ProxyMetadata: + """Child-side TaskMetadata replacement backed by IPC flush to main. + + Behaves like the default-namespace mapping the handler uses; named + namespaces are supported via the callable protocol. Data is plain JSON. + The full state rides the terminal message so main can persist any + mutations the handler did not explicitly flush. + """ + + def __init__(self, proto: _ChildProtocol, initial: dict, namespace: Optional[str] = None, + store: "Optional[dict[Any, dict]]" = None) -> None: + self._proto = proto + self._ns = namespace + # store maps namespace(None|str) -> dict; shared across facades. + self._store: "dict[Any, dict]" = store if store is not None else {None: dict(initial or {})} + if namespace not in self._store: + self._store[namespace] = {} + + def _data(self) -> dict: + return self._store[self._ns] + + def __getitem__(self, k): return self._data()[k] + def __setitem__(self, k, v): self._data()[k] = v + def __delitem__(self, k): del self._data()[k] + def __iter__(self): return iter(self._data()) + def __len__(self): return len(self._data()) + def __contains__(self, k): return k in self._data() + + def get(self, k, default=None): return self._data().get(k, default) + def pop(self, k, default=None): return self._data().pop(k, default) + + def __call__(self, namespace: str) -> "_ProxyMetadata": + return _ProxyMetadata(self._proto, {}, namespace=namespace, store=self._store) + + async def flush(self) -> None: + await self._proto.request(REQ_FLUSH, {"namespace": self._ns, "data": dict(self._data())}) + + async def _flush_all(self) -> None: + for ns, data in self._store.items(): + await self._proto.request(REQ_FLUSH, {"namespace": ns, "data": dict(data)}) + + def snapshot(self) -> dict: + # {namespace-or-"__default__": data} + return {("__default__" if ns is None else ns): dict(d) for ns, d in self._store.items()} + + +def _build_child_context(snapshot: dict, proto: _ChildProtocol) -> Any: + """Construct a child-local TaskContext-like object for the handler.""" + from ._context import TaskContext # pylint: disable=import-outside-toplevel + + meta = _ProxyMetadata(proto, snapshot.get("metadata") or {}) + + ctx: Any = TaskContext( + task_id=snapshot["task_id"], + session_id=snapshot.get("session_id") or "", + input=snapshot.get("input"), + metadata=meta, # type: ignore[arg-type] + retry_attempt=snapshot.get("retry_attempt", 0), + recovery_count=snapshot.get("recovery_count", 0), + entry_mode=snapshot.get("entry_mode", "fresh"), + is_steered_turn=snapshot.get("is_steered_turn", False), + input_id=snapshot.get("input_id"), + ) + # Replace the cancel/shutdown events with the hybrid (thread-settable) ones. + # (cancel / shutdown / _pending_count_provider / cause booleans are all in + # TaskContext.__slots__, so we set only those — no arbitrary attributes.) + ctx.cancel = proto.cancel # type: ignore[assignment] + ctx.shutdown = proto.shutdown # type: ignore[assignment] + ctx._pending_count_provider = lambda: proto.pending # type: ignore[attr-defined] + proto.bind_context(ctx) + return ctx, meta + + +async def _run_child(snapshot: dict) -> None: + loop = asyncio.get_running_loop() + + # --- fd setup: reserve fd1 for the protocol, redirect stdout->stderr ---- + protocol_out_fd = os.dup(1) + os.dup2(2, 1) # user print() now goes to stderr + in_fd = 0 + + proto = _ChildProtocol(in_fd=in_fd, out_fd=protocol_out_fd, loop=loop) + _install_stream_proxy(proto) + + # --- resolve the handler by name (Gap 2) -------------------------------- + handler = _resolve_handler(snapshot["handler_module"], snapshot["handler_name"]) + + ctx, meta = _build_child_context(snapshot, proto) + + # Start the control reader ONLY after ctx is bound, so an early CANCEL + # (buffered in the pipe before the handler starts) still has a ctx to set + # its cause booleans on. + proto.start_reader() + + try: + result = await handler(ctx) + term = _classify_terminal(result, meta) + except BaseException as exc: # noqa: BLE001 pylint: disable=broad-except + proto.terminal({ + "t": MSG_ERROR, + "exc_type": type(exc).__name__, + "exc_msg": str(exc)[:4000], + "traceback": traceback.format_exc()[:8000], + "metadata": meta.snapshot(), + }) + else: + proto.terminal(term) + finally: + proto.close() + + +def _classify_terminal(result: Any, meta: _ProxyMetadata) -> dict: + from ._context import _ExitForRecovery, _Suspended # pylint: disable=import-outside-toplevel + + md = meta.snapshot() + if isinstance(result, _ExitForRecovery): + return {"t": MSG_EXIT_FOR_RECOVERY, "metadata": md} + if isinstance(result, _Suspended): + return {"t": MSG_SUSPEND, "reason": result.reason, "output": result.output, "metadata": md} + return {"t": MSG_RESULT, "value": result, "metadata": md} + + +def _resolve_handler(module_name: str, handler_name: str) -> Callable[..., Awaitable[Any]]: + import importlib # pylint: disable=import-outside-toplevel + + importlib.import_module(module_name) + # Optional extra bootstrap modules. + extra = os.environ.get("AGENTSERVER_WORKER_BOOTSTRAP_MODULES", "") + for m in (x.strip() for x in extra.split(",") if x.strip()): + try: + importlib.import_module(m) + except Exception: # pylint: disable=broad-except + logger.warning("worker: failed to import bootstrap module %s", m, exc_info=True) + + from ._decorator import _REGISTERED_DESCRIPTORS # pylint: disable=import-outside-toplevel + + for name, fn, _opts in _REGISTERED_DESCRIPTORS: + if name == handler_name: + return fn + raise RuntimeError(f"worker: handler {handler_name!r} not found after importing {module_name!r}") + + +async def _run_child_reusable(first_snapshot: dict) -> None: + """Persistent child loop: run turns until shutdown/EOF (reuse mode). + + Same per-turn semantics as :func:`_run_child`, but the process stays alive + between turns. fd setup and the control-reader thread are established ONCE; + each turn resets cancel state (:meth:`_ChildProtocol.begin_turn`), rebuilds + a fresh ``ctx`` from the turn's snapshot, runs the handler, emits the + per-turn terminal message (WITHOUT closing stdout), then waits for the next + RUN. Exits on MSG_SHUTDOWN or control-channel EOF. + """ + loop = asyncio.get_running_loop() + + protocol_out_fd = os.dup(1) + os.dup2(2, 1) # user print() -> stderr; true fd1 reserved for protocol + proto = _ChildProtocol(in_fd=0, out_fd=protocol_out_fd, loop=loop) + _install_stream_proxy(proto) + proto.enable_turn_mode() + + snapshot: Optional[dict] = first_snapshot + reader_started = False + while snapshot is not None: + proto.begin_turn() + handler = _resolve_handler(snapshot["handler_module"], snapshot["handler_name"]) + ctx, meta = _build_child_context(snapshot, proto) + if not reader_started: + # Start the control reader ONLY after the first ctx is bound so an + # early buffered CANCEL still has a ctx to set its causes on (same + # invariant as the one-shot path). The reader persists for the life + # of the worker and also delivers later-turn RUN messages. + proto.start_reader() + reader_started = True + try: + result = await handler(ctx) + term = _classify_terminal(result, meta) + except BaseException as exc: # noqa: BLE001 pylint: disable=broad-except + term = { + "t": MSG_ERROR, + "exc_type": type(exc).__name__, + "exc_msg": str(exc)[:4000], + "traceback": traceback.format_exc()[:8000], + "metadata": meta.snapshot(), + } + proto.terminal(term) # per-turn terminal; stdout stays open for reuse + snapshot = await proto.next_run() + proto.close() + + +def _child_main() -> None: + # First message on stdin is the RUN payload. + run_msg = _blocking_read_msg(0) + if run_msg is None or run_msg.get("t") != MSG_RUN: + sys.stderr.write("worker: expected RUN message on stdin\n") + sys.exit(2) + try: + if run_msg.get("reuse"): + asyncio.run(_run_child_reusable(run_msg["snapshot"])) + else: + asyncio.run(_run_child(run_msg["snapshot"])) + except Exception: # pylint: disable=broad-except + traceback.print_exc() + sys.exit(1) + + +# =========================================================================== +# PARENT SIDE +# =========================================================================== + + +class IsolationBridge: + """Main-process callbacks the worker proxies back to. + + All persistence, streaming, and steering state stays in main; the child + only forwards requests here. + + :param apply_flush: ``async (namespace|None, data) -> None`` — persist one + metadata namespace (also keeps main's ``ctx.metadata`` in sync). + :param stream_emit: ``async (stream_id, payload, close) -> None``. + :param stream_close: ``async (stream_id) -> None``. + :param stream_last_cursor: ``async (stream_id) -> int|None``. + :param get_pending_count: ``() -> int``. + :param apply_final_metadata: ``(snapshot: dict) -> None`` — apply the + child's terminal metadata snapshot onto main's ``ctx.metadata``. + """ + + def __init__( + self, + *, + apply_flush: Callable[[Optional[str], dict], Awaitable[None]], + stream_emit: Callable[[str, Any, bool], Awaitable[None]], + stream_close: Callable[[str], Awaitable[None]], + stream_last_cursor: Callable[[str], Awaitable[Optional[int]]], + get_pending_count: Callable[[], int], + apply_final_metadata: Callable[[dict], None], + ) -> None: + self.apply_flush = apply_flush + self.stream_emit = stream_emit + self.stream_close = stream_close + self.stream_last_cursor = stream_last_cursor + self.get_pending_count = get_pending_count + self.apply_final_metadata = apply_final_metadata + + +class WorkerCrash(Exception): + """Raised by :meth:`IsolatedRun.outcome` when the child exited without a + terminal message and was NOT killed by us (OOM/segfault/unexpected).""" + + +class IsolatedRun: + """A launched worker process running one handler turn. + + ``outcome()`` reproduces ``await fn(ctx)``: returns the handler's return + value (or a ``_Suspended`` / ``_ExitForRecovery`` sentinel), or raises the + handler's exception (or :class:`WorkerCrash` on unexpected child death). + ``kill()`` force-stops the worker (the timeout hard cap). + """ + + def __init__(self, proc: "asyncio.subprocess.Process", bridge: IsolationBridge) -> None: + self._proc = proc + self._bridge = bridge + self._we_killed_it = False + self._terminal: Optional[dict] = None + self._reader_task: Optional[asyncio.Task[None]] = None + self._stderr_task: Optional[asyncio.Task[None]] = None + self._write_lock = asyncio.Lock() + + @property + def pid(self) -> int: + return self._proc.pid + + async def _send(self, obj: dict) -> None: + assert self._proc.stdin is not None + async with self._write_lock: + self._proc.stdin.write(_pack(obj)) + try: + await self._proc.stdin.drain() + except (ConnectionResetError, BrokenPipeError): + pass + + def kill(self) -> None: + self._we_killed_it = True + try: + self._proc.kill() + except ProcessLookupError: + pass + + async def signal_cancel(self, *, timeout_exceeded: bool, cancel_requested: bool) -> None: + await self._send({ + "t": MSG_CANCEL, + "timeout_exceeded": timeout_exceeded, + "cancel_requested": cancel_requested, + }) + + async def signal_shutdown(self) -> None: + await self._send({"t": MSG_SHUTDOWN}) + + # ---- the always-draining stdout reader (Gap 1/5: no deadlock) -------- + async def _read_loop(self) -> None: + assert self._proc.stdout is not None + reader = self._proc.stdout + while True: + try: + header = await reader.readexactly(_LEN.size) + except asyncio.IncompleteReadError: + break + (length,) = _LEN.unpack(header) + try: + body = await reader.readexactly(length) + except asyncio.IncompleteReadError: + break + msg = json.loads(body.decode("utf-8")) + await self._dispatch(msg) + if msg.get("t") in (MSG_RESULT, MSG_ERROR, MSG_SUSPEND, MSG_EXIT_FOR_RECOVERY): + self._terminal = msg + # keep draining until EOF so nothing blocks the child + + async def _dispatch(self, msg: dict) -> None: + t = msg.get("t") + if t == MSG_EMIT: + await self._bridge.stream_emit(msg["stream_id"], msg["payload"], msg.get("close", False)) + elif t == MSG_STREAM_CLOSE: + await self._bridge.stream_close(msg["stream_id"]) + elif t == MSG_REQ: + await self._handle_request(msg) + # terminal messages are captured in _read_loop + + async def _handle_request(self, msg: dict) -> None: + rid = msg.get("id") + kind = msg.get("kind") + value: Any = None + try: + if kind == REQ_FLUSH: + await self._bridge.apply_flush(msg.get("namespace"), msg.get("data") or {}) + elif kind == REQ_LAST_CURSOR: + value = await self._bridge.stream_last_cursor(msg["stream_id"]) + elif kind == REQ_PENDING_COUNT: + value = self._bridge.get_pending_count() + except Exception: # pylint: disable=broad-except + logger.warning("isolation: bridge request %s failed", kind, exc_info=True) + await self._send({"t": MSG_RESP, "id": rid, "value": value}) + + async def _forward_stderr(self) -> None: + if self._proc.stderr is None: + return # fork transport: child shares the parent's stderr directly + while True: + line = await self._proc.stderr.readline() + if not line: + break + logger.info("[worker %s] %s", self._proc.pid, line.decode("utf-8", "replace").rstrip()) + + async def outcome(self) -> Any: + # Wait for the reader to finish (child closed stdout) then the process. + if self._reader_task is not None: + try: + await self._reader_task + except Exception: # pylint: disable=broad-except + logger.warning("isolation: reader loop failed", exc_info=True) + await self._proc.wait() # reap + if self._stderr_task is not None: + self._stderr_task.cancel() + + term = self._terminal + if term is not None: + self._bridge.apply_final_metadata(term.get("metadata") or {}) + return _reconstruct_outcome(term) + # No terminal message. + if self._we_killed_it: + # Intentional hard-cap kill — the caller (manager) handles the + # lifecycle handoff; signal via a sentinel exception. + raise _WorkerKilled() + raise WorkerCrash(f"worker {self._proc.pid} exited (rc={self._proc.returncode}) without a terminal message") + + +class _WorkerKilled(Exception): + """Internal: the worker was killed by us for the timeout hard cap.""" + + +def _reconstruct_outcome(term: dict) -> Any: + from ._context import _ExitForRecovery, _Suspended # pylint: disable=import-outside-toplevel + + t = term["t"] + if t == MSG_RESULT: + return term.get("value") + if t == MSG_SUSPEND: + return _Suspended(reason=term.get("reason"), output=term.get("output")) + if t == MSG_EXIT_FOR_RECOVERY: + return _ExitForRecovery() + # MSG_ERROR + exc_type = term.get("exc_type", "Exception") + exc_msg = term.get("exc_msg", "") + if exc_type in ("CancelledError", "TaskCancelled"): + raise asyncio.CancelledError() + raise _IsolatedHandlerError(exc_type, exc_msg, term.get("traceback", "")) + + +class _IsolatedHandlerError(Exception): + """A handler exception reconstructed across the process boundary.""" + + def __init__(self, exc_type: str, message: str, tb: str) -> None: + super().__init__(f"{exc_type}: {message}") + self.exc_type = exc_type + self.original_message = message + self.remote_traceback = tb + + +def _worker_command(python_exe: Optional[str] = None) -> list[str]: + return [python_exe or sys.executable, "-m", "azure.ai.agentserver.core.tasks._isolation"] + + +async def start_isolated(snapshot: dict, bridge: IsolationBridge, *, python_exe: Optional[str] = None) -> IsolatedRun: + """Launch a worker process and begin pumping its protocol. + + :param snapshot: The ctx snapshot + handler_module/handler_name. + :param bridge: Main-process callbacks the worker proxies to. + :return: A running :class:`IsolatedRun`. + """ + proc = await asyncio.create_subprocess_exec( + *_worker_command(python_exe), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=os.environ.copy(), + ) + run = IsolatedRun(proc, bridge) + run._reader_task = asyncio.create_task(run._read_loop()) # noqa: SLF001 + run._stderr_task = asyncio.create_task(run._forward_stderr()) # noqa: SLF001 + await run._send({"t": MSG_RUN, "snapshot": snapshot}) # noqa: SLF001 + return run + + +# =========================================================================== +# PERSISTENT (per-chain reuse) WORKER — §13.6 +# =========================================================================== + + +class PersistentWorker: + """A long-lived worker process reused across the turns of one chain. + + Unlike :class:`IsolatedRun` (one process per turn), a ``PersistentWorker`` + imports the app **once** and runs many turns via :meth:`run_turn`. The child + stays alive between turns (waiting for the next RUN), so an N-turn chain + pays the ~1-2s import ONCE instead of N times. The main process still owns + persistence/streaming/steering (proxied per turn via a fresh + :class:`IsolationBridge`) and can :meth:`kill` the worker for the timeout + hard cap exactly like the per-turn path. + + Lifecycle is managed by the task manager: a registry keyed by ``task_id`` + holds the worker across suspend/resume; it is torn down on hard-cap kill, + crash, idle-TTL eviction, or manager shutdown. + """ + + def __init__(self, proc: "asyncio.subprocess.Process") -> None: + self._proc = proc + self._current_bridge: Optional[IsolationBridge] = None + self._turn_future: "Optional[asyncio.Future[dict]]" = None + self._reader_task: Optional[asyncio.Task[None]] = None + self._stderr_task: Optional[asyncio.Task[None]] = None + self._write_lock = asyncio.Lock() + self._alive = True + self._we_killed = False + self.in_flight = False + try: + self.last_active_monotonic = asyncio.get_running_loop().time() + except RuntimeError: + self.last_active_monotonic = 0.0 + + @property + def pid(self) -> int: + return self._proc.pid + + @property + def alive(self) -> bool: + return self._alive + + async def _send(self, obj: dict) -> None: + if self._proc.stdin is None: + return + async with self._write_lock: + try: + self._proc.stdin.write(_pack(obj)) + await self._proc.stdin.drain() + except (ConnectionResetError, BrokenPipeError): + pass + + async def run_turn(self, snapshot: dict, bridge: IsolationBridge) -> Any: + """Run one turn on this worker; reproduces ``await fn(ctx)``. + + :param snapshot: The ctx snapshot + handler_module/handler_name. + :param bridge: Per-turn main-process callbacks the worker proxies to. + :return: The handler's return value / suspend / exit sentinel. + :raises _WorkerKilled: The worker was hard-cap killed during this turn. + :raises WorkerCrash: The worker died unexpectedly during this turn. + """ + if not self._alive: + raise WorkerCrash("persistent worker is not alive") + loop = asyncio.get_running_loop() + self._current_bridge = bridge + self._turn_future = loop.create_future() + self.in_flight = True + try: + await self._send({"t": MSG_RUN, "snapshot": snapshot, "reuse": True}) + term = await self._turn_future + finally: + self.in_flight = False + self._turn_future = None + self._current_bridge = None + self.last_active_monotonic = loop.time() + bridge.apply_final_metadata(term.get("metadata") or {}) + return _reconstruct_outcome(term) + + def kill(self) -> None: + """Force-stop the worker (timeout hard cap). Idempotent.""" + self._we_killed = True + self._alive = False + try: + self._proc.kill() + except ProcessLookupError: + pass + + async def signal_cancel(self, *, timeout_exceeded: bool, cancel_requested: bool) -> None: + await self._send({ + "t": MSG_CANCEL, + "timeout_exceeded": timeout_exceeded, + "cancel_requested": cancel_requested, + }) + + async def signal_shutdown(self) -> None: + await self._send({"t": MSG_SHUTDOWN}) + + def idle_seconds(self, now: float) -> float: + return now - self.last_active_monotonic + + async def _read_loop(self) -> None: + assert self._proc.stdout is not None + reader = self._proc.stdout + while True: + try: + header = await reader.readexactly(_LEN.size) + except asyncio.IncompleteReadError: + break + (length,) = _LEN.unpack(header) + try: + body = await reader.readexactly(length) + except asyncio.IncompleteReadError: + break + msg = json.loads(body.decode("utf-8")) + t = msg.get("t") + if t in (MSG_RESULT, MSG_ERROR, MSG_SUSPEND, MSG_EXIT_FOR_RECOVERY): + fut = self._turn_future + if fut is not None and not fut.done(): + fut.set_result(msg) + else: + await self._dispatch(msg) + # stdout EOF — the child exited. + self._alive = False + fut = self._turn_future + if fut is not None and not fut.done(): + if self._we_killed: + fut.set_exception(_WorkerKilled()) + else: + fut.set_exception( + WorkerCrash(f"persistent worker {self._proc.pid} exited (rc={self._proc.returncode})") + ) + + async def _dispatch(self, msg: dict) -> None: + bridge = self._current_bridge + if bridge is None: + return + t = msg.get("t") + if t == MSG_EMIT: + await bridge.stream_emit(msg["stream_id"], msg["payload"], msg.get("close", False)) + elif t == MSG_STREAM_CLOSE: + await bridge.stream_close(msg["stream_id"]) + elif t == MSG_REQ: + await self._handle_request(msg, bridge) + + async def _handle_request(self, msg: dict, bridge: IsolationBridge) -> None: + rid = msg.get("id") + kind = msg.get("kind") + value: Any = None + try: + if kind == REQ_FLUSH: + await bridge.apply_flush(msg.get("namespace"), msg.get("data") or {}) + elif kind == REQ_LAST_CURSOR: + value = await bridge.stream_last_cursor(msg["stream_id"]) + elif kind == REQ_PENDING_COUNT: + value = bridge.get_pending_count() + except Exception: # pylint: disable=broad-except + logger.warning("persistent worker: bridge request %s failed", kind, exc_info=True) + await self._send({"t": MSG_RESP, "id": rid, "value": value}) + + async def _forward_stderr(self) -> None: + if self._proc.stderr is None: + return # fork transport: child shares the parent's stderr directly + while True: + line = await self._proc.stderr.readline() + if not line: + break + logger.info("[worker %s] %s", self._proc.pid, line.decode("utf-8", "replace").rstrip()) + + async def aclose(self) -> None: + """Graceful teardown: ask the child to exit, else kill; reap tasks.""" + if self._alive: + try: + await self.signal_shutdown() + except Exception: # pylint: disable=broad-except + pass + try: + await asyncio.wait_for(self._proc.wait(), timeout=5) + except Exception: # pylint: disable=broad-except + try: + self._proc.kill() + except ProcessLookupError: + pass + self._alive = False + for task in (self._reader_task, self._stderr_task): + if task is not None and not task.done(): + task.cancel() + + +async def start_persistent_worker(*, python_exe: Optional[str] = None) -> PersistentWorker: + """Launch a reusable worker process and begin pumping its protocol. + + No RUN is sent here; the first (and every subsequent) turn is started by + :meth:`PersistentWorker.run_turn`. + + :return: A ready :class:`PersistentWorker`. + """ + proc = await asyncio.create_subprocess_exec( + *_worker_command(python_exe), + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=os.environ.copy(), + ) + worker = PersistentWorker(proc) + worker._reader_task = asyncio.create_task(worker._read_loop()) # noqa: SLF001 + worker._stderr_task = asyncio.create_task(worker._forward_stderr()) # noqa: SLF001 + return worker + + +# =========================================================================== +# FORK BACKEND (Unix) — coexists with the spawn path +# =========================================================================== +# +# The spawn path launches a blank interpreter (`python -m ..._isolation`) that +# re-imports the app (~1-2s) and rebuilds ctx from a snapshot. The fork path +# instead ``os.fork()``s a child from the (already-imported) parent, so the app +# + all modules are inherited (no re-import, COW memory, ~ms start). We reuse +# ALL higher-level logic: the same ``_ChildProtocol`` (over one socket instead +# of std pipes), the same ``_build_child_context`` / ``_resolve_handler`` (whose +# ``import_module`` is a no-op hit on the inherited ``sys.modules``), and the +# same parent-side ``IsolatedRun`` / ``PersistentWorker`` (via an adapter that +# presents the mp process + socket like a subprocess). Fork-only concern: the +# child must sanitize the inherited async/threaded parent (fresh event loop, +# close the inherited peer socket, neutralize the inherited task-manager) — see +# ``_after_fork_child``. + + +def _after_fork_child(*, sock_fd: int, dead_fds: "list[int]") -> None: + """Sanitize a freshly-forked child before it runs any handler. + + The child inherited the parent's entire memory image — a running event loop, + reader threads' locks, open store/lease sockets, and the live task-manager + singleton. Make the child a clean compute unit: + + * neutralize the inherited task-manager so ``get_task_manager()`` raises just + like in the spawn child (no accidental direct store/lease access → + no split-brain), + * close the inherited peer (parent-side) socket fds so parent-death EOF works + and the fd is not leaked, + * drop the inherited event loop (a fresh one is created by ``asyncio.run``). + + :keyword sock_fd: The child's own IPC socket fd (kept open). + :keyword dead_fds: Inherited fds to close (e.g. the parent-side socketpair end). + """ + try: + from ._manager import set_task_manager # pylint: disable=import-outside-toplevel + + set_task_manager(None) + except Exception: # pylint: disable=broad-except + pass + for fd in dead_fds: + if fd == sock_fd: + continue + try: + os.close(fd) + except OSError: + pass + # The inherited asyncio loop must never be used; asyncio.run() in the child + # entry creates and owns a brand-new loop. + + +async def _fork_child_amain(sock_fd: int, reuse: bool) -> None: + """Child-side async entry for the fork backend. + + Uses one bidirectional socket (``sock_fd``) for the whole protocol instead of + the spawn path's stdin/stdout/stderr split. Turn-mode is always enabled so + the RUN message(s) arrive through the reader; one-shot processes exactly one + RUN then returns (process exit → parent EOF), reuse loops until shutdown/EOF. + """ + loop = asyncio.get_running_loop() + proto = _ChildProtocol(in_fd=sock_fd, out_fd=sock_fd, loop=loop) + _install_stream_proxy(proto) + proto.enable_turn_mode() + proto.start_reader() + + if reuse: + first = True + while True: + snapshot = await proto.next_run() + if snapshot is None: + break + await _fork_run_one_turn(proto, snapshot, reset=not first) + first = False + else: + snapshot = await proto.next_run() + if snapshot is not None: + await _fork_run_one_turn(proto, snapshot, reset=False) + proto.close() + + +async def _fork_run_one_turn(proto: _ChildProtocol, snapshot: dict, *, reset: bool) -> None: + """Run a single turn in the forked child and emit its terminal message.""" + if reset: + proto.begin_turn() + handler = _resolve_handler(snapshot["handler_module"], snapshot["handler_name"]) + ctx, meta = _build_child_context(snapshot, proto) # binds ctx + swaps IO handles + # A CANCEL may have arrived between the reader starting and ctx binding; + # re-sync the cause booleans (the cancel Event itself is shared via + # ctx.cancel = proto.cancel, so is_set() already reflects it). + ctx.timeout_exceeded = proto.timeout_exceeded + ctx.cancel_requested = proto.cancel_requested + try: + result = await handler(ctx) + term = _classify_terminal(result, meta) + except BaseException as exc: # noqa: BLE001 pylint: disable=broad-except + term = { + "t": MSG_ERROR, + "exc_type": type(exc).__name__, + "exc_msg": str(exc)[:4000], + "traceback": traceback.format_exc()[:8000], + "metadata": meta.snapshot(), + } + proto.terminal(term) + + +def _fork_child_entry(child_sock: "socket.socket", parent_fd: int, reuse: bool) -> None: + """Process target for a forked worker (runs in the child).""" + sock_fd = child_sock.fileno() + _after_fork_child(sock_fd=sock_fd, dead_fds=[parent_fd]) + try: + asyncio.run(_fork_child_amain(sock_fd, reuse)) + except Exception: # pylint: disable=broad-except + traceback.print_exc() + + +class _ForkProcAdapter: + """Presents a forked ``multiprocessing.Process`` + socket like a subprocess. + + Exposes the small surface ``IsolatedRun`` / ``PersistentWorker`` use + (``stdin`` writer, ``stdout`` reader, ``stderr`` None, ``pid``, ``kill()``, + ``wait()``, ``returncode``) so all their read-loop / dispatch / run_turn / + kill logic works unchanged over the fork transport. + """ + + def __init__( + self, + proc: "multiprocessing.process.BaseProcess", + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + self._proc = proc + self.stdin = writer + self.stdout = reader + self.stderr = None + + @property + def pid(self) -> int: + return self._proc.pid or -1 + + @property + def returncode(self) -> Optional[int]: + return self._proc.exitcode + + def kill(self) -> None: + try: + if self._proc.is_alive(): + self._proc.kill() + except (ProcessLookupError, ValueError, AssertionError): + pass + + async def wait(self) -> Optional[int]: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self._proc.join) + return self._proc.exitcode + + +async def _start_fork_process(reuse: bool) -> "tuple[_ForkProcAdapter, socket.socket]": + """Fork a worker and wrap its process+socket in a ``_ForkProcAdapter``.""" + parent_sock, child_sock = socket.socketpair() + ctx = multiprocessing.get_context("fork") + proc = ctx.Process( + target=_fork_child_entry, + args=(child_sock, parent_sock.fileno(), reuse), + name="agentserver-fork-worker", + daemon=False, + ) + proc.start() + child_sock.close() # parent keeps only its end + parent_sock.setblocking(False) + reader, writer = await asyncio.open_connection(sock=parent_sock) + return _ForkProcAdapter(proc, reader, writer), parent_sock + + +async def start_forked(snapshot: dict, bridge: IsolationBridge) -> IsolatedRun: + """Fork a one-shot worker (per-turn) — the fork analogue of :func:`start_isolated`. + + :param snapshot: The ctx snapshot + handler_module/handler_name. + :param bridge: Main-process callbacks the worker proxies to. + :return: A running :class:`IsolatedRun` over the fork transport. + """ + adapter, _sock = await _start_fork_process(reuse=False) + run = IsolatedRun(adapter, bridge) # type: ignore[arg-type] + run._reader_task = asyncio.create_task(run._read_loop()) # noqa: SLF001 + # No stderr task: the forked child shares the parent's stderr directly. + await run._send({"t": MSG_RUN, "snapshot": snapshot}) # noqa: SLF001 + return run + + +async def start_forked_worker() -> PersistentWorker: + """Fork a reusable worker — the fork analogue of :func:`start_persistent_worker`. + + No RUN is sent here; each turn is started by :meth:`PersistentWorker.run_turn`. + + :return: A ready :class:`PersistentWorker` over the fork transport. + """ + adapter, _sock = await _start_fork_process(reuse=True) + worker = PersistentWorker(adapter) # type: ignore[arg-type] + worker._reader_task = asyncio.create_task(worker._read_loop()) # noqa: SLF001 + # No stderr task: the forked child shares the parent's stderr directly. + return worker + + +if __name__ == "__main__": + _child_main() \ No newline at end of file diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py index 1cff24b64b0a..5768e7fc195f 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/tasks/_manager.py @@ -13,6 +13,7 @@ import asyncio # pylint: disable=do-not-import-asyncio import logging import os +import sys import traceback from collections.abc import Awaitable, Callable, Mapping from typing import Any, TypeVar @@ -34,7 +35,14 @@ _ref_key, _resolve_input_storage, ) -from ._decorator import TaskOptions, _deserialize_input, _resolve_effective_timeout, _serialize_input +from ._decorator import ( + TaskOptions, + _TIMEOUT_HARD_STOP_GRACE, + _deserialize_input, + _resolve_effective_timeout, + _resolve_hard_stop_grace, + _serialize_input, +) from ._exceptions import ( EtagConflict, OutputTooLarge, @@ -140,6 +148,94 @@ def _is_evicted(exc: BaseException) -> bool: # budget = max(0, opts.timeout - (now - _turn_started_at)). _TURN_STARTED_AT_KEY: str = "turn_started_at" +# Spec 037 #8 (hard execution cap) — top-level payload field storing the +# ISO-8601 UTC timestamp of the moment the per-turn timeout fired its +# cooperative cancel. Written when the watchdog fires; cleared atomically with +# the ``turn_started_at`` re-stamp at every turn-start boundary. Recovery treats +# a task carrying a marker (correlated: marker >= turn_started_at) as +# "already timed out" and finalizes it instead of re-running the turn. +_TIMEOUT_CANCELLED_AT_KEY: str = "timeout_cancelled_at" + +# Isolation feature flag. Process-isolated handler execution (the hard-cap +# enforcement mechanism) is opt-in via env while under development / for the +# hosted POC. When off, handlers run in-process as before (cooperative-only, +# no hard kill). Read live (not cached) so tests can toggle it per-case. +def _isolation_enabled() -> bool: + return os.environ.get("AGENTSERVER_TASK_ISOLATION", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +# Per-chain persistent worker reuse (§13.6). Layered ON TOP of isolation and +# applied ONLY to multi-turn chains: instead of spawning a fresh child per turn +# (paying the ~1-2s re-import N times for an N-turn chain), a single worker is +# created per ``task_id`` and reused across the chain's turns. Opt-in, read live. +# When off (default) OR for one-shot tasks, the per-turn-spawn path is used +# unchanged. A warm worker survives the suspend gap between turns; the idle-TTL +# reaper evicts workers idle beyond the TTL so long/indefinite parks release memory. +def _worker_reuse_enabled() -> bool: + return os.environ.get("AGENTSERVER_TASK_WORKER_REUSE", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _resolve_worker_idle_ttl() -> float: + """Seconds a reused per-chain worker may sit idle before eviction. + + Configured via ``AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS`` (default 300). + Values <= 0 disable the TTL (workers live until kill/crash/shutdown). + """ + raw = os.environ.get("AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS", "").strip() + if not raw: + return 300.0 + try: + return float(raw) + except ValueError: + logger.warning( + "invalid AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS=%r; using default 300s", raw + ) + return 300.0 + + +# Fork worker backend (§ fork-worker-design). Unix-only: when on AND isolation +# is on, worker children are created via ``fork`` (inherit the imported app — +# no re-import, COW memory, ~ms start) instead of ``create_subprocess_exec`` +# (spawn). Ignored on non-Linux (no fork / unsafe) → spawn is used. Opt-in, +# read live; composes with reuse. +def _worker_fork_enabled() -> bool: + if not sys.platform.startswith("linux"): + return False + return os.environ.get("AGENTSERVER_TASK_WORKER_FORK", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _apply_namespace(metadata: Any, namespace: Optional[str], data: dict) -> None: + """Replace one metadata namespace's contents (main-process side). + + Used to keep main's ``ctx.metadata`` in sync with an isolated worker's + flushes / terminal snapshot so the end-of-turn ``_flush_all`` is consistent. + + :param metadata: The main-process ``TaskMetadata``. + :param namespace: ``None`` for the default namespace, else the name. + :param data: The full namespace contents to install. + """ + try: + facade = metadata if namespace is None else metadata(namespace) + facade.clear() + facade.update(data) + except Exception: # pylint: disable=broad-except + logger.warning("failed to apply metadata namespace %r from worker", namespace, exc_info=True) + def _utc_now_iso() -> str: """Return current UTC time as an ISO-8601 string with a ``+00:00`` offset. @@ -344,6 +440,9 @@ class _ActiveTask: # pylint: disable=too-many-instance-attributes # slot or it is unsettable (the historic bug: it was read but never # storable). "_pending_input_count", + # Hard-cap: the currently-running isolated worker (if isolation is on), + # so the timeout watchdog can force-kill it. + "isolated_run", ) def __init__( @@ -373,6 +472,9 @@ def __init__( self.input_type = input_type self.opts = opts self.retry = retry + # Hard-cap: the in-flight isolated worker (set by _run_handler when + # isolation is enabled); the timeout watchdog force-kills it. + self.isolated_run: Any = None # ``asyncio.get_running_loop().time()`` value at the last successful # lease refresh -- updated by the renewal loop AND by every # payload PATCH that piggybacks lease ownership (see @@ -448,6 +550,13 @@ def __init__( # fresh one bound to the new turn's _turn_started_at. Cleared # on terminal exit. self._timeout_watchdogs: dict[str, asyncio.Task[None]] = {} + # §13.6 — per-chain persistent worker reuse registry. Keyed by + # task_id; survives across suspend/resume (unlike _active_tasks which + # is popped at every suspend boundary). Torn down on hard-cap kill, + # crash, idle-TTL eviction, or manager shutdown. Only populated when + # AGENTSERVER_TASK_WORKER_REUSE is on AND the task is multi-turn. + self._reuse_workers: dict[str, Any] = {} + self._worker_reaper_task: "asyncio.Task[None] | None" = None @staticmethod def _build_source(fn_name: str) -> dict[str, str]: @@ -809,6 +918,22 @@ async def shutdown(self) -> None: pass self._periodic_recovery_task = None + # §13.6 — stop the idle-TTL reaper and tear down all persistent + # reuse workers so no child processes are orphaned on shutdown. + if self._worker_reaper_task is not None: + self._worker_reaper_task.cancel() + try: + await self._worker_reaper_task + except (asyncio.CancelledError, Exception): # pylint: disable=broad-exception-caught + pass + self._worker_reaper_task = None + if self._reuse_workers: + for tid, worker in list(self._reuse_workers.items()): + try: + await worker.aclose() + except Exception: # pylint: disable=broad-exception-caught + logger.warning("failed to close reuse worker for task %s", tid, exc_info=True) + self._reuse_workers.clear() # Signal shutdown on all active contexts. Yield once so the bridge # tasks (running in the event loop) get a chance to observe the # shutdown event and notify their handlers before we proceed — @@ -1388,6 +1513,9 @@ async def _start_existing_task( # pylint: disable=too-many-locals,too-many-stat turn_start_payload: dict[str, Any] = {} if entry_mode != "recovered": turn_start_payload[_TURN_STARTED_AT_KEY] = _utc_now_iso() + # Clear the hard-cap marker atomically with the turn-start stamp so + # it is strictly per-turn (recovery correlates marker >= turn_start). + turn_start_payload[_TIMEOUT_CANCELLED_AT_KEY] = None # / SOT §11/§20: the framework does not write # payload["output"] at any point. No clear is needed on resume. @@ -1603,30 +1731,27 @@ async def _timeout_watchdog( cancel_event: asyncio.Event, ctx: "TaskContext[Any] | None" = None, *, + task_id: str | None = None, remaining_seconds: float | None = None, ) -> None: - """/: per-turn timeout watchdog. + """Per-turn timeout watchdog with a hard execution cap. - Cooperative-only. On firing, sets ``ctx.timeout_exceeded = True`` - then sets ``cancel_event`` and exits. Does NOT cancel the lease - renewal or force-stop the handler. An ignoring handler runs - until process death or external :meth:`TaskRun.cancel`. + Stage 1 (cooperative): when the per-turn budget elapses, sets + ``ctx.timeout_exceeded`` then ``cancel_event`` (ordering invariant) and + persists a durable ``timeout_cancelled_at`` marker. Stage 2 (hard cap): + after :func:`_resolve_hard_stop_grace` more time, if the handler is + still running it is FORCE-stopped via + :meth:`_hard_stop_timed_out_task`. - :param timeout_seconds: Total per-turn timeout budget (used as - the clock-skew clamp ceiling). - :type timeout_seconds: float + :param timeout_seconds: Total per-turn timeout budget (also the + clock-skew clamp ceiling). :param cancel_event: Event to set for cooperative cancel. - :type cancel_event: asyncio.Event :param ctx: TaskContext to set ``timeout_exceeded`` on BEFORE - ``cancel_event`` (ordering invariant). - :type ctx: TaskContext[Any] | None - :keyword remaining_seconds: Optional override for "time left in - this turn" — used on recovery to honor the persisted - turn-start timestamp. Clamped to - ``[0, timeout_seconds]`` for clock-skew safety. - When ``None``, the watchdog uses ``timeout_seconds`` directly - (fresh-entry / drain-re-entry case). - :paramtype remaining_seconds: float | None + ``cancel_event``. + :keyword task_id: The task id (enables marker persistence + hard cap). + When ``None`` the watchdog is cooperative-only (legacy behavior). + :keyword remaining_seconds: Optional "time left in this turn" override + for recovery; clamped to ``[0, timeout_seconds]``. """ if remaining_seconds is None: sleep_for = timeout_seconds @@ -1643,13 +1768,86 @@ async def _timeout_watchdog( if ctx is not None: ctx.timeout_exceeded = True cancel_event.set() + + if task_id is None: + logger.info( + "Timeout watchdog fired cooperative cancel (slept %.3fs of %.3fs " + "budget; cooperative-only)", + sleep_for, + timeout_seconds, + ) + return + + grace = _resolve_hard_stop_grace().total_seconds() logger.info( - "Timeout watchdog fired cooperative cancel (slept %.3fs of " - "%.3fs budget; cooperative-only — handler must check " - "ctx.cancel.is_set() and ctx.timeout_exceeded to wind down)", + "Timeout watchdog fired cooperative cancel for task %s (slept %.3fs of " + "%.3fs budget); will FORCE-stop after %.0fs hard-cap grace if the " + "handler does not wind down", + task_id, sleep_for, timeout_seconds, + grace, ) + # Stage 1b — persist the durable "this turn entered timeout + # cancellation" marker (the moment cancellation happened). + await self._persist_timeout_cancelled_marker(task_id) + # Stage 2 — grace, then hard stop. If the handler winds down first, the + # turn ends, _cancel_watchdog_for_turn cancels this sleep, and no hard + # stop happens. + await asyncio.sleep(grace) + await self._hard_stop_timed_out_task(task_id) + + async def _persist_timeout_cancelled_marker(self, task_id: str) -> None: + """Persist ``payload.timeout_cancelled_at`` = now (best-effort). + + Enables recovery to detect a turn that already entered timeout + cancellation and finalize it rather than re-running (nanny non-recovery). + + :param task_id: The task id. + """ + try: + await self._provider_update_locked( + task_id, + TaskPatchRequest(payload={_TIMEOUT_CANCELLED_AT_KEY: _utc_now_iso()}), + ) + except Exception: # pylint: disable=broad-except + logger.warning( + "failed to persist timeout-cancelled marker for task %s (recovery " + "backstop still applies)", + task_id, + exc_info=True, + ) + + async def _hard_stop_timed_out_task(self, task_id: str) -> None: + """Force-stop a handler that ignored its per-turn timeout (hard cap). + + Isolated: kills the child worker (``_WorkerKilled`` → the execute loop's + cancel finalization moves the record out of ``in_progress``). + Non-isolated fallback: cancels the in-process execution task. + + :param task_id: The task id to hard-stop. + """ + active = self._active_tasks.get(task_id) + if active is None: + return + run = getattr(active, "isolated_run", None) + if run is not None: + logger.warning( + "Timeout hard cap reached for task %s — killing isolated worker (pid %s). " + "This turn will be force-ended and NOT recovered.", + task_id, + getattr(run, "pid", "?"), + ) + run.kill() + return + if not active.execution_task.done(): + logger.warning( + "Timeout hard cap reached for task %s — cancelling in-process handler " + "(isolation off; best-effort).", + task_id, + ) + active.renewal_cancel.set() + active.execution_task.cancel() async def _execute_task( self, @@ -1745,6 +1943,7 @@ async def _spawn_watchdog_for_turn( timeout_seconds=timeout_seconds, cancel_event=ctx.cancel, ctx=ctx, + task_id=task_id, remaining_seconds=remaining, ) ) @@ -1818,6 +2017,357 @@ async def _compute_remaining_for_watchdog( ctx.cancel.set() return remaining + # ================================================================== + # Hard execution cap — process-isolated handler execution + # ================================================================== + + async def _run_handler( + self, + fn: Callable[..., Awaitable[Any]], + ctx: "TaskContext[Any]", + task_id: str, + opts: TaskOptions, + ) -> Any: + """Run the user handler, optionally in an isolated child process. + + When isolation is off (default) this is exactly ``await fn(ctx)``. + When on, the handler runs in a child process the timeout watchdog can + force-kill; a kill surfaces as ``asyncio.CancelledError`` so the + existing end-of-turn finalization (delete / suspend / drain) applies. + + :param fn: The task handler. + :param ctx: The task context (main-process copy). + :param task_id: The task id. + :param opts: The task options. + :return: The handler's return value (or suspend / exit sentinel). + """ + if not _isolation_enabled(): + return await fn(ctx) + + # §13.6 — reuse a per-chain persistent worker for MULTI-TURN chains + # when enabled: import once/chain instead of once/turn. One-shot tasks + # (single turn) gain nothing from reuse and keep the per-turn path. + if _worker_reuse_enabled() and getattr(opts, "_is_multi_turn", False): + return await self._run_handler_reused(fn, ctx, task_id, opts) + + from ._isolation import ( # pylint: disable=import-outside-toplevel + IsolatedRun, + WorkerCrash, + _WorkerKilled, + start_forked, + start_isolated, + ) + + # Fork backend (Unix) vs spawn backend — same IsolatedRun, different + # transport. Selected live per-turn. + _start = start_forked if _worker_fork_enabled() else start_isolated + + snapshot = self._build_worker_snapshot(fn, ctx, opts) + bridge = self._build_isolation_bridge(ctx, task_id) + + max_worker_attempts = 2 # Gap 4: bounded local re-invoke on crash. + for worker_attempt in range(max_worker_attempts): + run: IsolatedRun = await _start(snapshot, bridge) + active = self._active_tasks.get(task_id) + if active is not None: + active.isolated_run = run + # Bridge main-process cancel/shutdown → child. + cancel_bridge = asyncio.create_task(self._bridge_cancel_to_worker(ctx, run)) + try: + result = await run.outcome() + return result + except _WorkerKilled: + # Intentional timeout hard-cap kill (§5.2). Surface as + # CancelledError so the execute loop's existing end-of-turn + # finalization applies: one-shot -> delete; multi-turn -> + # transition + _try_drain_steering (queued input drains to the + # next turn, else the chain suspends). Verified by + # test_isolation_multiturn.py. + raise asyncio.CancelledError() + except WorkerCrash: + if worker_attempt + 1 >= max_worker_attempts: + logger.warning( + "isolated worker for task %s crashed %d times; failing the turn", + task_id, + worker_attempt + 1, + ) + raise + logger.warning( + "isolated worker for task %s crashed (attempt %d); re-launching", + task_id, + worker_attempt + 1, + ) + # Re-enter as a recovered turn on the retry. + snapshot["entry_mode"] = "recovered" + continue + finally: + cancel_bridge.cancel() + if active is not None: + active.isolated_run = None + + async def _run_handler_reused( + self, + fn: Callable[..., Awaitable[Any]], + ctx: "TaskContext[Any]", + task_id: str, + opts: TaskOptions, + ) -> Any: + """Run one turn on the chain's persistent worker (§13.6 reuse path). + + Mirrors the per-turn isolated path but reuses a warm worker keyed by + ``task_id`` across the chain's turns. A hard-cap kill or crash discards + the worker (so the next turn spawns a fresh one) and, for a kill, + surfaces as ``asyncio.CancelledError`` so the existing end-of-turn + finalization (drain / suspend) applies exactly as in the per-turn path. + + :param fn: The task handler. + :param ctx: The task context (main-process copy). + :param task_id: The task id (worker registry key). + :param opts: The task options. + :return: The handler's return value (or suspend / exit sentinel). + """ + from ._isolation import ( # pylint: disable=import-outside-toplevel + WorkerCrash, + _WorkerKilled, + ) + + snapshot = self._build_worker_snapshot(fn, ctx, opts) + bridge = self._build_isolation_bridge(ctx, task_id) + + max_worker_attempts = 2 # bounded local re-invoke on crash. + for worker_attempt in range(max_worker_attempts): + worker = await self._get_or_start_worker(task_id) + active = self._active_tasks.get(task_id) + if active is not None: + active.isolated_run = worker # hard-cap watchdog kills this + cancel_bridge = asyncio.create_task(self._bridge_cancel_to_worker(ctx, worker)) + try: + result = await worker.run_turn(snapshot, bridge) + return result + except _WorkerKilled: + # Timeout hard-cap kill — discard the (dead) worker so the next + # turn re-spawns, then surface as CancelledError for the + # existing drain/suspend finalization. + self._discard_worker(task_id) + raise asyncio.CancelledError() + except WorkerCrash: + self._discard_worker(task_id) + if worker_attempt + 1 >= max_worker_attempts: + logger.warning( + "reused worker for task %s crashed %d times; failing the turn", + task_id, + worker_attempt + 1, + ) + raise + logger.warning( + "reused worker for task %s crashed (attempt %d); re-launching", + task_id, + worker_attempt + 1, + ) + snapshot["entry_mode"] = "recovered" + continue + finally: + cancel_bridge.cancel() + if active is not None: + active.isolated_run = None + + async def _get_or_start_worker(self, task_id: str) -> Any: + """Return the chain's warm persistent worker, spawning one if needed. + + :param task_id: The worker registry key. + :return: A live ``PersistentWorker``. + """ + worker = self._reuse_workers.get(task_id) + if worker is not None and getattr(worker, "alive", False): + return worker + if worker is not None: # stale/dead entry + self._reuse_workers.pop(task_id, None) + + from ._isolation import ( # pylint: disable=import-outside-toplevel + start_forked_worker, + start_persistent_worker, + ) + + # Fork backend (Unix) vs spawn backend — same PersistentWorker, different + # transport. + worker = await (start_forked_worker() if _worker_fork_enabled() else start_persistent_worker()) + self._reuse_workers[task_id] = worker + logger.info( + "started persistent reuse worker (pid %s) for task %s", + getattr(worker, "pid", "?"), + task_id, + ) + self._ensure_worker_reaper() + return worker + + def _discard_worker(self, task_id: str) -> None: + """Kill (idempotent) and deregister the chain's persistent worker. + + :param task_id: The worker registry key. + """ + worker = self._reuse_workers.pop(task_id, None) + if worker is None: + return + try: + worker.kill() + except Exception: # pylint: disable=broad-except + logger.debug("error killing discarded worker for task %s", task_id, exc_info=True) + + def _ensure_worker_reaper(self) -> None: + """Start the idle-TTL reaper loop if not already running.""" + if self._worker_reaper_task is None or self._worker_reaper_task.done(): + self._worker_reaper_task = asyncio.create_task(self._worker_reaper_loop()) + + async def _worker_reaper_loop(self) -> None: + """Evict per-chain workers idle beyond the TTL (§13.6 warm-across-suspend). + + A worker parked between turns (chain suspended) is killed + deregistered + once idle for ``AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS`` so long or + indefinite parks release their ~100-300MB footprint. A worker actively + running a turn (``in_flight``) is never evicted. TTL <= 0 disables + eviction. Exits when shutdown is signalled or no workers remain. + """ + while not self._shutdown_event.is_set(): + ttl = _resolve_worker_idle_ttl() + poll = 30.0 if ttl <= 0 else max(5.0, min(ttl, 60.0)) + try: + await asyncio.wait_for(self._shutdown_event.wait(), timeout=poll) + return # shutdown + except asyncio.TimeoutError: + pass + except asyncio.CancelledError: + return + if not self._reuse_workers: + # Nothing to reap — let the loop exit; it is re-armed on the + # next _get_or_start_worker. + return + ttl = _resolve_worker_idle_ttl() + now = asyncio.get_running_loop().time() + self._reap_idle_workers(now, ttl) + + def _reap_idle_workers(self, now: float, ttl: float) -> int: + """Evict dead + idle-beyond-TTL workers. Returns the count evicted. + + Separated from the reaper loop so it can be driven deterministically in + tests. ``in_flight`` workers are never evicted; ``ttl <= 0`` disables + idle eviction (dead workers are still reaped). + + :param now: Current event-loop monotonic time. + :param ttl: Idle time-to-live in seconds. + :return: Number of workers evicted. + """ + evicted = 0 + for tid, worker in list(self._reuse_workers.items()): + try: + if not getattr(worker, "alive", True): + self._reuse_workers.pop(tid, None) + continue + if ttl <= 0 or getattr(worker, "in_flight", False): + continue + idle = worker.idle_seconds(now) + if idle >= ttl: + logger.info( + "idle-TTL evicting reuse worker (pid %s) for task %s " + "(idle %.0fs >= %.0fs)", + getattr(worker, "pid", "?"), + tid, + idle, + ttl, + ) + self._reuse_workers.pop(tid, None) + worker.kill() + evicted += 1 + except Exception: # pylint: disable=broad-except + logger.warning("worker reaper error for task %s", tid, exc_info=True) + return evicted + + def _build_worker_snapshot( + self, fn: Callable[..., Awaitable[Any]], ctx: "TaskContext[Any]", opts: TaskOptions + ) -> dict: + """Build the pickle-free ctx snapshot handed to the worker at fork.""" + return { + "handler_module": getattr(fn, "__module__", ""), + "handler_name": opts.name, + "task_id": ctx.task_id, + "input_id": ctx.input_id, + "session_id": getattr(ctx, "_session_id", "") or "", + "input": ctx.input, + "entry_mode": ctx.entry_mode, + "is_steered_turn": ctx.is_steered_turn, + "retry_attempt": ctx.retry_attempt, + "recovery_count": ctx.recovery_count, + # Default-namespace metadata seed (named namespaces are proxied via + # flush RPCs and ride the terminal snapshot). + "metadata": dict(ctx.metadata), + } + + def _build_isolation_bridge(self, ctx: "TaskContext[Any]", task_id: str) -> Any: + """Build the main-process callbacks the worker proxies back to.""" + from ._isolation import IsolationBridge # pylint: disable=import-outside-toplevel + from ..streaming import streams # pylint: disable=import-outside-toplevel + from ..streaming import EventStreamNotFoundError # pylint: disable=import-outside-toplevel + + flush_cb = self._make_metadata_flush(task_id) + + async def apply_flush(namespace: Optional[str], data: dict) -> None: + # Keep main's ctx.metadata in sync so the end-of-turn _flush_all is + # consistent, then persist to the store. + _apply_namespace(ctx.metadata, namespace, data) + await flush_cb(namespace, data) + + async def stream_emit(stream_id: str, payload: Any, close: bool) -> None: + stream = await streams.get_or_create(stream_id) + await stream.emit(payload, close=close) + + async def stream_close(stream_id: str) -> None: + try: + stream = await streams.get(stream_id) + except EventStreamNotFoundError: + return + await stream.close() + + async def stream_last_cursor(stream_id: str) -> Optional[int]: + try: + stream = await streams.get(stream_id) + except EventStreamNotFoundError: + return None + return await stream.last_cursor() + + def get_pending_count() -> int: + return self._make_pending_count_provider(task_id)() + + def apply_final_metadata(snapshot: dict) -> None: + for ns_key, data in snapshot.items(): + ns = None if ns_key == "__default__" else ns_key + _apply_namespace(ctx.metadata, ns, data) + + return IsolationBridge( + apply_flush=apply_flush, + stream_emit=stream_emit, + stream_close=stream_close, + stream_last_cursor=stream_last_cursor, + get_pending_count=get_pending_count, + apply_final_metadata=apply_final_metadata, + ) + + async def _bridge_cancel_to_worker(self, ctx: "TaskContext[Any]", run: Any) -> None: + """Forward main-process cancel/shutdown to the isolated worker.""" + async def fwd_cancel() -> None: + await ctx.cancel.wait() + await run.signal_cancel( + timeout_exceeded=bool(getattr(ctx, "timeout_exceeded", False)), + cancel_requested=bool(getattr(ctx, "cancel_requested", False)), + ) + + async def fwd_shutdown() -> None: + await ctx.shutdown.wait() + await run.signal_shutdown() + + try: + await asyncio.gather(fwd_cancel(), fwd_shutdown()) + except asyncio.CancelledError: + pass + async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-branches,too-many-nested-blocks,unused-argument,too-many-locals self, *, @@ -1870,7 +2420,7 @@ async def _execute_task_loop( # pylint: disable=too-many-statements,too-many-br ) ) try: - result = await fn(ctx) + result = await self._run_handler(fn, ctx, task_id, opts) finally: reset_request_context(request_context_token) @@ -2373,6 +2923,8 @@ async def _try_drain_steering( # pylint: disable=too-many-branches,too-many-sta # turn-start boundary — write a fresh _turn_started_at so the # respawned watchdog computes a full per-turn budget. payload[_TURN_STARTED_AT_KEY] = _utc_now_iso() + # Clear the hard-cap marker atomically with the new turn-start stamp. + payload[_TIMEOUT_CANCELLED_AT_KEY] = None payload["steering"] = steering # SOT §11/§20: the framework does not write payload["output"]; # no clear is needed at the drain transition. @@ -3041,6 +3593,12 @@ async def _recover_stale_tasks(self) -> None: # Look up stored opts for resumed-task configuration. fn_name = (task_info.source or {}).get("name", "") opts = self._resume_opts.get(fn_name) + # Hard execution cap (Spec 037 #8): a turn that already + # entered timeout cancellation must NOT be re-run. Finalize + # it terminally so the nanny stops recovering it. + if self._turn_timed_out(task_info, opts): + await self._finalize_timed_out_recovery(task_info, opts) + continue await self._start_existing_task( fn=fn, fn_name=task_info.agent_name, @@ -3056,6 +3614,68 @@ async def _recover_stale_tasks(self) -> None: exc_info=True, ) + def _turn_timed_out(self, task_info: TaskInfo, opts: "TaskOptions | None") -> bool: + """Detect whether the current turn already entered timeout cancellation. + + Primary signal: the persisted ``timeout_cancelled_at`` marker, correlated + with ``turn_started_at`` (marker >= turn-start, so a stale marker from a + prior turn is ignored). Backstop (marker missing due to a crash between + the cooperative cancel and the marker PATCH): the derived deadline + ``now >= turn_started_at + timeout``. + + :param task_info: The stale task record under recovery. + :param opts: The registered options (for the timeout budget). + :return: True iff the turn already timed out. + """ + payload = task_info.payload or {} + started = _parse_turn_started_at(payload.get(_TURN_STARTED_AT_KEY)) + marker = payload.get(_TIMEOUT_CANCELLED_AT_KEY) + if marker: + marker_ts = _parse_turn_started_at(marker) + if marker_ts is not None and (started is None or marker_ts >= started - 1.0): + return True + if started is not None: + import time # pylint: disable=import-outside-toplevel + + timeout_s = _resolve_effective_timeout(opts.timeout if opts else None).total_seconds() + if time.time() >= started + timeout_s: + return True + return False + + async def _finalize_timed_out_recovery(self, task_info: TaskInfo, opts: "TaskOptions | None") -> None: + """Finalize a timed-out turn on recovery instead of re-running it. + + One-shot / ephemeral → delete; multi-turn → suspend (both move the record + out of ``in_progress`` so the nanny stops recovering it). + + :param task_info: The stale task record. + :param opts: The registered options. + """ + task_id = task_info.id + is_multi_turn = bool(getattr(opts, "_is_multi_turn", False)) if opts else False + ephemeral = bool(getattr(opts, "ephemeral", True)) if opts else True + logger.warning( + "Recovery: task %s already exceeded its per-turn timeout hard cap; " + "finalizing (%s) without re-running the turn.", + task_id, + "suspend" if is_multi_turn else "delete", + ) + try: + if is_multi_turn: + await self._provider_update_locked( + task_id, + TaskPatchRequest( + status="suspended", + lease_owner=self._lease_owner, + lease_instance_id=self._instance_id, + lease_duration_seconds=0, + ), + ) + else: + await self._provider.delete(task_id, force=True) + except Exception: # pylint: disable=broad-exception-caught + logger.warning("Recovery finalize failed for timed-out task %s", task_id, exc_info=True) + def _find_resume_callback(self, task_info: TaskInfo) -> Callable[..., Any] | None: """Find a registered resume callback for a task. diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/task-timeout-isolation-design.md b/sdk/agentserver/azure-ai-agentserver-core/docs/task-timeout-isolation-design.md new file mode 100644 index 000000000000..be180c5fdc75 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/task-timeout-isolation-design.md @@ -0,0 +1,414 @@ +# Task timeout hard-cap via process isolation — unified design (spawn + fork) + +**Status:** spawn + per-chain reuse **implemented & tested** (unit + live westus2); fork backend +**implemented & tested** (unit on Linux/WSL + live westus2). Both are **opt-in, default off**. +This doc is the single source of truth: the **general design is common to both backends**; +the two options differ **only in the worker-internals** (§7). It supersedes/combines +`task-timeout-hardcap-design.md` and `fork-worker-design.md`. + +--- + +## 1. Problem & requirements +A durable/resilient task's per-turn `timeout` was historically **cooperative-only** — a handler +that ignored `ctx.cancel` could run forever (conformance C-TMO-6: "MUST NOT force-stop"). We +need a **hard execution cap**: when the timeout elapses and the handler ignores the cooperative +cancel, force-stop it after a bounded grace, **without**: + +- **R1 — Force-stop at will:** the SDK must be able to stop a runaway handler regardless of what + it's doing (CPU-bound loop, blocking syscall, C-extension call, or code that swallows + cancellation). +- **R2 — Don't disrupt the parent:** killing the handler must not affect the main container or + co-located tasks. +- **R3 — Durability preserved:** the task store, lease, recovery/nanny, streaming, and steering + must stay correct; a force-stopped turn must **not** be silently resurrected. + +**Non-goal:** capping concurrency. We do **not** introduce a policy concurrency cap (see §9). + +--- + +## 2. The only authoritative mechanism: a separate OS process you can kill +For **arbitrary native Python**, the *only* way to satisfy R1+R2 is to run the handler in a +**separate OS process and `kill()` it**. Everything else fails one requirement: + +| Option | Force-stop? | Parent-safe? | Verdict | +|---|---|---|---| +| `task.cancel()` / `asyncio.timeout()` / `wait_for()` | ❌ cooperative | ✅ | old behavior; stubborn handler ignores it | +| `PyThreadState_SetAsyncExc` / `sys.settrace` / `SIGALRM` | ⚠️ can't interrupt C/blocking; catchable | ❌ corrupts shared state | rejected (non-authoritative + unsafe) | +| Thread + kill | ⚠️ no clean thread-kill in CPython | ❌ shares GIL/memory | rejected | +| Subinterpreters (PEP 734) | ❌ can't force-stop busy interp | ❌ shares process | immature; doesn't solve killability | +| **Subprocess + SIGKILL (chosen)** | ✅ | ✅ | the mechanism | +| Instrumented runtime (WASM epoch/fuel) | ✅ | ✅ | N/A — not native Python | + +**`spawn` and `fork` are two ways to create that killable process** — same kill mechanism, +different creation internals. That difference is the entire subject of §7; everything else in +this doc is common. + +--- + +## 3. Architecture — parent owns authority, child is pure compute +**Parent (main process):** owns everything with side effects/authority — the **task store & +lease**, **timeout watchdog**, **recovery/nanny**, **streams registry**, **steering queue**, +the `_ActiveTask` registry and (reuse) the worker registry. It runs the retry loop, drain, +suspend/complete/delete transitions, and marker/recovery logic. + +**Child (worker process):** a pure compute unit that runs **only `fn(ctx)`**. It may compute +anything but may **not** own or directly touch any durable/shared resource. Every side effect +the handler triggers is **proxied to the parent** over IPC. The child has **no** direct +store/lease access (enforced — see §8). + +The developer contract is **`@task` / `@multi_turn_task` on `async def fn(ctx)` and nothing +else** — isolation/spawn/fork/reuse are invisible, toggled by platform env vars. + +--- + +## 4. IPC — the common wire protocol & contracts +The protocol is **identical for both backends** (only the transport bytes differ — §7). It is +**length-prefixed (4-byte big-endian) UTF-8 JSON**, with an id-correlated request/response +sub-protocol. + +### 4.1 Message families +**Parent → child (control):** +- `MSG_RUN {snapshot, reuse?}` — start a turn (the snapshot = ctx fields + handler module/name). +- `MSG_CANCEL {timeout_exceeded, cancel_requested, pending?}` — cooperative cancel. +- `MSG_SHUTDOWN` — process-scoped shutdown. +- `MSG_RESP {id, value}` — response to a child round-trip. + +**Child → parent:** +- `MSG_EMIT {stream_id, payload, close}` / `MSG_STREAM_CLOSE {stream_id}` — streaming. +- `MSG_REQ {id, kind, ...}` — round-trip: `flush` / `last_cursor` / `pending_count`. +- **Terminal (exactly one per turn):** `MSG_RESULT {value, metadata}` | + `MSG_SUSPEND {reason, output, metadata}` | `MSG_EXIT_FOR_RECOVERY {metadata}` | + `MSG_ERROR {exc_type, exc_msg, traceback, metadata}`. + +### 4.2 The `ctx` proxy (what the handler touches) +The child's `ctx` looks identical to in-process, but its IO-bound handles are proxies: +- `ctx.metadata` → `_ProxyMetadata` → `REQ_FLUSH` → parent persists; full snapshot also rides + the terminal message so un-flushed mutations aren't lost. +- streams / `ctx.emit` → `_ProxyStream` (via `streams._factory`) → `MSG_EMIT` → parent's real + streams. +- `ctx.cancel` / `ctx.shutdown` → `_HybridCancel` (a thread-settable bool + `asyncio.Event`, so + it works even for a CPU-bound handler that polls without awaiting). +- `ctx.pending_input_count` → provider backed by `proto.pending`. +- Pure-data fields (`task_id`, `input`, `entry_mode`, `retry_attempt`, `recovery_count`, …) are + plain values carried in the snapshot. + +### 4.3 Bidirectional `ctx` bridge (watchdog ↔ handler) +In-process, `ctx` is shared mutable state between the watchdog (writes `timeout_exceeded`, +`cancel`) and the handler (reads them). Across processes there is no shared memory, so the +"sharing" is **emulated by message passing**, in both directions, preserving the ordering +invariant: + +- **Parent → child:** `_bridge_cancel_to_worker` watches the parent `ctx.cancel`/`ctx.shutdown`; + on fire it sends `MSG_CANCEL`/`MSG_SHUTDOWN`. The child reader sets the **cause booleans + first**, then trips the cancel event (cause-before-cancel invariant, same as in-process). +- **Child → parent:** metadata flush (`REQ_FLUSH`), pending count (`REQ_PENDING_COUNT`), streams + (`MSG_EMIT`), and the terminal metadata snapshot. + +**Design rule:** any *new* field shared between the handler and the watchdog/framework must get +an explicit bridge message — it will not "just work" across the boundary. + +### 4.4 Terminal-outcome contract (how the parent learns the result) +The parent never infers the outcome from an exit code. The child sends a **structured terminal +message** (§4.1); the parent's `outcome()`/`run_turn` applies the final metadata snapshot then +`_reconstruct_outcome`: +`MSG_RESULT`→value, `MSG_SUSPEND`→`_Suspended`, `MSG_EXIT_FOR_RECOVERY`→`_ExitForRecovery`, +`MSG_ERROR`→re-raised reconstructed exception (type+message+traceback; `CancelledError` mapped). +The exit code / EOF is used **only** as a crash fallback (§6). + +--- + +## 5. Lifecycle — tiered timeout enforcement (common) +Per turn, the parent arms a **tiered watchdog** (a task independent of the handler): + +1. **Stage 1 — cooperative.** At the per-turn budget: set `ctx.timeout_exceeded = True` **then** + `ctx.cancel` (bridged to the child); persist a durable `timeout_cancelled_at` marker. +2. **Stage 2 — grace.** Wait `AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS` (default 1 hr). If + the handler winds down, the turn ends normally and no kill happens. +3. **Stage 3 — hard cap.** Still running → `run.kill()` (SIGKILL the child). The kill surfaces + (via socket EOF, no terminal message) as `_WorkerKilled` → `_run_handler` raises + `asyncio.CancelledError()` → the **existing** end-of-turn finalization runs: + - one-shot → **delete** (ephemeral); + - multi-turn → `_try_drain_steering`: if input is queued, **drain to the next turn**; else + **suspend** the chain. + +The **manager awaits the handler** in `_execute_task_loop` (`await self._run_handler(...)`). +Because the watchdog is a *separate* task and the kill closes the child → EOF, the await is +**guaranteed to unblock** (unlike in-process, where a runaway `await fn(ctx)` can hang forever). + +--- + +## 6. Outcome, crash, lease, recovery (common) +- **Normal terminal:** parent reconstructs the outcome (§4.4), persists status + (completed/suspended/failed), and **releases the lease**. The child never touches the + lease/store. +- **Intentional kill (hard cap):** no terminal + we killed it → `_WorkerKilled` → CancelledError + → finalization as in §5. +- **Unexpected death (OOM/segfault):** no terminal + we did *not* kill it → **`WorkerCrash`** → + bounded local re-invoke (`max_worker_attempts = 2`, re-entered as a recovered turn), then fail + the turn. +- **Recovery / nanny non-resurrection:** a force-stopped turn is detected via the persisted + `timeout_cancelled_at` marker (correlated to `turn_started_at`) **or** the derived backstop + `now ≥ turn_started_at + timeout`; recovery **finalizes** it instead of re-running, and moving + the record out of `in_progress` stops the external `StaleTaskRecoveryService` from reviving it. +- **Parent crash (orphan child):** durability is handled by the *normal* crash-recovery path — + lease renewal stops → lease expires → nanny reclaims the `in_progress` task → re-runs from the + last checkpoint in a fresh container/parent/child. The orphaned child **cannot corrupt the + store** (no direct access; its IPC writes hit a dead pipe and silently no-op) and receives a + **cooperative shutdown** via EOF on its control channel. See §10 for the orphan-cleanup gap. + +--- + +## 7. **Backend internals — the ONLY place spawn and fork diverge** +Everything in §1–§6 and §8–§12 is backend-agnostic. The two backends differ purely in **how the +worker process is created, how bytes are transported, and how the handler/ctx are obtained** — +then both plug into the **same** `IsolatedRun` (one-shot) / `PersistentWorker` (reuse) and the +same protocol. Selected once per turn by `_make_spawner()` / the fork flag. + +### 7.0 Shared worker classes (both backends run these) +- `_ChildProtocol` — child-side protocol endpoint (control-reader thread + writer). +- `IsolatedRun` — parent-side one-shot runner (`outcome()`, `kill()`, read-loop, dispatch). +- `PersistentWorker` — parent-side reusable runner (`run_turn()`, turn-loop, `kill()`). +- `IsolationBridge` — parent callbacks (flush/emit/last_cursor/pending/apply_final_metadata). +- `_build_child_context` / `_resolve_handler` / `_classify_terminal` / `_reconstruct_outcome`. + +### 7.1 SPAWN internals (`create_subprocess_exec`) — the shipped default +- **Create:** `asyncio.create_subprocess_exec(python, "-m", "...tasks._isolation")` → a **blank + interpreter**. +- **Get the app:** the child **re-imports** the handler's module (~1–2 s) via `_resolve_handler` + (`importlib.import_module`), then **resolves the handler by name** from + `_REGISTERED_DESCRIPTORS`. +- **Get ctx:** rebuilt from the JSON **snapshot** (nothing crosses a process boundary without + serialization). +- **Transport:** three std streams — `stdin` (control), `stdout` (protocol), `stderr` (user + logs, forwarded to the parent logger). The child dups fd1→fd2 first so user `print()` can't + corrupt the protocol; the true fd1 is a private protocol fd. +- **Process handle:** `asyncio.subprocess.Process` → `.pid`, `.kill()` (SIGKILL), `.wait()`. +- **Memory:** full private copy per worker (**~100–300 MB**, no COW). +- **OS:** portable — Linux, macOS, Windows. +- **Manager in child:** never set (`_manager is None`) → Tasks API unavailable "for free" (§8). + +### 7.2 FORK internals (`multiprocessing.get_context("fork")`) — Unix-only, opt-in +- **Create:** `mp.get_context("fork").Process(target=_fork_child_entry, args=(child_sock, + parent_fd, reuse))`; `proc.start()` calls `os.fork()`. The child **inherits the parent's whole + imported app** (COW). +- **Get the app:** **no re-import** — `_resolve_handler`'s `import_module` is a **free no-op hit + on the inherited `sys.modules`**; the handler object is already resident. Start cost **~1–5 + ms**. +- **Get ctx:** rebuilt from the snapshot via the **same** `_build_child_context` — chosen over + literally inheriting the live `ctx` because (a) multi-turn reuse can't inherit a fresh ctx per + turn, (b) lowest risk, (c) the perf/memory wins come from inheriting the *interpreter*, not the + ctx object. (Developer DX is identical either way.) +- **Sanitization — `_after_fork_child` (fork-only, mandatory):** the child inherited a live + async+threaded parent, so **first thing**: `set_task_manager(None)` (neutralize the inherited + manager → no split-brain, §8), **close the inherited peer-socket fd** (and other dead fds), and + let `asyncio.run()` create a **fresh event loop** (the inherited loop is never used). CPython's + `os.register_at_fork` resets GIL/import/logging locks. +- **Transport:** one `socket.socketpair()` — a single bidirectional channel carrying the **same** + length-prefixed-JSON protocol (`_ChildProtocol(in_fd == out_fd == socketfd)`). User logs go to + the inherited `stderr` directly. +- **Process handle:** a thin **`_ForkProcAdapter`** presents the `multiprocessing.Process` + the + asyncio-wrapped socket with the **exact surface** `IsolatedRun`/`PersistentWorker` consume + (`stdin` writer, `stdout` reader, `stderr=None`, `pid`, `kill()`, `wait()`, `returncode`) — so + those classes run over fork **unchanged**. `start_forked` / `start_forked_worker` are the fork + analogues of `start_isolated` / `start_persistent_worker`. +- **Memory:** COW — each worker is a small incremental cost, not a full copy. +- **OS:** **Linux only** (`_worker_fork_enabled()` returns False elsewhere → spawn). macOS fork + + native libs is unsafe; Windows has no fork. +- **Safe alternative (`forkserver`):** fork children from a dedicated single-threaded preloaded + server (no held locks, no live loop) — safest, but must snapshot+resolve (no live ctx). Kept as + a documented sub-mode; not the default. Neither fork nor forkserver adds a concurrency cap + (they're factories, not pools). + +### 7.3 Side-by-side +| Property | Spawn | Fork | +|---|---|---| +| Create | `create_subprocess_exec` (blank interp) | `mp fork` (inherit app, COW) | +| Re-import app | **yes (~1–2 s/turn)** | **no** (inherited) | +| Resolve handler | by name (after re-import) | by name = free (already imported) | +| Build ctx | from snapshot | from snapshot | +| Transport | 3 std pipes (stdin/stdout/stderr) | 1 `socketpair` | +| Process handle | `asyncio.subprocess.Process` | `mp.Process` via `_ForkProcAdapter` | +| Start latency | ~1–2 s | **~1–5 ms** | +| Memory/worker | full copy ~100–300 MB | **COW** (small increment) | +| OS | Linux/macOS/Windows | **Linux only** | +| Manager in child | `None` (never set) | **`None` (actively nulled)** | +| Extra safety step | none | `_after_fork_child` sanitization | +| Kill primitive | `proc.kill()` | `proc.kill()` (mp) via adapter | + +--- + +## 8. Enforcement boundary — no Tasks API inside a handler (common) +A handler's complete interface is `ctx` (metadata, streams, cancel/shutdown/timeout, `return` +to suspend, pending count). The **Task manager is framework infrastructure**; a handler calling +`get_task_manager()` / `other_task.start()` / `manager.provider.get()` is reaching below its +layer. This is: +- **Unnecessary** — everything a handler needs is in `ctx`; orchestration belongs in the + request/app layer (main process), where the manager is always present. +- **Recovery-unsafe** — such calls aren't checkpointed, so a crash+recovery **re-fires** them + (e.g. duplicate `start`). In-handler orchestration silently violates durability. + +Isolation **enforces** this boundary by construction: the child has no manager (spawn: never +set; fork: actively nulled in `_after_fork_child`) → in-handler Tasks-API calls raise +`TaskManagerNotInitialized`. For fork this also prevents **split-brain** (using the inherited +manager copy would make two processes write the same lease/record). + +**Only the narrow `ctx` surface is bridged** (metadata/streams/cancel/pending). We deliberately +do **not** bridge the general Tasks API — see §11 (optional/deferred, with the recovery-safety +caveat). + +--- + +## 9. Concurrency, performance, memory (common framing, per-backend numbers) +- **No policy cap.** Neither backend imposes a concurrency ceiling; the only limit is physical + memory / OS fd+PID ulimits. We explicitly reject a bounded shared pool (it would cap + concurrency, which is a non-goal). +- **Isolation is a parallelism win:** handlers run in separate processes → no shared-GIL + starvation from a CPU-bound handler. +- **Spawn cost:** ~1–2 s re-import + ~100–300 MB per concurrent turn ⇒ effective ceiling ~**8–10 + on a 2 Gi** container (memory-bound; overload risks OOM). Per-chain reuse (below) amortizes the + re-import to once/chain. +- **Fork cost:** ~ms start + COW memory ⇒ ceiling rises **back toward in-process density**. +- **Per-chain reuse + idle-TTL** (multi-turn only): a warm worker per `task_id`, reused across + turns, evicted after `AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS` idle. **This is a *spawn* + optimization** (import once/chain instead of once/turn); **fork barely needs it** (per-turn + fork is already ~ms + COW), so the recommended fork mode is per-turn fork with reuse off. Reuse + composes with fork if cross-turn warm RAM is explicitly wanted. + +--- + +## 10. Open items / known gaps (same for both backends) +1. **Orphan-child hard guarantee (unimplemented).** On parent crash the child gets a cooperative + shutdown (EOF → `ctx.shutdown`) and can't corrupt the store, but a runaway child isn't + *guaranteed* to die (no `PR_SET_PDEATHSIG`). In the hosted container model, parent crash = + container/PID-namespace teardown → children die anyway; the gap bites only if the parent dies + but the namespace survives. **Hardening:** set `prctl(PR_SET_PDEATHSIG, SIGKILL)` on Linux + (fork: in `_after_fork_child`; spawn: in the child at startup) so the kernel kills the child + on parent death. +2. **Tasks-API bridge (spec only; optional/deferred).** In-handler Tasks-API calls raise today + (by design, §8). If genuine in-handler orchestration is ever needed it requires a *replay-safe + orchestrator* feature, not just an RPC bridge (the bridge would make the call *work* but not + *recovery-safe*). See §11. +3. **Clearer error message.** Replace `TaskManagerNotInitialized` (when hit from inside a + handler) with a purpose-built "the Tasks API can't be called from a handler; use `ctx` / + orchestrate from the app layer" error. Cheap, worth doing. +4. **At-least-once external side effects.** Inherent to crash-recovery (not isolation-specific): + a recovered turn re-runs from checkpoint; handler idempotency for external effects is the + author's responsibility. The store is always safe. + +--- + +## 11. Tasks-API bridge — optional/deferred (summary) +*Backend-agnostic.* If we ever support in-handler orchestration safely, route the Tasks-API +surface (`start`/`get`/`list`/`cancel`; then `run`/`result`/streaming) to the parent's **real** +manager over a new `MSG_TASKAPI_REQ/RESP` family, via a child-side `_ProxyTaskManager` (returned +by `get_task_manager()`) + `_ProxyTaskRun`, with all execution (and thus all store/lease writes) +in the parent (split-brain-free). Requires refactoring the decorator's `start`/`run` to delegate +to a manager entry point the proxy can intercept, plus error-type reconstruction across the +boundary. **Caveat:** the bridge is transport only — it does not make in-handler orchestration +recovery-safe; that needs a purpose-built deterministic/replay-safe orchestrator. Ship only if a +real need arises. + +--- + +## 12. Selection, flags, backward-compat, testing +### 12.1 Flags (all opt-in; default off ⇒ exact legacy in-process behavior) +- `AGENTSERVER_TASK_ISOLATION` — turn isolation on. +- `AGENTSERVER_TASK_WORKER_FORK` — use fork backend (Linux only; else spawn). +- `AGENTSERVER_TASK_WORKER_REUSE` — per-chain reuse (multi-turn only). +- `AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS` (default 1 hr). +- `AGENTSERVER_TASK_WORKER_IDLE_TTL_SECONDS` (default 300). + +Selection is a **single OS-branch** at worker creation (`_run_handler` / `_get_or_start_worker` +pick `start_forked*` vs `start_isolated`/`start_persistent_worker`). Off ⇒ `_run_handler` +short-circuits to `await fn(ctx)` (zero overhead). Runtime **fallback to spawn** if a fork +hand-off fails — never fail a turn over the optimization. + +### 12.2 Backward-compat +- **Default off = no change**; the legacy cooperative-only, in-process path is untouched. +- **Persisted state is two-way compatible:** the only new field (`timeout_cancelled_at`) is + additive/optional (recovery tolerates its absence; old readers ignore it). The **execution + backend is never persisted** → a task can be checkpointed in-process and recovered under + spawn/fork (or vice versa) freely. +- **Crash-restart is the safest upgrade point:** a crash already wiped RAM, so switching backends + there loses nothing that wasn't already gone. +- **Behavior changes are gated behind the flags:** hard-cap enforcement (was cooperative-only) + and per-turn RAM reset (spawn/fork) / cross-turn RAM persistence (reuse) — transparent to + handlers that keep durable state in `ctx.metadata`. + +### 12.3 Validation status +- **Unit:** full tasks suite green with isolation off (654 passed); isolation suite (spawn + one-shot/multi-turn/reuse) green; **fork** suite (one-shot + multi-turn drain/suspend + reuse + same-pid / hard-kill respawn / idle-TTL) green on Linux/WSL. Windows: fork tests skip + (Linux-gated), spawn tests pass — no regression. +- **Live westus2 (hosted Foundry agent):** spawn validated earlier (one-shot hard-kill+delete, + not recovered; multi-turn drain). Fork validated this session — image built (with an + `httpx`-in-requirements fix for image dependency drift), agent version deployed with + `AGENTSERVER_TASK_WORKER_FORK=1`, one-shot + multi-turn runaway dispatched, and the parent + container **survived the fork-child hard-kills and kept serving** (the core R2 signal). + +--- + +## 13. Alternatives considered & rejected +This consolidates every significant option evaluated across the design, grouped by decision. + +### 13.1 How to force-stop a runaway handler (the core mechanism) +| Alternative | Why rejected | +|---|---| +| **Cooperative only** — `task.cancel()`, `asyncio.timeout()`, `wait_for()`, `CancellationToken`-style flag | Not authoritative — a handler that never awaits or swallows `CancelledError` ignores it. This *was* the old behavior and is exactly what we're fixing. | +| **In-process exception injection** — `PyThreadState_SetAsyncExc`, `sys.settrace` deadline hook, `signal.SIGALRM` | Non-authoritative **and** unsafe: cannot interrupt a C-extension call or blocking syscall, is catchable, and injecting at an arbitrary point corrupts shared in-process state (same reason .NET removed `Thread.Abort`, Java deprecated `Thread.stop`). | +| **Thread + kill** | CPython has no clean thread termination; a killed thread shares the GIL/heap so it can't be isolated or safely stopped. | +| **Subinterpreters** (PEP 734) | Can't force-stop a busy interpreter, and it still shares the process → violates "don't disrupt the parent." Immature. | +| **cgroups / one container per task** | Works, but it's the same "kill a separate process" principle at a much heavier orchestration layer — more infra, slower, out of the SDK's control. | +| **Instrumented runtime** (WASM epoch/fuel, e.g. Wasmtime) | Genuinely preempts guest code without killing the host, but requires running handlers as **WASM, not native Python** — inapplicable. | +| **✅ Subprocess + SIGKILL (chosen)** | The only mechanism that is both authoritative (kills regardless of what the handler is doing) and parent-safe, for arbitrary native Python. | + +### 13.2 Enforcing the timeout — why not `asyncio.timeout` / `wait_for` +Rejected as the *enforcement* primitive because they are (a) equally cooperative (defeated by the same handlers), (b) `wait_for` **awaits the task's actual completion after firing**, so an uncooperative task makes it hang — it can't even hand control back to escalate, and (c) an `asyncio.timeout` context manager would cancel *our own* watchdog coroutine, which must survive to run the grace→SIGKILL. Chosen instead: `Event` + cause-bool for Stage 1, and a plain `asyncio.sleep(grace)` → out-of-process SIGKILL for Stage 3. + +### 13.3 Worker process creation +| Alternative | Decision | +|---|---| +| **Spawn** (`create_subprocess_exec`) | **Chosen default** — portable (Linux/macOS/Windows), single code path, zero overhead when off. Cost: per-turn re-import + full-copy memory. | +| **Fork** (`multiprocessing fork`) | **Chosen, Linux-only, opt-in** — inherits the imported app (COW, ~ms start, no re-import). Rejected as *default* because it's Unix-only and forking a live async+threaded parent needs sanitization. | +| **forkserver** | **Documented sub-mode, not default** — safest fork variant (forks from a clean preloaded server, no held locks), but it pickles args and lacks the live `ctx`, so it reverts to snapshot+resolve. Kept as the fallback if a native lib proves fork-hostile. | +| **Direct fork inheriting the live `ctx` object** | **Rejected in favor of snapshot-built `ctx`** — multi-turn reuse can't inherit a fresh ctx per turn anyway; inheriting a live object graph adds shared-fd/mutation risk; the perf/memory wins come from inheriting the *interpreter*, not the ctx object; developer DX is identical. | + +### 13.4 Worker lifecycle / concurrency +| Alternative | Decision | +|---|---| +| **Per-turn spawn/fork** | **Chosen default.** Spawn amortized via optional reuse; fork is already ~ms so per-turn is optimal. | +| **Per-chain reuse + idle-TTL** | **Optional, multi-turn only.** A spawn optimization (import once/chain); fork barely needs it. Retained for spawn latency and (optionally) cross-turn warm RAM. | +| **Bounded shared worker pool (size P)** | **Rejected** — it imposes a concurrency **cap** of P, and capping concurrency is an explicit non-goal. | +| **Unbounded warm pool** | **Rejected** — self-defeating: keeps the import saving but re-introduces the exact memory/OOM ceiling the pool was meant to remove (and grows idle memory without eviction). If a fixed P is undesirable, use a memory-derived cap + idle-TTL, never a hard pool. | + +### 13.5 Handler ↔ Task manager +| Alternative | Decision | +|---|---| +| **Let the (fork) child use the inherited manager** | **Rejected** — the inherited manager is a frozen COW copy holding the parent's store/lease fds → **split-brain** (two processes writing the same record). The child's manager is actively nulled. | +| **Bridge the Tasks API to the parent over RPC** | **Deferred / not built** — in-handler Tasks-API use is unnecessary (everything is in `ctx`; orchestration belongs in the app layer) **and recovery-unsafe** (uncheckpointed calls double-fire on recovery). We **enforce** the boundary (child has no manager) rather than bridge it. If real in-handler orchestration is ever needed it requires a replay-safe orchestrator feature, not just this bridge. | + +### 13.6 Recovery detection of a force-stopped turn +| Alternative | Decision | +|---|---| +| **Persisted "timed-out" terminal state** | Not used — avoids a bespoke terminal state and extra transitions. | +| **✅ Durable `timeout_cancelled_at` marker + derived backstop** (`now ≥ turn_started_at + timeout`) | **Chosen** — recovery finalizes the turn instead of re-running; moving the record out of `in_progress` stops the external nanny from resurrecting it. | + +### 13.7 Orphan child on parent crash +| Alternative | Decision | +|---|---| +| **Cooperative shutdown via control-channel EOF (current)** | Implemented, but not a hard guarantee — a runaway orphan can linger unless the container/namespace is torn down. | +| **`PR_SET_PDEATHSIG` → SIGKILL (recommended, not yet implemented)** | The robust fix — the kernel kills the child the instant the parent dies. Flagged as the hardening item (§10). | + +--- + +## Appendix — decision records +- **Only authoritative stop = separate killable process** (§2); in-process injectors are + non-authoritative + unsafe; instrumented runtimes (WASM) are inapplicable to native Python. +- **Concurrency capping is not a goal** (§9); no bounded pool; unbounded per-turn/reuse; the only + limit is physical memory. +- **Spawn is the portable default; fork is a Linux-only perf/memory optimization** behind a flag, + reusing all shared worker code via `_ForkProcAdapter`. +- **Handler must not call the Tasks API** (§8) — unnecessary + recovery-unsafe; enforced by + construction, not bridged (bridge is optional/deferred, §11). +- **`ctx` is shared state bridged by message-passing** (§4.3); new shared fields need explicit + bridging. diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/_isolation_handlers.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/_isolation_handlers.py new file mode 100644 index 000000000000..9d3b7a886d6d --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/_isolation_handlers.py @@ -0,0 +1,102 @@ +"""Importable task handlers for the isolation integration test. + +The isolation worker (a child process) resolves the handler by importing its +``__module__`` and looking it up in ``_REGISTERED_DESCRIPTORS`` by name, so the +handlers under test must live in an importable module (not inline in a test). +""" +from __future__ import annotations + +import time +from datetime import timedelta + +from azure.ai.agentserver.core.tasks import TaskContext, task + + +@task(name="iso_echo", timeout=timedelta(seconds=30)) +async def iso_echo(ctx: TaskContext[dict]) -> dict: + ctx.metadata["seen"] = ctx.input.get("msg") + await ctx.metadata.flush() + return {"echoed": ctx.input, "task_id": ctx.task_id} + + +@task(name="iso_never", timeout=timedelta(seconds=1)) +async def iso_never(ctx: TaskContext[dict]) -> dict: + # Ignores cooperative cancel entirely -> must be hard-killed by the cap. + while True: + time.sleep(0.05) + + +@task(name="iso_coop", timeout=timedelta(seconds=3)) +async def iso_coop(ctx: TaskContext[dict]) -> dict: + # Cooperates: polls ctx.cancel and winds down promptly on timeout, so it + # must NEVER reach the hard-cap kill. Timeout is set comfortably above the + # worker spawn+import cost (~1-2s) so the wind-down window is deterministic. + for _ in range(2000): # generous ceiling + if ctx.cancel.is_set(): + return {"wound_down": True, "timeout_exceeded": ctx.timeout_exceeded} + time.sleep(0.05) + return {"wound_down": False} + + +from azure.ai.agentserver.core.tasks import multi_turn_task # noqa: E402 + + +@multi_turn_task(name="iso_mt", steerable=True, timeout=timedelta(seconds=2)) +async def iso_mt(ctx: "TaskContext[dict]") -> None: + """Multi-turn probe: records its input tag, then either runs away + (ignores cancel — must be hard-killed) or cooperates (winds down).""" + tag = ctx.input.get("tag") + mode = ctx.input.get("mode") + ctx.metadata["last_tag"] = tag + seen = list(ctx.metadata.get("seen_tags", [])) + seen.append(tag) + ctx.metadata["seen_tags"] = seen + await ctx.metadata.flush() + if mode == "runaway": + while True: + time.sleep(0.05) # ignore cancel entirely -> hard-killed + # cooperative: wind down promptly if cancelled, else finish quickly + for _ in range(50): + if ctx.cancel.is_set(): + return None + time.sleep(0.02) + return None + + +import os # noqa: E402 + + +@multi_turn_task(name="iso_reuse", steerable=True, timeout=timedelta(seconds=30)) +async def iso_reuse(ctx: "TaskContext[dict]") -> None: + """Reuse probe: records the child PID each turn so tests can assert the + SAME worker process ran consecutive turns (per-chain reuse).""" + pids = list(ctx.metadata.get("pids", [])) + pids.append(os.getpid()) + ctx.metadata["pids"] = pids + ctx.metadata["last_tag"] = ctx.input.get("tag") + await ctx.metadata.flush() + return None # implicit suspend — chain stays alive for the next turn + + +@multi_turn_task(name="iso_reuse_mt", steerable=True, timeout=timedelta(seconds=2)) +async def iso_reuse_mt(ctx: "TaskContext[dict]") -> None: + """Reuse + hard-cap probe: records child PID and seen tags; ``runaway`` + mode ignores cancel (must be hard-killed -> worker discarded -> next turn + runs in a NEW pid).""" + pids = list(ctx.metadata.get("pids", [])) + pids.append(os.getpid()) + ctx.metadata["pids"] = pids + tag = ctx.input.get("tag") + ctx.metadata["last_tag"] = tag + seen = list(ctx.metadata.get("seen_tags", [])) + seen.append(tag) + ctx.metadata["seen_tags"] = seen + await ctx.metadata.flush() + if ctx.input.get("mode") == "runaway": + while True: + time.sleep(0.05) # ignore cancel -> hard-killed + for _ in range(50): + if ctx.cancel.is_set(): + return None + time.sleep(0.02) + return None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation.py new file mode 100644 index 000000000000..03d7d8e7abe6 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation.py @@ -0,0 +1,82 @@ +"""Integration tests for process-isolated handler execution + the timeout hard cap. + +These run the real ``TaskManager`` (local file provider) with +``AGENTSERVER_TASK_ISOLATION=1`` so the handler executes in a child process. +Windows/macOS/Linux compatible (``asyncio.create_subprocess_exec``). +""" +from __future__ import annotations + +import asyncio +import os +import sys +import uuid +from pathlib import Path + +import pytest + +_HANDLER_DIR = str(Path(__file__).parent) + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +@pytest.fixture(autouse=True) +def _isolation_env(monkeypatch): + # Enable isolation + a short hard-cap grace so the cap fires quickly. + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + # The child subprocess must be able to import the handler module + the SDK. + existing = os.environ.get("PYTHONPATH", "") + core_root = str(Path(__file__).resolve().parents[2]) # azure-ai-agentserver-core + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, core_root, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +@pytest.mark.asyncio +async def test_isolated_echo_returns_result(tmp_path): + import _isolation_handlers as h # noqa: F401 (registers the tasks) + manager, mgr_mod = await _setup(tmp_path) + try: + result = await h.iso_echo.run(input={"msg": "hi"}) + assert result["echoed"] == {"msg": "hi"}, result + assert "task_id" in result, result + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_isolated_hard_cap_kills_runaway_oneshot(tmp_path): + from azure.ai.agentserver.core.tasks import TaskCancelled + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + run = await h.iso_never.start(input={"n": 1}) + task_id = run.task_id + # timeout=1s + grace=1s -> hard-killed within a few seconds. + with pytest.raises(TaskCancelled): + await asyncio.wait_for(run.result(), timeout=15) + # One-shot ephemeral: the record is deleted on the cancel finalization. + await asyncio.sleep(0.5) + info = await manager._provider.get(task_id) + assert info is None or info.status in ("completed", "failed", "suspended"), ( + f"runaway one-shot should be gone/terminal, got {getattr(info, 'status', None)}" + ) + finally: + await manager.shutdown() + mgr_mod._manager = None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_multiturn.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_multiturn.py new file mode 100644 index 000000000000..a5c999696dca --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_multiturn.py @@ -0,0 +1,183 @@ +"""MULTI-TURN + per-chain REUSE under the FORK isolation backend. + +Mirrors test_isolation_multiturn.py and test_isolation_reuse.py but forces the +fork worker backend (AGENTSERVER_TASK_WORKER_FORK=1). Linux-only. +""" +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="fork worker backend is Linux-only" +) + +_HANDLER_DIR = str(Path(__file__).parent) +_CORE_ROOT = str(Path(__file__).resolve().parents[2]) + + +@pytest.fixture(autouse=True) +def _iso_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_WORKER_FORK", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, _CORE_ROOT, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +def _meta(info): + return (info.payload or {}).get("metadata", {}) if info else {} + + +async def _wait_for(manager, tid, predicate, timeout=12.0): + deadline = asyncio.get_event_loop().time() + timeout + last = None + while asyncio.get_event_loop().time() < deadline: + info = await manager._provider.get(tid) + last = info + if predicate(info): + return info + await asyncio.sleep(0.1) + return last + + +# --------------------------- multi-turn (no reuse) --------------------------- + +@pytest.mark.asyncio +async def test_fork_multiturn_hardcap_drains_to_queued_next_turn(tmp_path): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "fork-chain-drain-1" + await h.iso_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + await asyncio.sleep(1.0) + await h.iso_mt.start(task_id=tid, input={"mode": "coop", "tag": "B"}) + await asyncio.sleep(9) + info = await manager._provider.get(tid) + seen = _meta(info).get("seen_tags", []) + assert "A" in seen, f"turn 1 (A) should have run; seen={seen}" + assert "B" in seen, f"turn 2 (B) should have run after drain; seen={seen}" + assert _meta(info).get("last_tag") == "B", f"last turn should be B; meta={_meta(info)}" + assert info is not None and info.status in ("suspended", "in_progress"), info.status + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_fork_multiturn_hardcap_no_queue_suspends(tmp_path): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "fork-chain-suspend-1" + run = await h.iso_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + from azure.ai.agentserver.core.tasks import TaskCancelled + with pytest.raises(TaskCancelled): + await asyncio.wait_for(run.result(), timeout=12) + await asyncio.sleep(0.5) + info = await manager._provider.get(tid) + assert info is not None, "multi-turn chain should NOT be deleted" + assert info.status == "suspended", f"expected suspended, got {info.status}" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +# ------------------------------- reuse (fork) -------------------------------- + +@pytest.fixture +def _reuse_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_WORKER_REUSE", "1") + yield + + +@pytest.mark.asyncio +async def test_fork_reuse_same_pid_across_turns(tmp_path, _reuse_env): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "fork-reuse-same-1" + await h.iso_reuse.start(task_id=tid, input={"tag": "A"}) + await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "A") + assert tid in manager._reuse_workers + pid_after_a = manager._reuse_workers[tid].pid + await h.iso_reuse.start(task_id=tid, input={"tag": "B"}) + info = await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "B") + pids = _meta(info).get("pids", []) + assert len(pids) == 2, f"expected 2 turns recorded; pids={pids}" + assert pids[0] == pids[1], f"turns should reuse ONE forked worker; pids={pids}" + assert pids[1] == pid_after_a + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_fork_reuse_hardkill_respawns_new_pid(tmp_path, _reuse_env): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "fork-reuse-kill-1" + await h.iso_reuse_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + await _wait_for(manager, tid, lambda i: "A" in _meta(i).get("seen_tags", [])) + pid_a = manager._reuse_workers[tid].pid + await h.iso_reuse_mt.start(task_id=tid, input={"mode": "coop", "tag": "B"}) + info = await _wait_for( + manager, tid, lambda i: "B" in _meta(i).get("seen_tags", []), timeout=15 + ) + seen = _meta(info).get("seen_tags", []) + pids = _meta(info).get("pids", []) + assert "A" in seen and "B" in seen, f"both turns should run; seen={seen}" + assert len(pids) == 2 and pids[0] != pids[1], f"hard-kill must respawn; pids={pids}" + assert pids[0] == pid_a + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_fork_reuse_idle_ttl_evicts_worker(tmp_path, _reuse_env): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "fork-reuse-ttl-1" + await h.iso_reuse.start(task_id=tid, input={"tag": "A"}) + await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "A") + worker = manager._reuse_workers.get(tid) + assert worker is not None and worker.alive + pid_a = worker.pid + now = asyncio.get_event_loop().time() + worker.last_active_monotonic = now - 3600 + evicted = manager._reap_idle_workers(now, ttl=1.0) + assert evicted == 1 + assert tid not in manager._reuse_workers + await h.iso_reuse.start(task_id=tid, input={"tag": "B"}) + info = await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "B") + pids = _meta(info).get("pids", []) + assert len(pids) == 2 and pids[1] != pid_a, f"post-eviction must use NEW worker; pids={pids}" + finally: + await manager.shutdown() + mgr_mod._manager = None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_oneshot.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_oneshot.py new file mode 100644 index 000000000000..b9463d881ca0 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_fork_oneshot.py @@ -0,0 +1,109 @@ +"""ONE-SHOT task timeout cancellation under the FORK isolation backend. + +Mirrors test_isolation_oneshot.py but forces the fork worker backend +(AGENTSERVER_TASK_WORKER_FORK=1). Fork is Unix-only, so the whole module is +skipped on non-Linux platforms. +""" +from __future__ import annotations + +import asyncio +import os +import sys +import time +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="fork worker backend is Linux-only" +) + +_HANDLER_DIR = str(Path(__file__).parent) +_CORE_ROOT = str(Path(__file__).resolve().parents[2]) + +_TURN_STARTED_AT_KEY = "turn_started_at" +_TIMEOUT_CANCELLED_AT_KEY = "timeout_cancelled_at" + + +@pytest.fixture(autouse=True) +def _iso_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_WORKER_FORK", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, _CORE_ROOT, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +@pytest.mark.asyncio +async def test_fork_oneshot_cooperative_winds_down_not_killed(tmp_path): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + result = await asyncio.wait_for(h.iso_coop.run(input={"n": 1}), timeout=20) + assert result == {"wound_down": True, "timeout_exceeded": True}, result + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_fork_oneshot_runaway_hard_killed_and_deleted(tmp_path): + from azure.ai.agentserver.core.tasks import TaskCancelled + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + run = await h.iso_never.start(input={"n": 1}) + task_id = run.task_id + t0 = time.monotonic() + with pytest.raises(TaskCancelled): + await asyncio.wait_for(run.result(), timeout=15) + elapsed = time.monotonic() - t0 + assert elapsed >= 1.5, f"killed too early ({elapsed:.2f}s) — grace not respected" + assert elapsed < 8.0, f"killed too late ({elapsed:.2f}s)" + await asyncio.sleep(0.4) + info = await manager._provider.get(task_id) + assert info is None, f"runaway one-shot should be deleted, got {getattr(info,'status',None)}" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_fork_oneshot_marker_persisted_during_grace(tmp_path): + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + run = await h.iso_never.start(input={"n": 1}) + task_id = run.task_id + await asyncio.sleep(1.4) + info = await manager._provider.get(task_id) + assert info is not None, "record should still exist mid-grace" + assert info.status == "in_progress", f"expected in_progress mid-grace, got {info.status}" + marker = (info.payload or {}).get(_TIMEOUT_CANCELLED_AT_KEY) + assert marker, f"timeout_cancelled_at marker should be persisted, payload={info.payload}" + with pytest.raises(Exception): + await asyncio.wait_for(run.result(), timeout=8) + await asyncio.sleep(0.4) + assert await manager._provider.get(task_id) is None, "should be deleted after kill" + finally: + await manager.shutdown() + mgr_mod._manager = None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_multiturn.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_multiturn.py new file mode 100644 index 000000000000..c5816dcf41ac --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_multiturn.py @@ -0,0 +1,99 @@ +"""Validate MULTI-TURN task timeout hard cap under isolation. + +Key scenario (§5.2): a steerable multi-turn turn that ignores cancellation is +hard-killed after timeout+grace, and if a steering input is queued the chain +DRAINS to the next turn (rather than terminating). If nothing is queued, the +chain suspends. +""" +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import pytest + +_HANDLER_DIR = str(Path(__file__).parent) +_CORE_ROOT = str(Path(__file__).resolve().parents[2]) + + +@pytest.fixture(autouse=True) +def _iso_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, _CORE_ROOT, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +def _meta(info): + return (info.payload or {}).get("metadata", {}) if info else {} + + +@pytest.mark.asyncio +async def test_multiturn_hardcap_drains_to_queued_next_turn(tmp_path): + """Runaway turn 1 + queued steering input B -> hard-kill -> turn 2 runs B.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "chain-drain-1" + # Turn 1: runaway (ignores cancel). + await h.iso_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + await asyncio.sleep(1.0) # let turn 1 start + flush tag A + # Steering input B (cooperative): queued while turn 1 runs. + await h.iso_mt.start(task_id=tid, input={"mode": "coop", "tag": "B"}) + # Turn1 timeout=2s + grace=1s -> kill ~3s, then drain -> turn 2 (B) runs. + await asyncio.sleep(9) + info = await manager._provider.get(tid) + meta = _meta(info) + seen = meta.get("seen_tags", []) + assert "A" in seen, f"turn 1 (A) should have run; seen={seen}" + assert "B" in seen, f"turn 2 (B) should have run after drain; seen={seen}" + assert meta.get("last_tag") == "B", f"last turn should be B; meta={meta}" + # Chain survived (not deleted); ends suspended after B completes. + assert info is not None and info.status in ("suspended", "in_progress"), info.status + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_multiturn_hardcap_no_queue_suspends(tmp_path): + """Runaway turn with NO queued input -> hard-kill -> chain suspends (not deleted).""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "chain-suspend-1" + run = await h.iso_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + # timeout=2s + grace=1s -> kill ~3s. No steering queued. + from azure.ai.agentserver.core.tasks import TaskCancelled + with pytest.raises(TaskCancelled): + await asyncio.wait_for(run.result(), timeout=12) + await asyncio.sleep(0.5) + info = await manager._provider.get(tid) + # Multi-turn: chain stays alive as suspended (nanny won't recover + # in_progress; a future .start() can resume). + assert info is not None, "multi-turn chain should NOT be deleted" + assert info.status == "suspended", f"expected suspended, got {info.status}" + finally: + await manager.shutdown() + mgr_mod._manager = None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_oneshot.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_oneshot.py new file mode 100644 index 000000000000..680edbdfa1b5 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_oneshot.py @@ -0,0 +1,161 @@ +"""Focused validation of ONE-SHOT task timeout cancellation under isolation. + +Scenarios (fast — timeout 1s, grace configurable per test): + 1. Cooperative one-shot winds down at the timeout -> completes, NOT hard-killed. + 2. Runaway one-shot (ignores cancel) -> hard-killed at timeout+grace, + caller gets TaskCancelled, record deleted (ephemeral). + 3. During the grace window the record carries the timeout_cancelled_at marker + and is still in_progress; after the kill it is gone. + 4. Timing: the cooperative window is respected (kill lands ~timeout+grace). + 5. Recovery: a stale in_progress one-shot whose turn already timed out is + finalized (deleted) instead of re-run. +""" +from __future__ import annotations + +import asyncio +import os +import sys +import time +import uuid +from pathlib import Path + +import pytest + +_HANDLER_DIR = str(Path(__file__).parent) +_CORE_ROOT = str(Path(__file__).resolve().parents[2]) + +_TURN_STARTED_AT_KEY = "turn_started_at" +_TIMEOUT_CANCELLED_AT_KEY = "timeout_cancelled_at" + + +@pytest.fixture(autouse=True) +def _iso_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, _CORE_ROOT, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +@pytest.mark.asyncio +async def test_oneshot_cooperative_winds_down_not_killed(tmp_path): + """A cooperative one-shot returns at the timeout and is NOT hard-killed.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + t0 = time.monotonic() + result = await asyncio.wait_for(h.iso_coop.run(input={"n": 1}), timeout=20) + elapsed = time.monotonic() - t0 + # The decisive signal: it returned wound_down=True (cooperative). A + # hard-cap kill would have raised TaskCancelled instead. Timing is a + # loose sanity bound (timeout=3s + worker spawn/import startup). + assert result == {"wound_down": True, "timeout_exceeded": True}, result + assert elapsed < 10.0, f"cooperative wind-down took too long: {elapsed:.2f}s" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_oneshot_runaway_hard_killed_and_deleted(tmp_path): + """A runaway one-shot is force-killed and the record is deleted.""" + from azure.ai.agentserver.core.tasks import TaskCancelled + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + run = await h.iso_never.start(input={"n": 1}) + task_id = run.task_id + t0 = time.monotonic() + with pytest.raises(TaskCancelled): + await asyncio.wait_for(run.result(), timeout=15) + elapsed = time.monotonic() - t0 + # Kill lands after timeout(1s)+grace(1s); cooperative window respected. + assert elapsed >= 1.5, f"killed too early ({elapsed:.2f}s) — grace not respected" + assert elapsed < 8.0, f"killed too late ({elapsed:.2f}s)" + await asyncio.sleep(0.4) + info = await manager._provider.get(task_id) + assert info is None, f"runaway one-shot should be deleted, got status={getattr(info,'status',None)}" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_oneshot_marker_persisted_during_grace(tmp_path): + """During the grace window the record carries the marker + stays in_progress.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + run = await h.iso_never.start(input={"n": 1}) + task_id = run.task_id + # timeout=1s -> cooperative cancel + marker at ~1s; kill at ~2s. + # Sample in the middle of the grace window. + await asyncio.sleep(1.4) + info = await manager._provider.get(task_id) + assert info is not None, "record should still exist mid-grace" + assert info.status == "in_progress", f"expected in_progress mid-grace, got {info.status}" + marker = (info.payload or {}).get(_TIMEOUT_CANCELLED_AT_KEY) + assert marker, f"timeout_cancelled_at marker should be persisted, payload={info.payload}" + # Let the hard cap fire. + with pytest.raises(Exception): + await asyncio.wait_for(run.result(), timeout=8) + await asyncio.sleep(0.4) + assert await manager._provider.get(task_id) is None, "should be deleted after kill" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_oneshot_recovery_finalizes_timed_out_turn(tmp_path): + """Recovery must finalize (delete) a stale in_progress one-shot whose turn + already timed out, rather than re-running it.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + opts = manager._resume_opts.get("iso_never") + assert opts is not None + # Craft a stale record: turn_started_at well in the past (> timeout). + from azure.ai.agentserver.core.tasks._manager import _utc_now_iso + from datetime import datetime, timezone, timedelta + old_ts = (datetime.now(timezone.utc) - timedelta(seconds=30)).isoformat() + + class _Info: + id = "stale-oneshot-1" + status = "in_progress" + payload = {_TURN_STARTED_AT_KEY: old_ts, _TIMEOUT_CANCELLED_AT_KEY: _utc_now_iso()} + source = {"name": "iso_never"} + + # _turn_timed_out should detect it via the marker. + assert manager._turn_timed_out(_Info(), opts) is True + + # And with only the derived backstop (no marker): + class _Info2(_Info): + payload = {_TURN_STARTED_AT_KEY: old_ts} + assert manager._turn_timed_out(_Info2(), opts) is True + + # A fresh turn (started now) must NOT be considered timed out. + class _InfoFresh(_Info): + payload = {_TURN_STARTED_AT_KEY: _utc_now_iso()} + assert manager._turn_timed_out(_InfoFresh(), opts) is False + finally: + await manager.shutdown() + mgr_mod._manager = None diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_reuse.py b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_reuse.py new file mode 100644 index 000000000000..c29d76643ab1 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/tasks/test_isolation_reuse.py @@ -0,0 +1,149 @@ +"""Validate the per-chain persistent worker reuse path (§13.6). + +Reuse is layered on top of isolation and applied ONLY to multi-turn chains: a +single child process is created per ``task_id`` and reused across the chain's +turns (import paid once, not per turn). A hard-cap kill or crash discards the +worker so the next turn re-spawns; a warm worker survives the suspend gap and +is evicted by the idle-TTL reaper. +""" +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import pytest + +_HANDLER_DIR = str(Path(__file__).parent) +_CORE_ROOT = str(Path(__file__).resolve().parents[2]) + + +@pytest.fixture(autouse=True) +def _iso_env(monkeypatch): + monkeypatch.setenv("AGENTSERVER_TASK_ISOLATION", "1") + monkeypatch.setenv("AGENTSERVER_TASK_WORKER_REUSE", "1") + monkeypatch.setenv("AGENTSERVER_TASK_TIMEOUT_HARDCAP_GRACE_SECONDS", "1") + existing = os.environ.get("PYTHONPATH", "") + monkeypatch.setenv("PYTHONPATH", os.pathsep.join([_HANDLER_DIR, _CORE_ROOT, existing])) + if _HANDLER_DIR not in sys.path: + sys.path.insert(0, _HANDLER_DIR) + yield + + +async def _setup(tmp_path): + from azure.ai.agentserver.core.tasks._local_provider import LocalFileTaskProvider + from azure.ai.agentserver.core.tasks._manager import TaskManager + import azure.ai.agentserver.core.tasks._manager as mgr_mod + + provider = LocalFileTaskProvider(Path(str(tmp_path))) + config = type("C", (), { + "agent_name": "iso-agent", "session_id": "iso-session", + "agent_version": "1.0.0", "is_hosted": False, + })() + manager = TaskManager(config=config, provider=provider) + mgr_mod._manager = manager + await manager.startup() + return manager, mgr_mod + + +def _meta(info): + return (info.payload or {}).get("metadata", {}) if info else {} + + +async def _wait_for(manager, tid, predicate, timeout=10.0): + deadline = asyncio.get_event_loop().time() + timeout + last = None + while asyncio.get_event_loop().time() < deadline: + info = await manager._provider.get(tid) + last = info + if predicate(info): + return info + await asyncio.sleep(0.1) + return last + + +@pytest.mark.asyncio +async def test_reuse_same_pid_across_turns(tmp_path): + """Two turns of one chain run in the SAME warm worker process.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "reuse-same-1" + await h.iso_reuse.start(task_id=tid, input={"tag": "A"}) + await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "A") + # A warm worker should be registered for the chain. + assert tid in manager._reuse_workers + pid_after_a = manager._reuse_workers[tid].pid + + # Resume the suspended chain with a second turn. + await h.iso_reuse.start(task_id=tid, input={"tag": "B"}) + info = await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "B") + pids = _meta(info).get("pids", []) + assert len(pids) == 2, f"expected 2 turns recorded; pids={pids}" + assert pids[0] == pids[1], f"turns should reuse ONE worker; pids={pids}" + assert pids[1] == pid_after_a, "registry worker pid should match the running child" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_reuse_hardkill_respawns_new_pid(tmp_path): + """Runaway turn A is hard-killed -> worker discarded -> queued turn B runs + in a NEW worker process (different pid).""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "reuse-kill-1" + await h.iso_reuse_mt.start(task_id=tid, input={"mode": "runaway", "tag": "A"}) + await _wait_for(manager, tid, lambda i: "A" in _meta(i).get("seen_tags", [])) + pid_a = manager._reuse_workers[tid].pid + # Queue turn B while A runs away. + await h.iso_reuse_mt.start(task_id=tid, input={"mode": "coop", "tag": "B"}) + # timeout=2s + grace=1s -> kill ~3s, then drain -> turn B in a new worker. + info = await _wait_for( + manager, tid, lambda i: "B" in _meta(i).get("seen_tags", []), timeout=15 + ) + seen = _meta(info).get("seen_tags", []) + pids = _meta(info).get("pids", []) + assert "A" in seen and "B" in seen, f"both turns should run; seen={seen}" + assert len(pids) == 2 and pids[0] != pids[1], ( + f"hard-kill must respawn a NEW worker; pids={pids}" + ) + assert pids[0] == pid_a, "first pid should be the killed worker" + finally: + await manager.shutdown() + mgr_mod._manager = None + + +@pytest.mark.asyncio +async def test_reuse_idle_ttl_evicts_worker(tmp_path): + """A warm worker idle beyond the TTL is evicted; next turn spawns fresh.""" + import _isolation_handlers as h # noqa: F401 + manager, mgr_mod = await _setup(tmp_path) + try: + tid = "reuse-ttl-1" + await h.iso_reuse.start(task_id=tid, input={"tag": "A"}) + await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "A") + worker = manager._reuse_workers.get(tid) + assert worker is not None and worker.alive + pid_a = worker.pid + + # Force the worker to look idle and drive one reap pass deterministically. + now = asyncio.get_event_loop().time() + worker.last_active_monotonic = now - 3600 # long idle + evicted = manager._reap_idle_workers(now, ttl=1.0) + assert evicted == 1, "idle worker should have been evicted" + assert tid not in manager._reuse_workers, "evicted worker must be deregistered" + + # Next turn spawns a fresh worker (different pid). + await h.iso_reuse.start(task_id=tid, input={"tag": "B"}) + info = await _wait_for(manager, tid, lambda i: _meta(i).get("last_tag") == "B") + pids = _meta(info).get("pids", []) + assert len(pids) == 2 and pids[1] != pid_a, ( + f"post-eviction turn must use a NEW worker; pids={pids}" + ) + finally: + await manager.shutdown() + mgr_mod._manager = None