diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d1f23a..9715586 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 distinguish `READY_TO_ATTEMPT`, `NOT_READY`, and `INDETERMINATE` without running an agent task. Registry, `AgentKit`, and provider-diagnostics helpers expose the same async checks. +- Absolute `AgentTask.deadline` values, `AgentKit.run(timeout=...)`, typed + `AgentTaskTimeoutError`, and `FinishReason.TIMED_OUT` provide bounded task + execution across provider startup and reused-process queueing. +- Built-in runtimes and `AgentKit.cancel()` now return immutable + `CancellationReceipt` values with explicit request dispositions; legacy + third-party `cancel()` methods returning `None` remain supported. ### Changed @@ -69,6 +75,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 during static preflight instead of surfacing later or being silently ignored. - A configured model allow-list now fails closed when model selection is provider-native and therefore cannot be verified locally. +- Cancellation now binds to one active task generation, unregisters without a + cancellable lock window, emits at most one timeout/cancellation terminal, and + gives reusable provider-process teardown a bounded grace period under repeated + cancellation. Runtime instances stay quarantined from new work while detached + teardown remains pending. ## 0.4.0 - 2026-07-02 diff --git a/README.md b/README.md index 9cb287d..a83a58e 100644 --- a/README.md +++ b/README.md @@ -178,5 +178,6 @@ built-in runtime populates it yet, so it is always an empty tuple today. - [Capability matrix](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/capability-matrix.md) - [API stability](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/api-stability.md) - [Live smoke tests](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/live-smoke.md) +- [Deadlines and cancellation](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/task-control.md) - [Mestre migration notes](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/mestre-migration.md) - [SDK evolution agent](https://github.com/ebarti/agent-runtime-kit/blob/main/docs/sdk-evolution-agent.md) diff --git a/docs/api-stability.md b/docs/api-stability.md index af1b1a4..37c9ef3 100644 --- a/docs/api-stability.md +++ b/docs/api-stability.md @@ -23,6 +23,11 @@ import the names from the top-level package instead. party can ship an adapter for a new runtime without forking the enum. - **Every runtime exposes the async lifecycle** (`aclose`, `async with`) declared by the `AgentRuntime` protocol. Stateless runtimes implement it as a no-op. +- **Cancellation receipts describe requests, not rollback.** Built-in runtimes + return an immutable `CancellationReceipt`; the protocol still permits `None` + for existing third-party runtimes, which `AgentKit` reports as + `LEGACY_UNCONFIRMED`. `AgentTask.deadline` is keyword-only and an expired task + raises `AgentTaskTimeoutError` with `FinishReason.TIMED_OUT` event semantics. - **`finish_reason` values** come from `FinishReason`. The field is typed `str` for forward-compatibility, so new reasons can be added without a type break; compare against `FinishReason` members rather than bare literals. diff --git a/docs/capability-matrix.md b/docs/capability-matrix.md index 2cc0844..a3a6c5d 100644 --- a/docs/capability-matrix.md +++ b/docs/capability-matrix.md @@ -12,6 +12,7 @@ | Permission mapping | `permission_mode` | approval mode + sandbox | capabilities + policies | | Streaming output events | Yes — incremental `output.delta` while the SDK runs | No — a single `output.delta` at the end (non-streaming SDK) | Yes — from response chunks | | Tool audit events | Yes — streamed from message blocks | Yes — emitted from parsed `TurnResult` items after the turn | Yes — from tool chunks | +| Deadlines and cancellation | Absolute deadline + coroutine cancellation | Absolute deadline + coroutine cancellation | Absolute deadline + coroutine cancellation | | `vendor.turn` events | No | No | Yes — from thought/unknown chunks | | Package availability | Sync, side-effect-free, package-only | Sync, side-effect-free, package-only | Sync, side-effect-free, package-only; never calls ADC | | Async readiness | Direct API key/OAuth/Bedrock bearer signal → `READY_TO_ATTEMPT`; provider/local chains → `INDETERMINATE` | Supported `AsyncCodex.account(refresh_token=False)` probe; absent account → `NOT_READY` | API key or off-loop ADC/project probe; missing → `NOT_READY` | @@ -33,6 +34,11 @@ For every provider, malformed schemas are rejected locally before dispatch and returned structured values are validated locally after the SDK completes. `parsed_output_available` distinguishes valid JSON `null` from no parsed value. +Deadlines cover provider startup and waits for a reused process. Cancellation +receipts acknowledge a request at the coroutine boundary; they do not imply +rollback of tool side effects that already completed. See +[Deadlines and cancellation](task-control.md). + ## Permission mapping A single portable `PermissionMode` maps to each vendor's native controls. The diff --git a/docs/task-control.md b/docs/task-control.md new file mode 100644 index 0000000..a9fca02 --- /dev/null +++ b/docs/task-control.md @@ -0,0 +1,61 @@ +# Deadlines and cancellation + +`AgentTask.deadline` is an absolute, timezone-aware `datetime`. It applies to +the whole provider operation, including SDK startup and waiting for a reused +provider process. `AgentKit.run(timeout=...)` is a convenience that converts a +finite, non-negative number of seconds (or `timedelta`) into one absolute +deadline at call time. + +```python +from agent_runtime_kit import AgentKit, AgentTaskTimeoutError + +try: + result = await kit.run("codex", goal="Refactor the parser", timeout=30) +except AgentTaskTimeoutError as exc: + print(exc.task_id, exc.deadline) +``` + +An expired deadline never starts the vendor SDK (and therefore emits no started +event). A deadline expiry emits an `agent.task.failed` event with +`finish_reason="timed_out"`, cancels the in-flight provider coroutine, gives its +cleanup a bounded five-second grace period, and raises `AgentTaskTimeoutError` +(which is both an `AgentRuntimeError` and `TimeoutError`). Direct adapter calls +honor `AgentTask(deadline=...)` too. If provider teardown remains wedged after +that grace, the runtime instance is quarantined and rejects new runs rather than +overlapping them; it becomes reusable when the detached cleanup finally settles. + +To cancel a task started through `AgentKit`, keep its task id and use the same +runtime instance or cached kind: + +```python +import asyncio + +task_id = "index-repository" +running = asyncio.create_task( + kit.run("claude", goal="Index the repository", task_id=task_id) +) +receipt = await kit.cancel("claude", task_id) + +try: + await running +except asyncio.CancelledError: + pass + +print(receipt.disposition) +``` + +`cancel()` does not construct a runtime that has not already been cached, and +it does not wait for the cancelled run to settle. Built-in adapters also expose +the same method directly. `CancellationReceipt.disposition` distinguishes a +new request, a repeated request, an inactive id, an unsupported hook, a failed +hook, and a legacy runtime that returned no receipt. + +An active `(runtime, task_id)` identifies one run generation. If a legacy +task-id-only cancellation hook is still settling, `AgentKit` keeps that id +reserved rather than allowing the delayed hook to target a replacement run. + +A `REQUESTED` receipt confirms only that cancellation was requested at the +runtime coroutine boundary. It does not promise rollback: commands, network +requests, or other tool side effects that completed before cancellation may be +permanent. Await the original run task to observe completion of provider +cleanup before reusing related resources. diff --git a/pyproject.toml b/pyproject.toml index 07977f4..24ae57f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -112,6 +112,7 @@ include = [ "docs/quickstart.md", "docs/sdk-evolution-agent-design.md", "docs/sdk-evolution-agent.md", + "docs/task-control.md", ] [tool.ruff] diff --git a/src/agent_runtime_kit/__init__.py b/src/agent_runtime_kit/__init__.py index 416207c..bd6a2ce 100644 --- a/src/agent_runtime_kit/__init__.py +++ b/src/agent_runtime_kit/__init__.py @@ -3,6 +3,7 @@ from agent_runtime_kit._errors import ( AgentRuntimeError, AgentRuntimeUnavailableError, + AgentTaskTimeoutError, OutputSchemaError, OutputTypeError, RuntimeNotRegisteredError, @@ -18,6 +19,8 @@ AgentTask, ArtifactRef, AvailabilityReason, + CancellationDisposition, + CancellationReceipt, EventSink, FilesystemAccess, FinishReason, @@ -66,8 +69,11 @@ "AgentRuntimeKind", "AgentRuntimeUnavailableError", "AgentTask", + "AgentTaskTimeoutError", "ArtifactRef", "AvailabilityReason", + "CancellationDisposition", + "CancellationReceipt", "COMPATIBILITY_MANIFEST", "DEFAULT_READINESS_TIMEOUT", "EventSink", diff --git a/src/agent_runtime_kit/_control.py b/src/agent_runtime_kit/_control.py new file mode 100644 index 0000000..45036ac --- /dev/null +++ b/src/agent_runtime_kit/_control.py @@ -0,0 +1,234 @@ +"""Portable task deadline and cancellation control. + +Cancellation is cooperative at the Python coroutine boundary. A successful +receipt means that cancellation was requested; it cannot promise that vendor +tools have rolled back side effects that already happened. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, TypeVar + +from agent_runtime_kit._errors import AgentRuntimeError, AgentTaskTimeoutError +from agent_runtime_kit._types import ( + AgentRuntimeKind, + AgentTask, + CancellationDisposition, + CancellationReceipt, + FinishReason, +) +from agent_runtime_kit.events import safe_emit, task_failed_event + +_T = TypeVar("_T") +_OPERATION_CLEANUP_TIMEOUT = 5.0 + + +@dataclass +class _ActiveRun: + task: asyncio.Task[Any] + cancellation_requested: bool = False + + +class RuntimeTaskController: + """Track active runs for one runtime and enforce their deadlines.""" + + def __init__(self, kind: AgentRuntimeKind | str) -> None: + self.kind = AgentRuntimeKind.coerce(kind) + self._active: dict[str, _ActiveRun] = {} + self._quarantined: set[asyncio.Future[Any]] = set() + + async def run( + self, + task: AgentTask, + operation: Callable[[], Awaitable[_T]], + ) -> _T: + """Run ``operation`` under task-id registration and deadline control.""" + + current = asyncio.current_task() + if current is None: # pragma: no cover - asyncio always supplies one here + raise AgentRuntimeError("task control requires an active asyncio task") + self._quarantined = {item for item in self._quarantined if not item.done()} + if self._quarantined: + raise AgentRuntimeError( + f"{self.kind} has a cancelled operation still cleaning up; " + "wait for cleanup before starting another run" + ) + # These registry operations deliberately contain no await. asyncio tasks + # share one event-loop thread, so the check-and-set is atomic without a + # cancellable lock acquisition. The same is true of cleanup below. + existing = self._active.get(task.task_id) + if existing is not None: + raise AgentRuntimeError( + f"task_id {task.task_id!r} is already active for {self.kind}" + ) + self._active[task.task_id] = _ActiveRun(current) + + try: + return await self._run_with_deadline(task, operation) + finally: + active = self._active.get(task.task_id) + if active is not None and active.task is current: + del self._active[task.task_id] + + async def cancel( + self, + task_id: str, + *, + _expected_task: asyncio.Task[Any] | None = None, + ) -> CancellationReceipt: + """Request cancellation of an active coroutine without blocking on it.""" + + active = self._active.get(task_id) + if active is None or ( + _expected_task is not None and active.task is not _expected_task + ): + return self._receipt(task_id, CancellationDisposition.NOT_ACTIVE) + if active.cancellation_requested: + return self._receipt(task_id, CancellationDisposition.ALREADY_REQUESTED) + active.cancellation_requested = True + accepted = active.task.cancel() + if not accepted: + return self._receipt( + task_id, + CancellationDisposition.FAILED, + "the active asyncio task had already completed", + ) + return self._receipt(task_id, CancellationDisposition.REQUESTED) + + async def _run_with_deadline( + self, + task: AgentTask, + operation: Callable[[], Awaitable[_T]], + ) -> _T: + deadline = task.deadline + if deadline is None: + try: + return await operation() + except asyncio.CancelledError: + await self._emit_interrupted(task) + raise + + remaining = (deadline - datetime.now(tz=timezone.utc)).total_seconds() + if remaining <= 0: + # An already-expired task never starts, including its started event. + # Emit the terminal first so cancellation of a blocking observability + # sink cannot strand the lifecycle after a misleading start. + await self._emit_timed_out(task, deadline) + raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) + + operation_task: asyncio.Future[_T] = asyncio.ensure_future(operation()) + timed_out = False + try: + done, _ = await asyncio.wait( + {operation_task}, + timeout=remaining, + return_when=asyncio.FIRST_COMPLETED, + ) + if operation_task in done: + return await operation_task + + # Transition before emitting: cancellation of a slow sink must never + # turn one timeout into a second, contradictory interrupted terminal. + timed_out = True + if not await _cancel_and_settle(operation_task): + self._quarantine(operation_task) + await self._emit_timed_out(task, deadline) + raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) + except asyncio.CancelledError: + if not await _cancel_and_settle(operation_task): + self._quarantine(operation_task) + if not timed_out: + await self._emit_interrupted(task) + raise + + async def _emit_interrupted(self, task: AgentTask) -> None: + await safe_emit( + task, + task_failed_event( + task, + self.kind, + error="task cancellation requested", + finish_reason=FinishReason.INTERRUPTED.value, + ), + ) + + async def _emit_timed_out(self, task: AgentTask, deadline: datetime) -> None: + await safe_emit( + task, + task_failed_event( + task, + self.kind, + error=f"task exceeded deadline {deadline.isoformat()}", + finish_reason=FinishReason.TIMED_OUT.value, + ), + ) + + def _receipt( + self, + task_id: str, + disposition: CancellationDisposition, + message: str | None = None, + ) -> CancellationReceipt: + return CancellationReceipt( + task_id=task_id, + kind=self.kind, + disposition=disposition, + message=message, + ) + + def _quarantine(self, operation: asyncio.Future[Any]) -> None: + self._quarantined.add(operation) + + def release(completed: asyncio.Future[Any]) -> None: + self._quarantined.discard(completed) + _consume_operation_result(completed) + + operation.add_done_callback(release) + + +async def run_legacy_with_deadline( + task: AgentTask, + kind: AgentRuntimeKind | str, + operation: Callable[[], Awaitable[_T]], +) -> _T: + """Enforce deadlines for third-party runtimes without native task control.""" + + # A fresh controller is intentional: AgentKit owns cancellation tracking for + # legacy runtimes, while this helper owns only this invocation's deadline and + # terminal event semantics. + controller = RuntimeTaskController(kind) + return await controller.run(task, operation) + + +async def _cancel_and_settle(operation: asyncio.Future[Any]) -> bool: + """Give a cancelled operation bounded, repeat-cancel-safe cleanup time.""" + + operation.cancel() + deadline = asyncio.get_running_loop().time() + _OPERATION_CLEANUP_TIMEOUT + while not operation.done(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + operation.cancel() + return False + try: + await asyncio.wait_for(asyncio.shield(operation), timeout=remaining) + except asyncio.CancelledError: + continue + except asyncio.TimeoutError: + operation.cancel() + return False + except Exception: + break + _consume_operation_result(operation) + return True + + +def _consume_operation_result(operation: asyncio.Future[Any]) -> None: + try: + operation.exception() + except BaseException: + return diff --git a/src/agent_runtime_kit/_errors.py b/src/agent_runtime_kit/_errors.py index 509e51a..0e10caa 100644 --- a/src/agent_runtime_kit/_errors.py +++ b/src/agent_runtime_kit/_errors.py @@ -2,6 +2,8 @@ from __future__ import annotations +from datetime import datetime + from agent_runtime_kit._types import AgentRuntimeKind, runtime_kind_value @@ -17,6 +19,24 @@ def __init__(self, kind: AgentRuntimeKind | str, message: str) -> None: super().__init__(message) +class AgentTaskTimeoutError(AgentRuntimeError, TimeoutError): + """A task exceeded its absolute deadline.""" + + def __init__( + self, + kind: AgentRuntimeKind | str, + task_id: str, + deadline: datetime, + ) -> None: + self.kind: AgentRuntimeKind | str = AgentRuntimeKind.coerce(kind) + self.task_id = task_id + self.deadline = deadline + super().__init__( + f"{runtime_kind_value(self.kind)} task {task_id!r} exceeded deadline " + f"{deadline.isoformat()}" + ) + + class UnsupportedTaskInputError(AgentRuntimeError, ValueError): """A runtime was asked to honor an input it does not support.""" diff --git a/src/agent_runtime_kit/_kit.py b/src/agent_runtime_kit/_kit.py index 53f070f..9fa6649 100644 --- a/src/agent_runtime_kit/_kit.py +++ b/src/agent_runtime_kit/_kit.py @@ -5,10 +5,14 @@ import asyncio import inspect from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from dataclasses import fields as dataclass_fields +from datetime import datetime, timedelta, timezone +from math import isfinite from pathlib import Path from typing import Any, TypeVar, cast, overload +from agent_runtime_kit._control import RuntimeTaskController, run_legacy_with_deadline from agent_runtime_kit._errors import OutputTypeError from agent_runtime_kit._schema import json_schema_for, parse_as from agent_runtime_kit._types import ( @@ -17,6 +21,8 @@ AgentRuntime, AgentRuntimeKind, AgentTask, + CancellationDisposition, + CancellationReceipt, EventSink, FilesystemAccess, FinishReason, @@ -51,6 +57,15 @@ def __repr__(self) -> str: _UNSET = _UnsetType() + +@dataclass +class _KitActiveRun: + agent: AgentRuntime + task: asyncio.Task[Any] + cancellation_requested: bool = False + cancel_hook_complete: bool = False + run_complete: bool = False + # Short spellings for the built-in kinds, resolved only by AgentKit (the # registry itself stays alias-free; the full kind strings always work). KIND_ALIASES: dict[str, AgentRuntimeKind] = { @@ -94,6 +109,7 @@ def __init__( self._runtimes: dict[AgentRuntimeKind | str, AgentRuntime] = {} self._handlers: list[tuple[str, _EventHandler]] = [] self._cache_lock = asyncio.Lock() + self._active_runs: dict[tuple[int, str], _KitActiveRun] = {} @property def registry(self) -> RuntimeRegistry: @@ -193,6 +209,8 @@ async def run( system: str | None = ..., model: str | None = ..., reasoning_effort: str | None = ..., + deadline: datetime | None = ..., + timeout: float | timedelta | None = ..., working_directory: Path | str | None = ..., permissions: PermissionProfile | PermissionMode | str | None = ..., filesystem: FilesystemAccess | str | None = ..., @@ -220,6 +238,8 @@ async def run( system: str | None = ..., model: str | None = ..., reasoning_effort: str | None = ..., + deadline: datetime | None = ..., + timeout: float | timedelta | None = ..., working_directory: Path | str | None = ..., permissions: PermissionProfile | PermissionMode | str | None = ..., filesystem: FilesystemAccess | str | None = ..., @@ -246,6 +266,8 @@ async def run( system: str | None | _UnsetType = _UNSET, model: str | None | _UnsetType = _UNSET, reasoning_effort: str | None | _UnsetType = _UNSET, + deadline: datetime | None | _UnsetType = _UNSET, + timeout: float | timedelta | None = None, working_directory: Path | str | None | _UnsetType = _UNSET, permissions: PermissionProfile | PermissionMode | str | None | _UnsetType = _UNSET, filesystem: FilesystemAccess | str | None | _UnsetType = _UNSET, @@ -273,6 +295,18 @@ async def run( an exception. """ + effective_deadline = _resolve_deadline(deadline, timeout) + if ( + task is not None + and task.deadline is not None + and effective_deadline is not _UNSET + and effective_deadline is not None + ): + raise ValueError( + "task.deadline is mutually exclusive with run(timeout=...) or " + "run(deadline=...)" + ) + if output_type is not None and output_schema is not None: raise ValueError("output_type and output_schema are mutually exclusive") schema = json_schema_for(output_type) if output_type is not None else output_schema @@ -284,6 +318,7 @@ async def run( system=system, model=model, reasoning_effort=reasoning_effort, + deadline=effective_deadline, working_directory=working_directory, permissions=permissions, filesystem=filesystem, @@ -331,6 +366,7 @@ async def run( "system": None if system is _UNSET else system, "model": None if model is _UNSET else model, "reasoning_effort": None if reasoning_effort is _UNSET else reasoning_effort, + "deadline": None if effective_deadline is _UNSET else effective_deadline, "working_directory": _as_path(normalized_working_directory), "mcp_servers": tuple(normalized_mcp_servers), "permissions": _normalize_permissions( @@ -352,10 +388,205 @@ async def run( built = AgentTask(**task_kwargs) agent = runtime if not isinstance(runtime, str) else await self._runtime_for(runtime) - result = await agent.run(built) - if output_type is None: - return result - return _parse_result(output_type, result) + key = (id(agent), built.task_id) + current = asyncio.current_task() + if current is None: # pragma: no cover - asyncio always supplies one here + raise RuntimeError("AgentKit.run requires an active asyncio task") + # No await between lookup and insertion: this is atomic among asyncio + # tasks and, unlike an asyncio.Lock acquisition, cannot be interrupted + # while registering or unregistering a run. + if key in self._active_runs: + raise ValueError(f"task_id {built.task_id!r} is already active on this runtime") + self._active_runs[key] = _KitActiveRun(agent=agent, task=current) + try: + if isinstance(getattr(agent, "_task_controller", None), RuntimeTaskController): + result = await agent.run(built) + else: + result = await run_legacy_with_deadline( + built, + agent.kind, + lambda: agent.run(built), + ) + if output_type is None: + return result + return _parse_result(output_type, result) + finally: + active = self._active_runs.get(key) + if active is not None and active.task is current: + active.run_complete = True + # Keep a tombstone while a legacy cancellation hook is still in + # flight. This prevents a delayed task-id-only hook from + # cancelling a newer run that reused the same id (ABA race). + if not active.cancellation_requested or active.cancel_hook_complete: + del self._active_runs[key] + + async def cancel( + self, + runtime: AgentRuntimeKind | str | AgentRuntime, + task_id: str, + ) -> CancellationReceipt: + """Request cancellation without constructing an uncached runtime. + + A receipt acknowledges the request only. It cannot promise that tools + or other external side effects have been rolled back. + """ + + if not task_id.strip(): + raise ValueError("task_id must be non-blank") + if isinstance(runtime, str): + kind = self._normalize_kind(runtime) + agent = self._runtimes.get(kind) + if agent is None: + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.NOT_ACTIVE, + ) + else: + agent = runtime + kind = agent.kind + + key = (id(agent), task_id) + active = self._active_runs.get(key) + if active is None: + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.NOT_ACTIVE, + ) + if active.cancellation_requested: + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.ALREADY_REQUESTED, + ) + active.cancellation_requested = True + + controller = getattr(agent, "_task_controller", None) + if isinstance(controller, RuntimeTaskController): + try: + # Bind the request to the exact asyncio task generation. A + # task-id-only lookup could otherwise cancel a newer same-id run. + receipt = await controller.cancel(task_id, _expected_task=active.task) + if receipt.disposition in { + CancellationDisposition.NOT_ACTIVE, + CancellationDisposition.UNSUPPORTED, + }: + accepted = active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=( + CancellationDisposition.REQUESTED + if accepted + else CancellationDisposition.FAILED + ), + message=( + "runtime controller no longer owned this run; local " + "coroutine cancellation was requested" + if accepted + else "the active asyncio task had already completed" + ), + ) + return receipt + except asyncio.CancelledError: + # The cancellation caller may itself be cancelled. The target + # still receives a request so the state cannot become poisoned. + active.task.cancel() + raise + except Exception as exc: + active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.FAILED, + message=( + "runtime cancellation controller failed; local cancellation " + f"requested: {exc}" + ), + ) + finally: + self._complete_cancel_hook(key, active) + + # Older third-party runtimes expose only a task-id hook (or no hook). + # Cancel the exact task first, then keep its registry tombstone until the + # hook returns so a delayed hook cannot target a replacement generation. + accepted = active.task.cancel() + if not accepted: + self._complete_cancel_hook(key, active) + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.FAILED, + message="the active asyncio task had already completed", + ) + + cancel_method = getattr(agent, "cancel", None) + if not callable(cancel_method): + self._complete_cancel_hook(key, active) + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.REQUESTED, + message=( + "runtime has no cancellation hook; local coroutine " + "cancellation was requested" + ), + ) + + try: + receipt = await cancel_method(task_id) + except asyncio.CancelledError: + # Local cancellation was already issued; propagate cancellation of + # this request without leaving an unretryable, still-running target. + raise + except Exception as exc: + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.FAILED, + message=( + "runtime cancellation hook failed; local cancellation " + f"requested: {exc}" + ), + ) + finally: + self._complete_cancel_hook(key, active) + + if isinstance(receipt, CancellationReceipt): + if receipt.disposition in { + CancellationDisposition.NOT_ACTIVE, + CancellationDisposition.UNSUPPORTED, + }: + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.REQUESTED, + message=( + "runtime hook could not confirm the task; local coroutine " + "cancellation was requested" + ), + ) + return receipt + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.LEGACY_UNCONFIRMED, + message=( + "legacy runtime returned no cancellation receipt; local coroutine " + "cancellation was requested" + ), + ) + + def _complete_cancel_hook( + self, + key: tuple[int, str], + active: _KitActiveRun, + ) -> None: + active.cancel_hook_complete = True + current = self._active_runs.get(key) + if current is active and active.run_complete: + del self._active_runs[key] async def aclose(self) -> None: """Close every runtime this hub constructed and cached.""" @@ -401,6 +632,7 @@ def _merge_into_task( *, schema: Mapping[str, Any] | None, event_sink: EventSink | None, + deadline: datetime | None | _UnsetType, **field_kwargs: Any, ) -> AgentTask: conflicting = sorted(name for name, value in field_kwargs.items() if value is not _UNSET) @@ -410,6 +642,8 @@ def _merge_into_task( + ", ".join(conflicting) ) replacements: dict[str, Any] = {} + if deadline is not _UNSET: + replacements["deadline"] = deadline if schema is not None: replacements["output_schema"] = schema effective_sink = event_sink if event_sink is not None else task.event_sink @@ -468,6 +702,30 @@ def _as_path(value: Path | str | None) -> Path | None: return Path(value) +def _resolve_deadline( + deadline: datetime | None | _UnsetType, + timeout: float | timedelta | None, +) -> datetime | None | _UnsetType: + if timeout is not None and deadline is not _UNSET and deadline is not None: + raise ValueError("timeout and deadline are mutually exclusive") + if timeout is None: + return deadline + if isinstance(timeout, bool): + raise ValueError("timeout must be a finite non-negative number or timedelta") + if isinstance(timeout, timedelta): + seconds = timeout.total_seconds() + elif isinstance(timeout, (int, float)): + seconds = float(timeout) + else: + raise TypeError("timeout must be a number, timedelta, or None") + if not isfinite(seconds) or seconds < 0: + raise ValueError("timeout must be finite and non-negative") + try: + return datetime.now(tz=timezone.utc) + timedelta(seconds=seconds) + except OverflowError as exc: + raise ValueError("timeout is too large to represent as a deadline") from exc + + def _normalize_permissions( permissions: PermissionProfile | PermissionMode | str | None, filesystem: FilesystemAccess | str | None, diff --git a/src/agent_runtime_kit/_runtime.py b/src/agent_runtime_kit/_runtime.py index b03d9ce..3e47f74 100644 --- a/src/agent_runtime_kit/_runtime.py +++ b/src/agent_runtime_kit/_runtime.py @@ -5,12 +5,14 @@ from collections.abc import Mapping from typing import Any +from agent_runtime_kit._control import RuntimeTaskController from agent_runtime_kit._schema import resolve_structured_output from agent_runtime_kit._types import ( AgentCapabilities, AgentResult, AgentRuntimeKind, AgentTask, + CancellationReceipt, RuntimeAvailability, RuntimeReadiness, TaskSupportReport, @@ -54,6 +56,7 @@ def __init__( self._output = output self._metadata = dict(metadata or {}) self.cancelled_task_ids: set[str] = set() + self._task_controller = RuntimeTaskController(self.kind) def availability(self) -> RuntimeAvailability: """Fake runtime is always available.""" @@ -77,6 +80,11 @@ async def check_readiness(self) -> RuntimeReadiness: ) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled fake task.""" + + return await self._task_controller.run(task, lambda: self._run_task(task)) + + async def _run_task(self, task: AgentTask) -> AgentResult: """Return a deterministic result after validating capabilities.""" await safe_emit(task, task_started_event(task, self.kind)) @@ -145,10 +153,11 @@ async def run(self, task: AgentTask) -> AgentResult: await safe_emit(task, task_completed_event(task, self.kind, result)) return result - async def cancel(self, task_id: str) -> None: - """Record cancellation requests for assertions.""" + async def cancel(self, task_id: str) -> CancellationReceipt: + """Record and request cancellation for an active fake task.""" self.cancelled_task_ids.add(task_id) + return await self._task_controller.cancel(task_id) async def aclose(self) -> None: """No-op: the fake runtime owns no vendor process.""" diff --git a/src/agent_runtime_kit/_types.py b/src/agent_runtime_kit/_types.py index 1b69af2..7a55066 100644 --- a/src/agent_runtime_kit/_types.py +++ b/src/agent_runtime_kit/_types.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field +from datetime import datetime, timezone from enum import Enum from math import isfinite from pathlib import Path @@ -205,6 +206,50 @@ class FinishReason(str, Enum): MAX_TURNS = "max_turns" MAX_TOKENS = "max_tokens" INTERRUPTED = "interrupted" + TIMED_OUT = "timed_out" + + +class CancellationDisposition(str, Enum): + """Outcome of a portable task-cancellation request. + + A cancellation receipt describes the request, not transactional rollback: + provider tools may already have produced external side effects. + """ + + __str__ = str.__str__ + __format__ = str.__format__ # type: ignore[assignment] + + REQUESTED = "requested" + ALREADY_REQUESTED = "already_requested" + NOT_ACTIVE = "not_active" + UNSUPPORTED = "unsupported" + FAILED = "failed" + LEGACY_UNCONFIRMED = "legacy_unconfirmed" + + +@dataclass(frozen=True) +class CancellationReceipt: + """Immutable acknowledgement returned by ``cancel()`` implementations.""" + + task_id: str + kind: AgentRuntimeKind | str + disposition: CancellationDisposition + message: str | None = None + + def __post_init__(self) -> None: + _require_nonblank(self.task_id, "CancellationReceipt.task_id") + object.__setattr__(self, "kind", AgentRuntimeKind.coerce(self.kind)) + object.__setattr__( + self, + "disposition", + _coerce_enum( + CancellationDisposition, + self.disposition, + "CancellationReceipt.disposition", + ), + ) + if self.message is not None: + _require_nonblank(self.message, "CancellationReceipt.message") class FilesystemAccess(str, Enum): @@ -673,6 +718,9 @@ class AgentTask: # working_directory, ...). model: str | None = field(default=None, kw_only=True) reasoning_effort: str | None = field(default=None, kw_only=True) + # An absolute, timezone-aware boundary. Adapters normalize it to UTC and + # include startup and process-reuse queueing in the deadline. + deadline: datetime | None = field(default=None, kw_only=True) working_directory: Path | None = None mcp_servers: tuple[McpServerConfig, ...] = () permissions: PermissionProfile = field(default_factory=PermissionProfile) @@ -694,6 +742,12 @@ def __post_init__(self) -> None: _require_nonblank(self.model, "AgentTask.model") if self.reasoning_effort is not None: _require_nonblank(self.reasoning_effort, "AgentTask.reasoning_effort") + if self.deadline is not None: + if not isinstance(self.deadline, datetime): + raise ValueError("AgentTask.deadline must be a datetime") + if self.deadline.tzinfo is None or self.deadline.utcoffset() is None: + raise ValueError("AgentTask.deadline must be timezone-aware") + object.__setattr__(self, "deadline", self.deadline.astimezone(timezone.utc)) if isinstance(self.sdk_executions, bool) or not isinstance(self.sdk_executions, int): raise ValueError("AgentTask.sdk_executions must be a positive integer") if self.sdk_executions < 1: @@ -813,7 +867,7 @@ def availability(self) -> RuntimeAvailability: async def run(self, task: AgentTask) -> AgentResult: """Execute one task.""" - async def cancel(self, task_id: str) -> None: + async def cancel(self, task_id: str) -> CancellationReceipt | None: """Request cancellation for a task if supported.""" async def aclose(self) -> None: diff --git a/src/agent_runtime_kit/adapters/_common.py b/src/agent_runtime_kit/adapters/_common.py index f49e2e7..841a0d1 100644 --- a/src/agent_runtime_kit/adapters/_common.py +++ b/src/agent_runtime_kit/adapters/_common.py @@ -2,15 +2,16 @@ from __future__ import annotations +import asyncio import inspect import os -from collections.abc import Iterable, Mapping +from collections.abc import Awaitable, Iterable, Mapping from dataclasses import dataclass from importlib import metadata from math import isfinite from typing import Any, Literal -from agent_runtime_kit._errors import UnsupportedTaskInputError +from agent_runtime_kit._errors import AgentRuntimeError, UnsupportedTaskInputError from agent_runtime_kit._schema import ( STRUCTURED_OUTPUT_MISSING as STRUCTURED_OUTPUT_MISSING, ) @@ -28,6 +29,46 @@ from agent_runtime_kit.compatibility import compatibility_for ModelSource = Literal["task", "metadata", "constructor", "provider-native"] +DEFAULT_VENDOR_CLEANUP_TIMEOUT = 5.0 + + +class VendorCleanupTimeoutError(TimeoutError): + """A bounded cleanup grace expired while teardown continues in quarantine.""" + + def __init__(self, timeout: float, pending: asyncio.Future[Any]) -> None: + self.timeout = timeout + self.pending = pending + super().__init__(f"vendor cleanup did not finish within {timeout:g} seconds") + + +class VendorCleanupQuarantine: + """Fail new work fast while detached provider teardown still owns resources.""" + + def __init__(self) -> None: + self._pending: set[asyncio.Future[Any]] = set() + + def track(self, error: BaseException | None) -> None: + if not isinstance(error, VendorCleanupTimeoutError): + return + pending = error.pending + if pending.done(): + _consume_cleanup_result(pending) + return + self._pending.add(pending) + + def release(completed: asyncio.Future[Any]) -> None: + self._pending.discard(completed) + _consume_cleanup_result(completed) + + pending.add_done_callback(release) + + def ensure_ready(self, kind: AgentRuntimeKind | str) -> None: + self._pending = {item for item in self._pending if not item.done()} + if self._pending: + raise AgentRuntimeError( + f"{kind} has provider cleanup still pending; wait before reusing " + "this runtime instance" + ) @dataclass(frozen=True) @@ -374,3 +415,64 @@ async def close_vendor_resource( result = close() if hasattr(result, "__await__"): await result + + +async def finish_vendor_cleanup( + operation: Awaitable[Any], + *, + timeout: float = DEFAULT_VENDOR_CLEANUP_TIMEOUT, +) -> BaseException | None: + """Finish teardown even if the owning run receives repeated cancellation. + + The caller is already handling the original run exception (usually + ``CancelledError``). A second ``Task.cancel()`` must not abandon provider + teardown and expose a half-closed reusable process to the next run. The + grace period is bounded so a wedged vendor close cannot make cancellation + permanently non-returning. The original exception remains authoritative; a + cleanup failure is returned for logging instead of masking it. + """ + + if ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not isfinite(float(timeout)) + or timeout <= 0 + ): + raise ValueError("cleanup timeout must be a positive finite number") + + cleanup: asyncio.Future[Any] = asyncio.ensure_future(operation) + deadline = asyncio.get_running_loop().time() + float(timeout) + while not cleanup.done(): + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + cleanup.cancel() + cleanup.add_done_callback(_consume_cleanup_result) + return VendorCleanupTimeoutError(float(timeout), cleanup) + try: + await asyncio.wait_for(asyncio.shield(cleanup), timeout=remaining) + except asyncio.CancelledError: + # shield keeps the cleanup task alive. Consume any number of + # additional cancellation requests until teardown settles or the + # bounded grace period ends. + continue + except asyncio.TimeoutError: + cleanup.cancel() + cleanup.add_done_callback(_consume_cleanup_result) + return VendorCleanupTimeoutError(float(timeout), cleanup) + except Exception: + # The failure is retrieved and returned below so it can be logged + # without replacing the run's original exception. + break + try: + return cleanup.exception() + except asyncio.CancelledError as exc: + return exc + + +def _consume_cleanup_result(cleanup: asyncio.Future[Any]) -> None: + """Retrieve a detached cleanup result so asyncio never reports it unhandled.""" + + try: + cleanup.exception() + except BaseException: + return diff --git a/src/agent_runtime_kit/adapters/antigravity.py b/src/agent_runtime_kit/adapters/antigravity.py index 92efbf6..e470f7e 100644 --- a/src/agent_runtime_kit/adapters/antigravity.py +++ b/src/agent_runtime_kit/adapters/antigravity.py @@ -13,6 +13,7 @@ from pathlib import Path from typing import Any +from agent_runtime_kit._control import RuntimeTaskController from agent_runtime_kit._errors import AgentRuntimeUnavailableError, UnsupportedTaskInputError from agent_runtime_kit._types import ( AgentCapabilities, @@ -20,6 +21,7 @@ AgentRuntimeKind, AgentTask, AvailabilityReason, + CancellationReceipt, FilesystemAccess, PermissionMode, ReadinessStatus, @@ -31,10 +33,12 @@ Usage, ) from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, filter_supported_kwargs, fingerprint_item, + finish_vendor_cleanup, model_support_issue, optional_int, optional_str, @@ -72,7 +76,7 @@ class AntigravityAgentRuntime: structured_output=True, streaming=True, tool_audit=True, - cancellation=False, + cancellation=True, tool_filters=True, ) @@ -113,6 +117,8 @@ def __init__( self._sdk_process_reuse_count = 0 self._agent_lock = asyncio.Lock() self._agent_run_lock = asyncio.Lock() + self._task_controller = RuntimeTaskController(self.kind) + self._cleanup_quarantine = VendorCleanupQuarantine() async def __aenter__(self) -> AntigravityAgentRuntime: return self @@ -227,6 +233,12 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: return TaskSupportReport(kind=self.kind, issues=tuple(issues)) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled Antigravity task.""" + + self._cleanup_quarantine.ensure_ready(self.kind) + return await self._task_controller.run(task, lambda: self._run_task(task)) + + async def _run_task(self, task: AgentTask) -> AgentResult: """Execute one task with Antigravity.""" await safe_emit(task, task_started_event(task, self.kind)) @@ -269,10 +281,10 @@ async def run(self, task: AgentTask) -> AgentResult: await safe_emit(task, task_completed_event(task, self.kind, result)) return result - async def cancel(self, task_id: str) -> None: - """Antigravity cancellation is not exposed through this portable adapter yet.""" + async def cancel(self, task_id: str) -> CancellationReceipt: + """Request cancellation at the adapter coroutine boundary.""" - del task_id + return await self._task_controller.cancel(task_id) async def aclose(self) -> None: """Close any reusable Antigravity agent process owned by this runtime. @@ -281,6 +293,7 @@ async def aclose(self) -> None: in-flight ``run()`` instead of closing the agent mid-turn. """ + self._cleanup_quarantine.ensure_ready(self.kind) async with self._agent_run_lock: await self._close_agent() @@ -408,9 +421,9 @@ async def _invoke( # next run() never reuses a poisoned process. The run lock is # already held, so close under the agent lock only and never # let cleanup mask the original error. - try: - await self._close_agent() - except Exception as close_exc: + close_exc = await finish_vendor_cleanup(self._close_agent()) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Antigravity agent after run failure: %s", close_exc, @@ -554,9 +567,9 @@ async def _persistent_agent( try: agent = await enter() if callable(enter) else context except BaseException: - try: - await close_vendor_resource(context) - except Exception as close_exc: + close_exc = await finish_vendor_cleanup(close_vendor_resource(context)) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Antigravity agent after startup failure: %s", close_exc ) diff --git a/src/agent_runtime_kit/adapters/claude.py b/src/agent_runtime_kit/adapters/claude.py index 4273049..c13f51d 100644 --- a/src/agent_runtime_kit/adapters/claude.py +++ b/src/agent_runtime_kit/adapters/claude.py @@ -10,6 +10,7 @@ from math import isfinite from typing import Any +from agent_runtime_kit._control import RuntimeTaskController from agent_runtime_kit._errors import AgentRuntimeUnavailableError from agent_runtime_kit._types import ( AgentCapabilities, @@ -17,6 +18,7 @@ AgentRuntimeKind, AgentTask, AvailabilityReason, + CancellationReceipt, FilesystemAccess, PermissionMode, PermissionProfile, @@ -29,11 +31,13 @@ ) from agent_runtime_kit.adapters._common import ( STRUCTURED_OUTPUT_MISSING, + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, field_value, filter_supported_kwargs, fingerprint_value, + finish_vendor_cleanup, metadata_str, model_support_issue, optional_int, @@ -69,7 +73,7 @@ class ClaudeAgentRuntime: structured_output=True, streaming=True, tool_audit=True, - cancellation=False, + cancellation=True, budget=True, reasoning_effort=True, tool_filters=True, @@ -102,6 +106,8 @@ def __init__( self._sdk_process_reuse_count = 0 self._client_lock = asyncio.Lock() self._client_run_lock = asyncio.Lock() + self._task_controller = RuntimeTaskController(self.kind) + self._cleanup_quarantine = VendorCleanupQuarantine() async def __aenter__(self) -> ClaudeAgentRuntime: return self @@ -164,6 +170,12 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: ) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled Claude task.""" + + self._cleanup_quarantine.ensure_ready(self.kind) + return await self._task_controller.run(task, lambda: self._run_task(task)) + + async def _run_task(self, task: AgentTask) -> AgentResult: """Execute one task with Claude Agent SDK, streaming events as they arrive.""" await safe_emit(task, task_started_event(task, self.kind)) @@ -219,10 +231,10 @@ async def run(self, task: AgentTask) -> AgentResult: await safe_emit(task, task_completed_event(task, self.kind, result)) return result - async def cancel(self, task_id: str) -> None: - """Claude ``query`` calls do not expose a portable cancellation handle.""" + async def cancel(self, task_id: str) -> CancellationReceipt: + """Request cancellation at the adapter coroutine boundary.""" - del task_id + return await self._task_controller.cancel(task_id) async def aclose(self) -> None: """Close any reusable Claude SDK client process owned by this runtime. @@ -231,6 +243,7 @@ async def aclose(self) -> None: in-flight ``run()`` to finish instead of closing the client mid-stream. """ + self._cleanup_quarantine.ensure_ready(self.kind) async with self._client_run_lock: await self._close_client() @@ -291,9 +304,9 @@ async def _run_with_client( # poisoned process is never handed to the next run(). The run # lock is already held, so close under the client lock only and # never let cleanup mask the original error. - try: - await self._close_client() - except Exception as close_exc: + close_exc = await finish_vendor_cleanup(self._close_client()) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Claude client after run failure: %s", close_exc ) @@ -323,9 +336,11 @@ async def _persistent_client( if callable(enter): client = await enter() except BaseException: - try: - await close_vendor_resource(client, try_disconnect=True) - except Exception as close_exc: + close_exc = await finish_vendor_cleanup( + close_vendor_resource(client, try_disconnect=True) + ) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Claude client after startup failure: %s", close_exc ) diff --git a/src/agent_runtime_kit/adapters/codex.py b/src/agent_runtime_kit/adapters/codex.py index b14ede6..5d0ffc3 100644 --- a/src/agent_runtime_kit/adapters/codex.py +++ b/src/agent_runtime_kit/adapters/codex.py @@ -8,6 +8,7 @@ from collections.abc import Mapping from typing import Any +from agent_runtime_kit._control import RuntimeTaskController from agent_runtime_kit._errors import AgentRuntimeUnavailableError from agent_runtime_kit._types import ( AgentCapabilities, @@ -15,6 +16,7 @@ AgentRuntimeKind, AgentTask, AvailabilityReason, + CancellationReceipt, FilesystemAccess, PermissionMode, ReadinessStatus, @@ -25,10 +27,12 @@ Usage, ) from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, field_value, filter_supported_kwargs, + finish_vendor_cleanup, metadata_str, model_support_issue, optional_int, @@ -74,7 +78,7 @@ class CodexAgentRuntime: structured_output=True, streaming=False, tool_audit=True, - cancellation=False, + cancellation=True, reasoning_effort=True, ) @@ -111,6 +115,8 @@ def __init__( self._sdk_process_reuse_count = 0 self._codex_client_lock = asyncio.Lock() self._codex_run_lock = asyncio.Lock() + self._task_controller = RuntimeTaskController(self.kind) + self._cleanup_quarantine = VendorCleanupQuarantine() async def __aenter__(self) -> CodexAgentRuntime: return self @@ -192,6 +198,12 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: ) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled Codex task.""" + + self._cleanup_quarantine.ensure_ready(self.kind) + return await self._task_controller.run(task, lambda: self._run_task(task)) + + async def _run_task(self, task: AgentTask) -> AgentResult: """Execute one task with the Codex SDK.""" await safe_emit(task, task_started_event(task, self.kind)) @@ -239,10 +251,10 @@ async def run(self, task: AgentTask) -> AgentResult: await safe_emit(task, task_completed_event(task, self.kind, result)) return result - async def cancel(self, task_id: str) -> None: - """Codex SDK cancellation is not exposed through this portable adapter yet.""" + async def cancel(self, task_id: str) -> CancellationReceipt: + """Request cancellation at the adapter coroutine boundary.""" - del task_id + return await self._task_controller.cancel(task_id) async def aclose(self) -> None: """Close any reusable Codex app-server process owned by this runtime. @@ -251,6 +263,7 @@ async def aclose(self) -> None: in-flight ``run()`` instead of closing the process mid-turn. """ + self._cleanup_quarantine.ensure_ready(self.kind) async with self._codex_run_lock: await self._close_codex_client() @@ -359,9 +372,9 @@ async def _run_codex( # inside the run lock so a queued run cannot observe the # doomed client between failure and eviction; close under # the client lock only and never mask the original error. - try: - await self._close_codex_client() - except Exception as close_exc: + close_exc = await finish_vendor_cleanup(self._close_codex_client()) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Codex app-server after run failure: %s", close_exc ) @@ -447,9 +460,9 @@ async def _persistent_codex_client( try: client = await enter() if callable(enter) else context except BaseException: - try: - await close_vendor_resource(context) - except Exception as close_exc: + close_exc = await finish_vendor_cleanup(close_vendor_resource(context)) + self._cleanup_quarantine.track(close_exc) + if close_exc is not None: logger.warning( "failed to close Codex app-server after startup failure: %s", close_exc ) diff --git a/tests/test_control.py b/tests/test_control.py new file mode 100644 index 0000000..51f2369 --- /dev/null +++ b/tests/test_control.py @@ -0,0 +1,662 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta, timezone + +import pytest + +import agent_runtime_kit._control as control_module +from agent_runtime_kit import ( + AgentKit, + AgentResult, + AgentRuntimeError, + AgentRuntimeKind, + AgentTask, + AgentTaskTimeoutError, + CancellationDisposition, + CancellationReceipt, + FakeAgentRuntime, + FinishReason, + RuntimeAvailability, + RuntimeRegistry, +) +from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, + finish_vendor_cleanup, +) +from agent_runtime_kit.testing import RecordingEventSink + + +class LegacyBlockingRuntime: + """A pre-receipt third-party runtime used to verify compatibility.""" + + kind = AgentRuntimeKind.FAKE + capabilities = FakeAgentRuntime().capabilities + + def __init__(self) -> None: + self.started = asyncio.Event() + self.release = asyncio.Event() + self.run_calls = 0 + self.cancel_calls = 0 + self.cancelled = False + + def availability(self) -> RuntimeAvailability: + return RuntimeAvailability.ok(self.kind) + + async def run(self, task: AgentTask) -> AgentResult: + del task + self.run_calls += 1 + self.started.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled = True + raise + return AgentResult(output="done") + + async def cancel(self, task_id: str) -> None: + del task_id + self.cancel_calls += 1 + + async def aclose(self) -> None: + return None + + async def __aenter__(self) -> LegacyBlockingRuntime: + return self + + async def __aexit__(self, _exc_type: object, _exc: object, _tb: object) -> None: + await self.aclose() + + +class BlockingFakeRuntime(FakeAgentRuntime): + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + self.cancelled = False + + async def _run_task(self, task: AgentTask) -> AgentResult: + self.started.set() + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled = True + raise + return AgentResult(output=task.goal) + + +class SerialFakeRuntime(FakeAgentRuntime): + def __init__(self) -> None: + super().__init__() + self.lock = asyncio.Lock() + self.holder_started = asyncio.Event() + self.release_holder = asyncio.Event() + self.holder_cancelled = False + + async def _run_task(self, task: AgentTask) -> AgentResult: + async with self.lock: + if task.goal == "holder": + self.holder_started.set() + try: + await self.release_holder.wait() + except asyncio.CancelledError: + self.holder_cancelled = True + raise + return AgentResult(output=task.goal) + + +class GatedCancellationRuntime: + """Runtime whose cancellation hook can be delayed across run generations.""" + + kind = AgentRuntimeKind.FAKE + capabilities = FakeAgentRuntime().capabilities + + def __init__(self) -> None: + self._active: dict[str, asyncio.Task[object]] = {} + self.first_started = asyncio.Event() + self.release_first = asyncio.Event() + self.second_started = asyncio.Event() + self.release_second = asyncio.Event() + self.cancel_entered = asyncio.Event() + self.release_cancel = asyncio.Event() + self.cancel_calls = 0 + self.second_cancelled = False + + def availability(self) -> RuntimeAvailability: + return RuntimeAvailability.ok(self.kind) + + async def run(self, task: AgentTask) -> AgentResult: + current = asyncio.current_task() + assert current is not None + self._active[task.task_id] = current + try: + return await self._run_task(task) + finally: + if self._active.get(task.task_id) is current: + del self._active[task.task_id] + + async def _run_task(self, task: AgentTask) -> AgentResult: + if task.goal == "first": + self.first_started.set() + await self.release_first.wait() + return AgentResult(output="first") + self.second_started.set() + try: + await self.release_second.wait() + except asyncio.CancelledError: + self.second_cancelled = True + raise + return AgentResult(output="second") + + async def cancel(self, task_id: str) -> CancellationReceipt: + self.cancel_calls += 1 + self.cancel_entered.set() + await self.release_cancel.wait() + target = self._active.get(task_id) + if target is None: + return CancellationReceipt( + task_id=task_id, + kind=self.kind, + disposition=CancellationDisposition.NOT_ACTIVE, + ) + target.cancel() + return CancellationReceipt( + task_id=task_id, + kind=self.kind, + disposition=CancellationDisposition.REQUESTED, + ) + + async def aclose(self) -> None: + return None + + async def __aenter__(self) -> GatedCancellationRuntime: + return self + + async def __aexit__(self, _exc_type: object, _exc: object, _tb: object) -> None: + await self.aclose() + + +class BlockingTimeoutSink: + """Record a timeout terminal, then pause its delivery to expose races.""" + + def __init__(self) -> None: + self.events: list[object] = [] + self.timeout_seen = asyncio.Event() + self.release_timeout = asyncio.Event() + + async def emit(self, event: object) -> None: + assert isinstance(event, dict) + self.events.append(event) + attributes = event.get("attributes", {}) + if ( + event.get("name") == "agent.task.failed" + and isinstance(attributes, dict) + and attributes.get("finish_reason") == FinishReason.TIMED_OUT + ): + self.timeout_seen.set() + await self.release_timeout.wait() + + +class CleanupBlockingRuntime(FakeAgentRuntime): + """Expose whether a repeated cancel interrupts operation cleanup.""" + + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.cleanup_started = asyncio.Event() + self.release_cleanup = asyncio.Event() + self.cleanup_interrupted = False + + async def _run_task(self, task: AgentTask) -> AgentResult: + self.started.set() + try: + await asyncio.Event().wait() + finally: + self.cleanup_started.set() + try: + await self.release_cleanup.wait() + except asyncio.CancelledError: + self.cleanup_interrupted = True + raise + + +class StubbornCleanupRuntime(FakeAgentRuntime): + """Keep a deadline-cancelled operation alive until a test releases it.""" + + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.cleanup_started = asyncio.Event() + self.release_cleanup = asyncio.Event() + + async def _run_task(self, task: AgentTask) -> AgentResult: + if task.goal == "replacement": + return AgentResult(output="replacement") + self.started.set() + try: + await asyncio.Event().wait() + finally: + self.cleanup_started.set() + try: + await self.release_cleanup.wait() + except asyncio.CancelledError: + await self.release_cleanup.wait() + + +def test_deadline_requires_timezone_and_normalizes_to_utc() -> None: + with pytest.raises(ValueError, match="must be a datetime"): + AgentTask(goal="x", deadline="tomorrow") # type: ignore[arg-type] + + with pytest.raises(ValueError, match="timezone-aware"): + AgentTask(goal="x", deadline=datetime(2030, 1, 1)) + + plus_two = timezone(timedelta(hours=2)) + task = AgentTask(goal="x", deadline=datetime(2030, 1, 1, 12, tzinfo=plus_two)) + + assert task.deadline == datetime(2030, 1, 1, 10, tzinfo=timezone.utc) + + +@pytest.mark.parametrize("timeout", [-1, float("inf"), float("nan"), True]) +@pytest.mark.asyncio +async def test_kit_rejects_invalid_timeouts(timeout: float) -> None: + kit = AgentKit(register_default_adapters=False) + + with pytest.raises(ValueError, match="timeout"): + await kit.run(FakeAgentRuntime(), goal="x", timeout=timeout) + + +@pytest.mark.asyncio +async def test_timeout_and_existing_deadline_are_mutually_exclusive() -> None: + task = AgentTask( + goal="x", + deadline=datetime.now(tz=timezone.utc) + timedelta(seconds=5), + ) + + with pytest.raises(ValueError, match="mutually exclusive"): + await AgentKit(register_default_adapters=False).run( + FakeAgentRuntime(), task=task, timeout=1 + ) + + +@pytest.mark.asyncio +async def test_expired_deadline_never_invokes_legacy_runtime() -> None: + runtime = LegacyBlockingRuntime() + sink = RecordingEventSink() + deadline = datetime.now(tz=timezone.utc) - timedelta(seconds=1) + + with pytest.raises(AgentTaskTimeoutError) as caught: + await AgentKit(register_default_adapters=False).run( + runtime, + goal="too late", + task_id="expired", + deadline=deadline, + event_sink=sink, + ) + + assert runtime.run_calls == 0 + assert caught.value.task_id == "expired" + assert caught.value.deadline == deadline + assert isinstance(caught.value, TimeoutError) + assert [event["name"] for event in sink.events] == ["agent.task.failed"] + assert sink.events[-1]["attributes"]["finish_reason"] == FinishReason.TIMED_OUT + + +@pytest.mark.asyncio +async def test_kit_timeout_cancels_legacy_runtime_and_next_run_recovers() -> None: + runtime = LegacyBlockingRuntime() + kit = AgentKit(register_default_adapters=False) + + with pytest.raises(AgentTaskTimeoutError): + await kit.run(runtime, goal="slow", timeout=0.01) + + assert runtime.cancelled is True + runtime.started.clear() + runtime.release.set() + result = await kit.run(runtime, goal="next", timeout=1) + assert result.output == "done" + + +@pytest.mark.asyncio +async def test_kit_cancels_cached_runtime_without_constructing_another() -> None: + runtime = LegacyBlockingRuntime() + constructions = 0 + + def factory() -> LegacyBlockingRuntime: + nonlocal constructions + constructions += 1 + return runtime + + registry = RuntimeRegistry() + registry.register(AgentRuntimeKind.FAKE, factory) + kit = AgentKit(registry=registry) + running = asyncio.create_task(kit.run("fake", goal="wait", task_id="cancel-me")) + await runtime.started.wait() + + receipt = await kit.cancel("fake", "cancel-me") + repeated = await kit.cancel("fake", "cancel-me") + + assert receipt.disposition is CancellationDisposition.LEGACY_UNCONFIRMED + assert repeated.disposition is CancellationDisposition.ALREADY_REQUESTED + assert runtime.cancel_calls == 1 + assert constructions == 1 + with pytest.raises(asyncio.CancelledError): + await running + assert runtime.cancelled is True + + +@pytest.mark.asyncio +async def test_cancel_of_uncached_kind_is_not_active_and_has_no_side_effects() -> None: + constructions = 0 + + def factory() -> FakeAgentRuntime: + nonlocal constructions + constructions += 1 + return FakeAgentRuntime() + + registry = RuntimeRegistry() + registry.register(AgentRuntimeKind.FAKE, factory) + kit = AgentKit(registry=registry) + + receipt = await kit.cancel("fake", "missing") + + assert receipt.disposition is CancellationDisposition.NOT_ACTIVE + assert constructions == 0 + + +@pytest.mark.asyncio +async def test_direct_builtin_cancel_returns_receipts_and_emits_interrupted() -> None: + runtime = BlockingFakeRuntime() + sink = RecordingEventSink() + running = asyncio.create_task( + runtime.run(AgentTask(goal="wait", task_id="direct", event_sink=sink)) + ) + await runtime.started.wait() + + receipt = await runtime.cancel("direct") + repeated = await runtime.cancel("direct") + + assert isinstance(receipt, CancellationReceipt) + assert receipt.disposition is CancellationDisposition.REQUESTED + assert repeated.disposition is CancellationDisposition.ALREADY_REQUESTED + with pytest.raises(asyncio.CancelledError): + await running + assert runtime.cancelled is True + terminal = [event for event in sink.events if event["name"] == "agent.task.failed"] + assert len(terminal) == 1 + assert terminal[0]["attributes"]["finish_reason"] == FinishReason.INTERRUPTED + + runtime.started.clear() + runtime.release.set() + recovered = await runtime.run(AgentTask(goal="next", task_id="direct")) + assert recovered.output == "next" + + +@pytest.mark.asyncio +async def test_queued_deadline_does_not_interrupt_reuse_holder() -> None: + runtime = SerialFakeRuntime() + holder = asyncio.create_task( + runtime.run(AgentTask(goal="holder", task_id="holder")) + ) + await runtime.holder_started.wait() + + with pytest.raises(AgentTaskTimeoutError): + await runtime.run( + AgentTask( + goal="queued", + task_id="queued", + deadline=datetime.now(tz=timezone.utc) + timedelta(milliseconds=10), + ) + ) + + assert runtime.holder_cancelled is False + assert not holder.done() + runtime.release_holder.set() + assert (await holder).output == "holder" + + +@pytest.mark.asyncio +async def test_external_task_cancellation_emits_one_interrupted_terminal() -> None: + runtime = BlockingFakeRuntime() + sink = RecordingEventSink() + running = asyncio.create_task( + runtime.run(AgentTask(goal="wait", task_id="external", event_sink=sink)) + ) + await runtime.started.wait() + + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + + terminal = [event for event in sink.events if event["name"] == "agent.task.failed"] + assert len(terminal) == 1 + assert terminal[0]["attributes"]["finish_reason"] == FinishReason.INTERRUPTED + + +@pytest.mark.asyncio +async def test_delayed_cancel_hook_cannot_target_later_same_id_run() -> None: + runtime = GatedCancellationRuntime() + kit = AgentKit(register_default_adapters=False) + first = asyncio.create_task(kit.run(runtime, goal="first", task_id="same")) + await runtime.first_started.wait() + + cancelling = asyncio.create_task(kit.cancel(runtime, "same")) + await runtime.cancel_entered.wait() + with pytest.raises(asyncio.CancelledError): + await first + + # The completed generation remains reserved until its delayed task-id hook + # settles, so no replacement can become the hook's accidental target. + with pytest.raises(ValueError, match="already active"): + await kit.run(runtime, goal="second", task_id="same") + + runtime.release_cancel.set() + receipt = await cancelling + assert receipt.disposition is CancellationDisposition.REQUESTED + + second = asyncio.create_task(kit.run(runtime, goal="second", task_id="same")) + await runtime.second_started.wait() + assert runtime.second_cancelled is False + runtime.release_second.set() + assert (await second).output == "second" + + +@pytest.mark.asyncio +async def test_cancelling_cancel_request_still_cancels_target_and_releases_generation() -> None: + runtime = GatedCancellationRuntime() + kit = AgentKit(register_default_adapters=False) + running = asyncio.create_task(kit.run(runtime, goal="first", task_id="same")) + await runtime.first_started.wait() + + cancelling = asyncio.create_task(kit.cancel(runtime, "same")) + await runtime.cancel_entered.wait() + cancelling.cancel() + + with pytest.raises(asyncio.CancelledError): + await cancelling + with pytest.raises(asyncio.CancelledError): + await running + + replacement = asyncio.create_task(kit.run(runtime, goal="second", task_id="same")) + await runtime.second_started.wait() + runtime.release_second.set() + assert (await replacement).output == "second" + + +@pytest.mark.asyncio +async def test_timeout_racing_external_cancel_emits_only_timed_out_terminal() -> None: + runtime = BlockingFakeRuntime() + sink = BlockingTimeoutSink() + running = asyncio.create_task( + runtime.run( + AgentTask( + goal="wait", + task_id="timeout-race", + deadline=datetime.now(tz=timezone.utc) + timedelta(milliseconds=10), + event_sink=sink, # type: ignore[arg-type] + ) + ) + ) + await sink.timeout_seen.wait() + + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + + terminals = [ + event + for event in sink.events + if isinstance(event, dict) and event.get("name") == "agent.task.failed" + ] + assert len(terminals) == 1 + assert terminals[0]["attributes"]["finish_reason"] == FinishReason.TIMED_OUT + + +@pytest.mark.asyncio +async def test_timeout_wins_over_external_cancel_during_operation_cleanup() -> None: + runtime = CleanupBlockingRuntime() + sink = RecordingEventSink() + running = asyncio.create_task( + runtime.run( + AgentTask( + goal="wait", + task_id="cleanup-timeout-race", + deadline=datetime.now(tz=timezone.utc) + timedelta(milliseconds=10), + event_sink=sink, + ) + ) + ) + await runtime.cleanup_started.wait() + + running.cancel() + await asyncio.sleep(0) + runtime.release_cleanup.set() + + with pytest.raises(AgentTaskTimeoutError): + await running + terminal = [event for event in sink.events if event["name"] == "agent.task.failed"] + assert len(terminal) == 1 + assert terminal[0]["attributes"]["finish_reason"] == FinishReason.TIMED_OUT + + +@pytest.mark.asyncio +async def test_repeated_runtime_cancel_does_not_interrupt_operation_cleanup() -> None: + runtime = CleanupBlockingRuntime() + running = asyncio.create_task( + runtime.run(AgentTask(goal="wait", task_id="cleanup")) + ) + await runtime.started.wait() + + first = await runtime.cancel("cleanup") + await runtime.cleanup_started.wait() + repeated = await runtime.cancel("cleanup") + + assert first.disposition is CancellationDisposition.REQUESTED + assert repeated.disposition is CancellationDisposition.ALREADY_REQUESTED + assert runtime.cleanup_interrupted is False + runtime.release_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await running + assert runtime.cleanup_interrupted is False + + +@pytest.mark.asyncio +async def test_vendor_cleanup_helper_survives_repeated_task_cancellation() -> None: + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + cleanup_finished = False + + async def cleanup() -> None: + nonlocal cleanup_finished + cleanup_started.set() + await release_cleanup.wait() + cleanup_finished = True + + async def owner() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + assert await finish_vendor_cleanup(cleanup()) is None + raise + + running = asyncio.create_task(owner()) + await asyncio.sleep(0) + running.cancel() + await cleanup_started.wait() + running.cancel() + await asyncio.sleep(0) + + assert cleanup_finished is False + release_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await running + assert cleanup_finished is True + + +@pytest.mark.asyncio +async def test_vendor_cleanup_helper_has_bounded_liveness() -> None: + cleanup_started = asyncio.Event() + cleanup_cancelled = asyncio.Event() + release_cleanup = asyncio.Event() + + async def wedged_cleanup() -> None: + cleanup_started.set() + try: + await release_cleanup.wait() + except asyncio.CancelledError: + cleanup_cancelled.set() + await release_cleanup.wait() + + outcome = await finish_vendor_cleanup(wedged_cleanup(), timeout=0.01) + await asyncio.sleep(0) + + assert cleanup_started.is_set() + assert cleanup_cancelled.is_set() + assert isinstance(outcome, TimeoutError) + quarantine = VendorCleanupQuarantine() + quarantine.track(outcome) + with pytest.raises(AgentRuntimeError, match="cleanup still pending"): + quarantine.ensure_ready(AgentRuntimeKind.FAKE) + release_cleanup.set() + await asyncio.sleep(0) + quarantine.ensure_ready(AgentRuntimeKind.FAKE) + + +@pytest.mark.asyncio +async def test_detached_operation_quarantines_runtime_until_cleanup_settles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(control_module, "_OPERATION_CLEANUP_TIMEOUT", 0.01) + runtime = StubbornCleanupRuntime() + + with pytest.raises(AgentTaskTimeoutError): + await runtime.run( + AgentTask( + goal="stubborn", + task_id="old", + deadline=datetime.now(tz=timezone.utc) + timedelta(milliseconds=10), + ) + ) + + assert runtime.cleanup_started.is_set() + with pytest.raises(AgentRuntimeError, match="still cleaning up"): + await runtime.run(AgentTask(goal="replacement", task_id="new")) + + runtime.release_cleanup.set() + await asyncio.sleep(0) + await asyncio.sleep(0) + assert ( + await runtime.run(AgentTask(goal="replacement", task_id="new")) + ).output == "replacement" + + +@pytest.mark.asyncio +async def test_cancel_after_completed_run_is_truthfully_not_active() -> None: + runtime = LegacyBlockingRuntime() + runtime.release.set() + kit = AgentKit(register_default_adapters=False) + + assert (await kit.run(runtime, goal="done", task_id="finished")).output == "done" + receipt = await kit.cancel(runtime, "finished") + + assert receipt.disposition is CancellationDisposition.NOT_ACTIVE + assert runtime.cancel_calls == 0