From d47f9d20622b1f668109cacaffd6e4797d491d85 Mon Sep 17 00:00:00 2001 From: Eloi Date: Fri, 10 Jul 2026 14:05:10 +0200 Subject: [PATCH 1/4] feat: add task deadlines and cancellation --- CHANGELOG.md | 6 + README.md | 1 + docs/api-stability.md | 5 + docs/capability-matrix.md | 6 + docs/task-control.md | 54 ++++ pyproject.toml | 1 + src/agent_runtime_kit/__init__.py | 6 + src/agent_runtime_kit/_control.py | 176 +++++++++++ src/agent_runtime_kit/_errors.py | 20 ++ src/agent_runtime_kit/_kit.py | 200 ++++++++++++- src/agent_runtime_kit/_runtime.py | 13 +- src/agent_runtime_kit/_types.py | 54 +++- src/agent_runtime_kit/adapters/antigravity.py | 16 +- src/agent_runtime_kit/adapters/claude.py | 16 +- src/agent_runtime_kit/adapters/codex.py | 16 +- tests/test_control.py | 283 ++++++++++++++++++ 16 files changed, 854 insertions(+), 19 deletions(-) create mode 100644 docs/task-control.md create mode 100644 src/agent_runtime_kit/_control.py create mode 100644 tests/test_control.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d1f23a..e1b4a7e 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 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..77c319d --- /dev/null +++ b/docs/task-control.md @@ -0,0 +1,54 @@ +# 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. A deadline expiry emits an +`agent.task.failed` event with `finish_reason="timed_out"`, cancels the in-flight +provider coroutine, waits for its cleanup, and raises `AgentTaskTimeoutError` +(which is both an `AgentRuntimeError` and `TimeoutError`). Direct adapter calls +honor `AgentTask(deadline=...)` too. + +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. + +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..600ea44 --- /dev/null +++ b/src/agent_runtime_kit/_control.py @@ -0,0 +1,176 @@ +"""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, task_started_event + +_T = TypeVar("_T") + + +@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._lock = asyncio.Lock() + + 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") + async with self._lock: + 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: + async with self._lock: + 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) -> CancellationReceipt: + """Request cancellation of an active coroutine without blocking on it.""" + + async with self._lock: + active = self._active.get(task_id) + if active is None: + 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: + # The adapter body normally emits started. It is deliberately not + # invoked for an already-expired task, so preserve a coherent pair. + await safe_emit(task, task_started_event(task, self.kind)) + await self._emit_timed_out(task, deadline) + raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) + + operation_task: asyncio.Future[_T] = asyncio.ensure_future(operation()) + try: + done, _ = await asyncio.wait( + {operation_task}, + timeout=remaining, + return_when=asyncio.FIRST_COMPLETED, + ) + if operation_task in done: + return await operation_task + + operation_task.cancel() + await asyncio.gather(operation_task, return_exceptions=True) + await self._emit_timed_out(task, deadline) + raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) + except asyncio.CancelledError: + operation_task.cancel() + await asyncio.gather(operation_task, return_exceptions=True) + 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, + ) + + +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) 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..3720d19 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,13 @@ def __repr__(self) -> str: _UNSET = _UnsetType() + +@dataclass +class _KitActiveRun: + agent: AgentRuntime + task: asyncio.Task[Any] + cancellation_requested: 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 +107,8 @@ 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] = {} + self._active_lock = asyncio.Lock() @property def registry(self) -> RuntimeRegistry: @@ -193,6 +208,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 +237,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 +265,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 +294,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 +317,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 +365,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 +387,140 @@ 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") + async with self._active_lock: + 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: + async with self._active_lock: + active = self._active_runs.get(key) + if active is not None and active.task is current: + 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) + async with self._active_lock: + 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 + + cancel_method = getattr(agent, "cancel", None) + if callable(cancel_method): + try: + receipt = await cancel_method(task_id) + except Exception as exc: + active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.FAILED, + message=( + "runtime cancellation hook failed; local cancellation " + f"requested: {exc}" + ), + ) + if isinstance(receipt, CancellationReceipt): + if receipt.disposition in { + CancellationDisposition.NOT_ACTIVE, + CancellationDisposition.UNSUPPORTED, + }: + active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=CancellationDisposition.REQUESTED, + message=( + "runtime hook could not find the task; local coroutine " + "cancellation was requested" + ), + ) + return receipt + accepted = active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=( + CancellationDisposition.LEGACY_UNCONFIRMED + if accepted + else CancellationDisposition.FAILED + ), + message=( + "legacy runtime returned no cancellation receipt; local coroutine " + "cancellation was requested" + if accepted + else "the active asyncio task had already completed" + ), + ) + + accepted = active.task.cancel() + return CancellationReceipt( + task_id=task_id, + kind=kind, + disposition=( + CancellationDisposition.REQUESTED + if accepted + else CancellationDisposition.FAILED + ), + message=( + "runtime has no cancellation hook; local coroutine cancellation was requested" + if accepted + else "the active asyncio task had already completed" + ), + ) async def aclose(self) -> None: """Close every runtime this hub constructed and cached.""" @@ -401,6 +566,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 +576,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 +636,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..0ccd242 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,10 @@ 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 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 +865,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/antigravity.py b/src/agent_runtime_kit/adapters/antigravity.py index 92efbf6..aa7ad17 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, @@ -72,7 +74,7 @@ class AntigravityAgentRuntime: structured_output=True, streaming=True, tool_audit=True, - cancellation=False, + cancellation=True, tool_filters=True, ) @@ -113,6 +115,7 @@ 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) async def __aenter__(self) -> AntigravityAgentRuntime: return self @@ -227,6 +230,11 @@ 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.""" + + 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 +277,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. diff --git a/src/agent_runtime_kit/adapters/claude.py b/src/agent_runtime_kit/adapters/claude.py index 4273049..ee4d3a1 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, @@ -69,7 +71,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 +104,7 @@ 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) async def __aenter__(self) -> ClaudeAgentRuntime: return self @@ -164,6 +167,11 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: ) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled Claude task.""" + + 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 +227,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. diff --git a/src/agent_runtime_kit/adapters/codex.py b/src/agent_runtime_kit/adapters/codex.py index b14ede6..73a5eca 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, @@ -74,7 +76,7 @@ class CodexAgentRuntime: structured_output=True, streaming=False, tool_audit=True, - cancellation=False, + cancellation=True, reasoning_effort=True, ) @@ -111,6 +113,7 @@ 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) async def __aenter__(self) -> CodexAgentRuntime: return self @@ -192,6 +195,11 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: ) async def run(self, task: AgentTask) -> AgentResult: + """Execute one deadline- and cancellation-controlled Codex task.""" + + 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 +247,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. diff --git a/tests/test_control.py b/tests/test_control.py new file mode 100644 index 0000000..d32828a --- /dev/null +++ b/tests/test_control.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import asyncio +from datetime import datetime, timedelta, timezone + +import pytest + +from agent_runtime_kit import ( + AgentKit, + AgentResult, + AgentRuntimeKind, + AgentTask, + AgentTaskTimeoutError, + CancellationDisposition, + CancellationReceipt, + FakeAgentRuntime, + FinishReason, + RuntimeAvailability, + RuntimeRegistry, +) +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) + + +def test_deadline_requires_timezone_and_normalizes_to_utc() -> None: + 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.started", + "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 + + +@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 From c6d7455e08c9cc3af612161dc491d935fb2ff1b6 Mon Sep 17 00:00:00 2001 From: Eloi Date: Fri, 10 Jul 2026 14:25:10 +0200 Subject: [PATCH 2/4] fix: harden cancellation races --- CHANGELOG.md | 3 + docs/task-control.md | 4 + src/agent_runtime_kit/_control.py | 55 ++-- src/agent_runtime_kit/_kit.py | 190 ++++++++---- src/agent_runtime_kit/_types.py | 2 + src/agent_runtime_kit/adapters/_common.py | 31 +- src/agent_runtime_kit/adapters/antigravity.py | 11 +- src/agent_runtime_kit/adapters/claude.py | 13 +- src/agent_runtime_kit/adapters/codex.py | 11 +- tests/test_control.py | 270 ++++++++++++++++++ 10 files changed, 487 insertions(+), 103 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e1b4a7e..001a77c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -75,6 +75,9 @@ 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 + finishes reusable provider-process teardown under repeated cancellation. ## 0.4.0 - 2026-07-02 diff --git a/docs/task-control.md b/docs/task-control.md index 77c319d..8effb3a 100644 --- a/docs/task-control.md +++ b/docs/task-control.md @@ -47,6 +47,10 @@ 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 diff --git a/src/agent_runtime_kit/_control.py b/src/agent_runtime_kit/_control.py index 600ea44..3348bd6 100644 --- a/src/agent_runtime_kit/_control.py +++ b/src/agent_runtime_kit/_control.py @@ -38,7 +38,6 @@ class RuntimeTaskController: def __init__(self, kind: AgentRuntimeKind | str) -> None: self.kind = AgentRuntimeKind.coerce(kind) self._active: dict[str, _ActiveRun] = {} - self._lock = asyncio.Lock() async def run( self, @@ -50,33 +49,40 @@ async def run( 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") - async with self._lock: - 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) + # 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: - async with self._lock: - active = self._active.get(task.task_id) - if active is not None and active.task is current: - del self._active[task.task_id] + 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) -> CancellationReceipt: + 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.""" - async with self._lock: - active = self._active.get(task_id) - if active is None: - 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() + 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, @@ -107,6 +113,7 @@ async def _run_with_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}, @@ -118,12 +125,16 @@ async def _run_with_deadline( operation_task.cancel() await asyncio.gather(operation_task, return_exceptions=True) + # Transition before emitting: cancellation of a slow sink must never + # turn one timeout into a second, contradictory interrupted terminal. + timed_out = True await self._emit_timed_out(task, deadline) raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) except asyncio.CancelledError: operation_task.cancel() await asyncio.gather(operation_task, return_exceptions=True) - await self._emit_interrupted(task) + if not timed_out: + await self._emit_interrupted(task) raise async def _emit_interrupted(self, task: AgentTask) -> None: diff --git a/src/agent_runtime_kit/_kit.py b/src/agent_runtime_kit/_kit.py index 3720d19..9fa6649 100644 --- a/src/agent_runtime_kit/_kit.py +++ b/src/agent_runtime_kit/_kit.py @@ -63,6 +63,8 @@ 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). @@ -108,7 +110,6 @@ def __init__( self._handlers: list[tuple[str, _EventHandler]] = [] self._cache_lock = asyncio.Lock() self._active_runs: dict[tuple[int, str], _KitActiveRun] = {} - self._active_lock = asyncio.Lock() @property def registry(self) -> RuntimeRegistry: @@ -391,12 +392,12 @@ async def run( 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") - async with self._active_lock: - 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) + # 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) @@ -410,9 +411,13 @@ async def run( return result return _parse_result(output_type, result) finally: - async with self._active_lock: - active = self._active_runs.get(key) - if active is not None and active.task is current: + 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( @@ -442,86 +447,147 @@ async def cancel( kind = agent.kind key = (id(agent), task_id) - async with self._active_lock: - 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 + 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 - cancel_method = getattr(agent, "cancel", None) - if callable(cancel_method): + controller = getattr(agent, "_task_controller", None) + if isinstance(controller, RuntimeTaskController): try: - receipt = await cancel_method(task_id) - except Exception as exc: - active.task.cancel() - return CancellationReceipt( - task_id=task_id, - kind=kind, - disposition=CancellationDisposition.FAILED, - message=( - "runtime cancellation hook failed; local cancellation " - f"requested: {exc}" - ), - ) - if isinstance(receipt, CancellationReceipt): + # 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, }: - active.task.cancel() + accepted = active.task.cancel() return CancellationReceipt( task_id=task_id, kind=kind, - disposition=CancellationDisposition.REQUESTED, + disposition=( + CancellationDisposition.REQUESTED + if accepted + else CancellationDisposition.FAILED + ), message=( - "runtime hook could not find the task; local coroutine " - "cancellation was requested" + "runtime controller no longer owned this run; local " + "coroutine cancellation was requested" + if accepted + else "the active asyncio task had already completed" ), ) return receipt - accepted = active.task.cancel() + 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.LEGACY_UNCONFIRMED - if accepted - else CancellationDisposition.FAILED - ), + 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=( - "legacy runtime returned no cancellation receipt; local coroutine " + "runtime has no cancellation hook; local coroutine " "cancellation was requested" - if accepted - else "the active asyncio task had already completed" ), ) - accepted = active.task.cancel() + 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.REQUESTED - if accepted - else CancellationDisposition.FAILED - ), + disposition=CancellationDisposition.LEGACY_UNCONFIRMED, message=( - "runtime has no cancellation hook; local coroutine cancellation was requested" - if accepted - else "the active asyncio task had already completed" + "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.""" diff --git a/src/agent_runtime_kit/_types.py b/src/agent_runtime_kit/_types.py index 0ccd242..7a55066 100644 --- a/src/agent_runtime_kit/_types.py +++ b/src/agent_runtime_kit/_types.py @@ -743,6 +743,8 @@ def __post_init__(self) -> None: 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)) diff --git a/src/agent_runtime_kit/adapters/_common.py b/src/agent_runtime_kit/adapters/_common.py index f49e2e7..87e4b18 100644 --- a/src/agent_runtime_kit/adapters/_common.py +++ b/src/agent_runtime_kit/adapters/_common.py @@ -2,9 +2,10 @@ 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 @@ -374,3 +375,31 @@ async def close_vendor_resource( result = close() if hasattr(result, "__await__"): await result + + +async def finish_vendor_cleanup(operation: Awaitable[Any]) -> 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 + original exception remains authoritative; a cleanup failure is returned for + logging instead of masking it. + """ + + cleanup: asyncio.Future[Any] = asyncio.ensure_future(operation) + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + # shield keeps the cleanup task alive. Consume any number of + # additional cancellation requests until teardown has settled. + continue + 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 diff --git a/src/agent_runtime_kit/adapters/antigravity.py b/src/agent_runtime_kit/adapters/antigravity.py index aa7ad17..6e71d04 100644 --- a/src/agent_runtime_kit/adapters/antigravity.py +++ b/src/agent_runtime_kit/adapters/antigravity.py @@ -37,6 +37,7 @@ empty_completion_error, filter_supported_kwargs, fingerprint_item, + finish_vendor_cleanup, model_support_issue, optional_int, optional_str, @@ -416,9 +417,8 @@ 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()) + if close_exc is not None: logger.warning( "failed to close Antigravity agent after run failure: %s", close_exc, @@ -562,9 +562,8 @@ 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)) + 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 ee4d3a1..babb68f 100644 --- a/src/agent_runtime_kit/adapters/claude.py +++ b/src/agent_runtime_kit/adapters/claude.py @@ -36,6 +36,7 @@ field_value, filter_supported_kwargs, fingerprint_value, + finish_vendor_cleanup, metadata_str, model_support_issue, optional_int, @@ -299,9 +300,8 @@ 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()) + if close_exc is not None: logger.warning( "failed to close Claude client after run failure: %s", close_exc ) @@ -331,9 +331,10 @@ 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) + ) + 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 73a5eca..e92c1c5 100644 --- a/src/agent_runtime_kit/adapters/codex.py +++ b/src/agent_runtime_kit/adapters/codex.py @@ -31,6 +31,7 @@ empty_completion_error, field_value, filter_supported_kwargs, + finish_vendor_cleanup, metadata_str, model_support_issue, optional_int, @@ -367,9 +368,8 @@ 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()) + if close_exc is not None: logger.warning( "failed to close Codex app-server after run failure: %s", close_exc ) @@ -455,9 +455,8 @@ 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)) + 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 index d32828a..49bce95 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -18,6 +18,7 @@ RuntimeAvailability, RuntimeRegistry, ) +from agent_runtime_kit.adapters._common import finish_vendor_cleanup from agent_runtime_kit.testing import RecordingEventSink @@ -99,7 +100,125 @@ async def _run_task(self, task: AgentTask) -> AgentResult: 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 + + 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)) @@ -241,6 +360,11 @@ async def test_direct_builtin_cancel_returns_receipts_and_emits_interrupted() -> 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: @@ -281,3 +405,149 @@ async def test_external_task_cancellation_emits_one_interrupted_terminal() -> No 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_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_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 From 225666e71f53917a88336e023a2dfd8ea06cfb48 Mon Sep 17 00:00:00 2001 From: Eloi Date: Fri, 10 Jul 2026 14:33:41 +0200 Subject: [PATCH 3/4] fix: bound cancellation cleanup --- docs/task-control.md | 7 +-- src/agent_runtime_kit/_control.py | 46 +++++++++++++++---- src/agent_runtime_kit/adapters/_common.py | 48 +++++++++++++++++-- tests/test_control.py | 56 +++++++++++++++++++++-- 4 files changed, 137 insertions(+), 20 deletions(-) diff --git a/docs/task-control.md b/docs/task-control.md index 8effb3a..abb307e 100644 --- a/docs/task-control.md +++ b/docs/task-control.md @@ -15,9 +15,10 @@ except AgentTaskTimeoutError as exc: print(exc.task_id, exc.deadline) ``` -An expired deadline never starts the vendor SDK. A deadline expiry emits an -`agent.task.failed` event with `finish_reason="timed_out"`, cancels the in-flight -provider coroutine, waits for its cleanup, and raises `AgentTaskTimeoutError` +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. diff --git a/src/agent_runtime_kit/_control.py b/src/agent_runtime_kit/_control.py index 3348bd6..b5e16ba 100644 --- a/src/agent_runtime_kit/_control.py +++ b/src/agent_runtime_kit/_control.py @@ -21,9 +21,10 @@ CancellationReceipt, FinishReason, ) -from agent_runtime_kit.events import safe_emit, task_failed_event, task_started_event +from agent_runtime_kit.events import safe_emit, task_failed_event _T = TypeVar("_T") +_OPERATION_CLEANUP_TIMEOUT = 5.0 @dataclass @@ -106,9 +107,9 @@ async def _run_with_deadline( remaining = (deadline - datetime.now(tz=timezone.utc)).total_seconds() if remaining <= 0: - # The adapter body normally emits started. It is deliberately not - # invoked for an already-expired task, so preserve a coherent pair. - await safe_emit(task, task_started_event(task, self.kind)) + # 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) @@ -123,16 +124,14 @@ async def _run_with_deadline( if operation_task in done: return await operation_task - operation_task.cancel() - await asyncio.gather(operation_task, return_exceptions=True) # Transition before emitting: cancellation of a slow sink must never # turn one timeout into a second, contradictory interrupted terminal. timed_out = True + await _cancel_and_settle(operation_task) await self._emit_timed_out(task, deadline) raise AgentTaskTimeoutError(self.kind, task.task_id, deadline) except asyncio.CancelledError: - operation_task.cancel() - await asyncio.gather(operation_task, return_exceptions=True) + await _cancel_and_settle(operation_task) if not timed_out: await self._emit_interrupted(task) raise @@ -185,3 +184,34 @@ async def run_legacy_with_deadline( # terminal event semantics. controller = RuntimeTaskController(kind) return await controller.run(task, operation) + + +async def _cancel_and_settle(operation: asyncio.Future[Any]) -> None: + """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() + operation.add_done_callback(_consume_operation_result) + return + try: + await asyncio.wait_for(asyncio.shield(operation), timeout=remaining) + except asyncio.CancelledError: + continue + except asyncio.TimeoutError: + operation.cancel() + operation.add_done_callback(_consume_operation_result) + return + except Exception: + break + _consume_operation_result(operation) + + +def _consume_operation_result(operation: asyncio.Future[Any]) -> None: + try: + operation.exception() + except BaseException: + return diff --git a/src/agent_runtime_kit/adapters/_common.py b/src/agent_runtime_kit/adapters/_common.py index 87e4b18..315e4ad 100644 --- a/src/agent_runtime_kit/adapters/_common.py +++ b/src/agent_runtime_kit/adapters/_common.py @@ -29,6 +29,7 @@ from agent_runtime_kit.compatibility import compatibility_for ModelSource = Literal["task", "metadata", "constructor", "provider-native"] +DEFAULT_VENDOR_CLEANUP_TIMEOUT = 5.0 @dataclass(frozen=True) @@ -377,24 +378,52 @@ async def close_vendor_resource( await result -async def finish_vendor_cleanup(operation: Awaitable[Any]) -> BaseException | None: +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 - original exception remains authoritative; a cleanup failure is returned for - logging instead of masking it. + 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 TimeoutError( + f"vendor cleanup did not finish within {float(timeout):g} seconds" + ) try: - await asyncio.shield(cleanup) + 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 has settled. + # 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 TimeoutError( + f"vendor cleanup did not finish within {float(timeout):g} seconds" + ) except Exception: # The failure is retrieved and returned below so it can be logged # without replacing the run's original exception. @@ -403,3 +432,12 @@ async def finish_vendor_cleanup(operation: Awaitable[Any]) -> BaseException | No 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/tests/test_control.py b/tests/test_control.py index 49bce95..e65b2b0 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -269,10 +269,7 @@ async def test_expired_deadline_never_invokes_legacy_runtime() -> None: 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.started", - "agent.task.failed", - ] + assert [event["name"] for event in sink.events] == ["agent.task.failed"] assert sink.events[-1]["attributes"]["finish_reason"] == FinishReason.TIMED_OUT @@ -486,6 +483,33 @@ async def test_timeout_racing_external_cancel_emits_only_timed_out_terminal() -> 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() @@ -540,6 +564,30 @@ async def owner() -> None: 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) + release_cleanup.set() + await asyncio.sleep(0) + + @pytest.mark.asyncio async def test_cancel_after_completed_run_is_truthfully_not_active() -> None: runtime = LegacyBlockingRuntime() From dbca5cd6c6e2154f9d337c34cf60e50b5644bc96 Mon Sep 17 00:00:00 2001 From: Eloi Date: Fri, 10 Jul 2026 14:38:56 +0200 Subject: [PATCH 4/4] fix: quarantine pending teardown --- CHANGELOG.md | 4 +- docs/task-control.md | 4 +- src/agent_runtime_kit/_control.py | 31 ++++++--- src/agent_runtime_kit/adapters/_common.py | 49 ++++++++++++--- src/agent_runtime_kit/adapters/antigravity.py | 6 ++ src/agent_runtime_kit/adapters/claude.py | 6 ++ src/agent_runtime_kit/adapters/codex.py | 6 ++ tests/test_control.py | 63 ++++++++++++++++++- 8 files changed, 152 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 001a77c..9715586 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -77,7 +77,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 - finishes reusable provider-process teardown under repeated cancellation. + 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/docs/task-control.md b/docs/task-control.md index abb307e..a9fca02 100644 --- a/docs/task-control.md +++ b/docs/task-control.md @@ -20,7 +20,9 @@ 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. +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: diff --git a/src/agent_runtime_kit/_control.py b/src/agent_runtime_kit/_control.py index b5e16ba..45036ac 100644 --- a/src/agent_runtime_kit/_control.py +++ b/src/agent_runtime_kit/_control.py @@ -39,6 +39,7 @@ class RuntimeTaskController: 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, @@ -50,6 +51,12 @@ async def run( 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. @@ -127,11 +134,13 @@ async def _run_with_deadline( # Transition before emitting: cancellation of a slow sink must never # turn one timeout into a second, contradictory interrupted terminal. timed_out = True - await _cancel_and_settle(operation_task) + 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: - await _cancel_and_settle(operation_task) + if not await _cancel_and_settle(operation_task): + self._quarantine(operation_task) if not timed_out: await self._emit_interrupted(task) raise @@ -171,6 +180,15 @@ def _receipt( 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, @@ -186,7 +204,7 @@ async def run_legacy_with_deadline( return await controller.run(task, operation) -async def _cancel_and_settle(operation: asyncio.Future[Any]) -> None: +async def _cancel_and_settle(operation: asyncio.Future[Any]) -> bool: """Give a cancelled operation bounded, repeat-cancel-safe cleanup time.""" operation.cancel() @@ -195,19 +213,18 @@ async def _cancel_and_settle(operation: asyncio.Future[Any]) -> None: remaining = deadline - asyncio.get_running_loop().time() if remaining <= 0: operation.cancel() - operation.add_done_callback(_consume_operation_result) - return + return False try: await asyncio.wait_for(asyncio.shield(operation), timeout=remaining) except asyncio.CancelledError: continue except asyncio.TimeoutError: operation.cancel() - operation.add_done_callback(_consume_operation_result) - return + return False except Exception: break _consume_operation_result(operation) + return True def _consume_operation_result(operation: asyncio.Future[Any]) -> None: diff --git a/src/agent_runtime_kit/adapters/_common.py b/src/agent_runtime_kit/adapters/_common.py index 315e4ad..841a0d1 100644 --- a/src/agent_runtime_kit/adapters/_common.py +++ b/src/agent_runtime_kit/adapters/_common.py @@ -11,7 +11,7 @@ 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, ) @@ -32,6 +32,45 @@ 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) class ModelSelection: """Effective model value and the layer that selected it.""" @@ -408,9 +447,7 @@ async def finish_vendor_cleanup( if remaining <= 0: cleanup.cancel() cleanup.add_done_callback(_consume_cleanup_result) - return TimeoutError( - f"vendor cleanup did not finish within {float(timeout):g} seconds" - ) + return VendorCleanupTimeoutError(float(timeout), cleanup) try: await asyncio.wait_for(asyncio.shield(cleanup), timeout=remaining) except asyncio.CancelledError: @@ -421,9 +458,7 @@ async def finish_vendor_cleanup( except asyncio.TimeoutError: cleanup.cancel() cleanup.add_done_callback(_consume_cleanup_result) - return TimeoutError( - f"vendor cleanup did not finish within {float(timeout):g} seconds" - ) + 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. diff --git a/src/agent_runtime_kit/adapters/antigravity.py b/src/agent_runtime_kit/adapters/antigravity.py index 6e71d04..e470f7e 100644 --- a/src/agent_runtime_kit/adapters/antigravity.py +++ b/src/agent_runtime_kit/adapters/antigravity.py @@ -33,6 +33,7 @@ Usage, ) from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, filter_supported_kwargs, @@ -117,6 +118,7 @@ def __init__( 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 @@ -233,6 +235,7 @@ def validate_task(self, task: AgentTask) -> TaskSupportReport: 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: @@ -290,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() @@ -418,6 +422,7 @@ async def _invoke( # already held, so close under the agent lock only and never # let cleanup mask the original error. 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", @@ -563,6 +568,7 @@ async def _persistent_agent( agent = await enter() if callable(enter) else context except BaseException: 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 babb68f..c13f51d 100644 --- a/src/agent_runtime_kit/adapters/claude.py +++ b/src/agent_runtime_kit/adapters/claude.py @@ -31,6 +31,7 @@ ) from agent_runtime_kit.adapters._common import ( STRUCTURED_OUTPUT_MISSING, + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, field_value, @@ -106,6 +107,7 @@ def __init__( 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 @@ -170,6 +172,7 @@ 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: @@ -240,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() @@ -301,6 +305,7 @@ async def _run_with_client( # lock is already held, so close under the client lock only and # never let cleanup mask the original error. 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 @@ -334,6 +339,7 @@ async def _persistent_client( 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 e92c1c5..5d0ffc3 100644 --- a/src/agent_runtime_kit/adapters/codex.py +++ b/src/agent_runtime_kit/adapters/codex.py @@ -27,6 +27,7 @@ Usage, ) from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, close_vendor_resource, empty_completion_error, field_value, @@ -115,6 +116,7 @@ def __init__( 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 @@ -198,6 +200,7 @@ 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: @@ -260,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() @@ -369,6 +373,7 @@ async def _run_codex( # doomed client between failure and eviction; close under # the client lock only and never mask the original error. 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 @@ -456,6 +461,7 @@ async def _persistent_codex_client( client = await enter() if callable(enter) else context except BaseException: 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 index e65b2b0..51f2369 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -5,9 +5,11 @@ import pytest +import agent_runtime_kit._control as control_module from agent_runtime_kit import ( AgentKit, AgentResult, + AgentRuntimeError, AgentRuntimeKind, AgentTask, AgentTaskTimeoutError, @@ -18,7 +20,10 @@ RuntimeAvailability, RuntimeRegistry, ) -from agent_runtime_kit.adapters._common import finish_vendor_cleanup +from agent_runtime_kit.adapters._common import ( + VendorCleanupQuarantine, + finish_vendor_cleanup, +) from agent_runtime_kit.testing import RecordingEventSink @@ -215,6 +220,29 @@ async def _run_task(self, task: AgentTask) -> AgentResult: 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] @@ -584,8 +612,41 @@ async def wedged_cleanup() -> None: 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