From 9d0904fa8d81de09d1dede891c66b82ded5986c0 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Fri, 23 Oct 2026 16:26:43 +0000 Subject: [PATCH] feat(server): recover tasks abandoned by a crashed replica In cluster mode, a replica that dies mid-turn leaves its task WORKING in the shared store, and no other replica ever finishes it. - A turn now holds a lease on its task (TaskLeaseStore, with a database implementation and migration). The running replica renews it on a heartbeat and releases it when the turn ends, including when the task waits for input. - fail_abandoned_tasks / AbandonedTaskSweeper mark tasks FAILED once their lease expires. Every replica can run the sweeper; each expired lease is claimed by one of them. A task already waiting for input recorded the end of its turn, so only its lease is dropped. - A follow-up arriving while another replica holds the lease is rejected with UnsupportedOperationError instead of starting a second executor. If the task already waits for input, the other replica is about to release the lease, so the follow-up waits for it for up to one heartbeat interval. - A replica that loses its lease stops its executor, and its late events are dropped. - The local task a replica created for a rejected follow-up is released as soon as it is idle, so a later subscription on that replica follows the shared event stream instead of a task that never runs there. - DefaultRequestHandler.aclose() marks turns still running on this replica FAILED. On by default when an event stream is configured. The lease store is optional; without it, only aclose() behaves differently. The clustermode sample wires it up and documents how to try a crash. Closes #1324 --- samples/clustermode/README.md | 25 + samples/clustermode/cluster_common.py | 14 +- samples/clustermode/server.py | 42 +- src/a2a/migrations/README.md | 2 + .../versions/b0f01b04ff26_add_task_leases.py | 71 ++ src/a2a/server/agent_execution/active_task.py | 195 ++++- .../agent_execution/active_task_registry.py | 21 +- src/a2a/server/cluster/__init__.py | 14 + .../server/cluster/database_lease_store.py | 174 ++++ src/a2a/server/cluster/lease.py | 56 ++ src/a2a/server/cluster/recovery.py | 218 +++++ src/a2a/server/models.py | 37 + .../default_request_handler_v2.py | 59 +- tests/integration/test_multiserver.py | 139 ++++ .../versions/test_migration_b0f01b04ff26.py | 150 ++++ tests/server/cluster/conftest.py | 102 ++- .../cluster/test_database_lease_store.py | 247 ++++++ tests/server/cluster/test_handler_wiring.py | 97 +++ .../server/cluster/test_lease_multireplica.py | 742 ++++++++++++++++++ tests/server/cluster/test_recovery.py | 300 +++++++ 20 files changed, 2677 insertions(+), 28 deletions(-) create mode 100644 src/a2a/migrations/versions/b0f01b04ff26_add_task_leases.py create mode 100644 src/a2a/server/cluster/database_lease_store.py create mode 100644 src/a2a/server/cluster/lease.py create mode 100644 src/a2a/server/cluster/recovery.py create mode 100644 tests/migrations/versions/test_migration_b0f01b04ff26.py create mode 100644 tests/server/cluster/test_database_lease_store.py create mode 100644 tests/server/cluster/test_lease_multireplica.py create mode 100644 tests/server/cluster/test_recovery.py diff --git a/samples/clustermode/README.md b/samples/clustermode/README.md index 1ec16a567..a7b3d5b8b 100644 --- a/samples/clustermode/README.md +++ b/samples/clustermode/README.md @@ -46,6 +46,27 @@ Expected output ends with: OK: one task was started, observed, and cancelled across three different replicas. ``` +## Recovering from a crashed replica + +Each turn holds a lease in the shared `task_leases` table, renewed every +`--heartbeat-interval` seconds (2 s here). If a replica dies mid-turn, its lease +expires after `--lease-ttl` seconds (10 s here) and the `AbandonedTaskSweeper` +on a surviving replica marks the task FAILED. A follow-up sent while another +replica holds the lease is rejected instead of starting a second executor. + +To see it, start the cluster with slower ticks, send a message to one replica, +then kill that replica's process: + +```bash +python -m samples.clustermode.run_cluster --replicas 2 --ticks 60 +pkill -9 -f 'samples.clustermode.server.*--port 41241' +``` + +A `GetTask` against port 41242 shows the task in WORKING until the lease +expires, then FAILED with the status message "The replica running this task +stopped responding." When a replica is stopped normally, `handler.aclose()` +fails its running turns right away. + ## Files | File | Purpose | @@ -66,9 +87,13 @@ handler = DefaultRequestHandler( task_store=VersionedDatabaseTaskStore(engine=..., create_table=False), agent_card=agent_card, event_stream=DatabaseTaskEventStream(engine=..., create_table=False), + lease_store=DatabaseTaskLeaseStore(engine=..., create_table=False), ) ``` +The lease store is optional; without it, a task whose replica crashes mid-turn +stays WORKING. + Omit `event_stream` (and use a plain `InMemoryTaskStore`) and you get the ordinary single-process behaviour — the multi-server path is fully opt-in. ``` diff --git a/samples/clustermode/cluster_common.py b/samples/clustermode/cluster_common.py index 6d13093bc..58437a2e8 100644 --- a/samples/clustermode/cluster_common.py +++ b/samples/clustermode/cluster_common.py @@ -24,7 +24,10 @@ from a2a.helpers.proto_helpers import new_task_from_user_message from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.agent_execution.context import RequestContext -from a2a.server.cluster import VersionedDatabaseTaskStore +from a2a.server.cluster import ( + DatabaseTaskLeaseStore, + VersionedDatabaseTaskStore, +) from a2a.server.cluster.database_event_stream import DatabaseTaskEventStream from a2a.server.events.event_queue import EventQueue from a2a.server.models import Base @@ -60,7 +63,7 @@ def default_dsn() -> str: async def init_schema(dsn: str) -> None: - """Creates the shared tables (task + task_events) once, up front.""" + """Creates the shared tables (tasks, task_events, task_leases) once.""" engine = create_async_engine(dsn) try: async with engine.begin() as conn: @@ -86,6 +89,13 @@ def build_cluster_backends( return store, stream +def build_lease_store(dsn: str) -> DatabaseTaskLeaseStore: + """Builds the shared turn-lease store used to recover crashed turns.""" + return DatabaseTaskLeaseStore( + engine=create_async_engine(dsn), create_table=False + ) + + def build_agent_card(base_url: str) -> AgentCard: """The agent card advertised by every replica (same logical agent).""" return AgentCard( diff --git a/samples/clustermode/server.py b/samples/clustermode/server.py index 8a2d2e0bc..66c66719c 100644 --- a/samples/clustermode/server.py +++ b/samples/clustermode/server.py @@ -14,6 +14,7 @@ from fastapi import FastAPI +from a2a.server.cluster.recovery import AbandonedTaskSweeper from a2a.server.request_handlers import DefaultRequestHandler from a2a.server.routes import ( add_a2a_routes_to_fastapi, @@ -24,6 +25,7 @@ SlowEchoAgent, build_agent_card, build_cluster_backends, + build_lease_store, default_dsn, init_schema, ) @@ -32,7 +34,14 @@ logger = logging.getLogger(__name__) -async def serve(host: str, port: int, replica_id: str, ticks: int) -> None: +async def serve( # noqa: PLR0913 + host: str, + port: int, + replica_id: str, + ticks: int, + lease_ttl: float, + heartbeat_interval: float, +) -> None: """Runs one replica bound to a shared store + event stream.""" dsn = default_dsn() await init_schema(dsn) @@ -42,12 +51,21 @@ async def serve(host: str, port: int, replica_id: str, ticks: int) -> None: store, stream = build_cluster_backends(dsn) await store.initialize() await stream.initialize() + leases = build_lease_store(dsn) + await leases.initialize() handler = DefaultRequestHandler( agent_executor=SlowEchoAgent(replica_id, ticks=ticks), task_store=store, agent_card=agent_card, event_stream=stream, + lease_store=leases, + lease_ttl=lease_ttl, + heartbeat_interval=heartbeat_interval, + ) + # Fails tasks whose replica died mid-turn, once their lease expires. + sweeper = AbandonedTaskSweeper( + leases, store, interval=heartbeat_interval, event_stream=stream ) jsonrpc_routes = create_jsonrpc_routes( @@ -67,7 +85,14 @@ async def serve(host: str, port: int, replica_id: str, ticks: int) -> None: 'Replica %s listening on %s (dsn=%s)', replica_id, base_url, dsn ) config = uvicorn.Config(app, host=host, port=port, log_level='warning') - await uvicorn.Server(config).serve() + sweeper.start() + try: + await uvicorn.Server(config).serve() + finally: + await sweeper.aclose() + # Records FAILED for turns still running here, instead of leaving + # them for the sweeper on another replica. + await handler.aclose() def main() -> None: @@ -78,9 +103,20 @@ def main() -> None: parser.add_argument('--port', type=int, default=41241) parser.add_argument('--replica-id', default='replica-0') parser.add_argument('--ticks', type=int, default=10) + parser.add_argument('--lease-ttl', type=float, default=10.0) + parser.add_argument('--heartbeat-interval', type=float, default=2.0) args = parser.parse_args() with contextlib.suppress(KeyboardInterrupt): - asyncio.run(serve(args.host, args.port, args.replica_id, args.ticks)) + asyncio.run( + serve( + args.host, + args.port, + args.replica_id, + args.ticks, + args.lease_ttl, + args.heartbeat_interval, + ) + ) if __name__ == '__main__': diff --git a/src/a2a/migrations/README.md b/src/a2a/migrations/README.md index 4767363ff..499f9f6af 100644 --- a/src/a2a/migrations/README.md +++ b/src/a2a/migrations/README.md @@ -40,6 +40,8 @@ Always run this command after installing or upgrading the SDK to ensure your dat Revision `b5e3d1c8a2f7` creates two additive tables, `task_versions` and `task_events`, used only by the clustered multi-server stores (`a2a.server.cluster`). It leaves the `tasks` table untouched, so deployments that use the default `DatabaseTaskStore` need not run it and are unaffected. Run it only when adopting a `VersionedTaskStore` for multi-server deployment. +Revision `b0f01b04ff26` creates the additive `task_leases` table, used only by `DatabaseTaskLeaseStore` (`a2a.server.cluster`) to recover tasks whose replica stopped while running them. Run it only when adopting a lease store. + ```bash uv run a2a-db ``` diff --git a/src/a2a/migrations/versions/b0f01b04ff26_add_task_leases.py b/src/a2a/migrations/versions/b0f01b04ff26_add_task_leases.py new file mode 100644 index 000000000..c28eb1af2 --- /dev/null +++ b/src/a2a/migrations/versions/b0f01b04ff26_add_task_leases.py @@ -0,0 +1,71 @@ +"""add task_leases table + +Revision ID: b0f01b04ff26 +Revises: b5e3d1c8a2f7 +Create Date: 2026-10-23 16:00:00.000000 + +""" + +import logging + +from collections.abc import Sequence +from typing import Union + +import sqlalchemy as sa + + +try: + from alembic import context, op +except ImportError as e: + raise ImportError( + "A2A migrations require the 'db-cli' extra. Install with: 'pip install a2a-sdk[db-cli]'." + ) from e + +from a2a.migrations.migration_utils import table_exists + + +# revision identifiers, used by Alembic. +revision: str = 'b0f01b04ff26' +down_revision: Union[str, Sequence[str], None] = 'b5e3d1c8a2f7' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema: create the task_leases table.""" + leases_table = context.config.get_main_option( + 'task_leases_table', 'task_leases' + ) + + if context.is_offline_mode() or not table_exists(leases_table): + op.create_table( + leases_table, + sa.Column('task_id', sa.String(36), primary_key=True), + sa.Column('owner', sa.String(255), nullable=True), + sa.Column('token', sa.String(32), nullable=False), + sa.Column('expires_at_ms', sa.BigInteger(), nullable=False), + ) + op.create_index( + f'ix_{leases_table}_expires_at_ms', + leases_table, + ['expires_at_ms'], + ) + else: + logging.info( + "Table '%s' already exists. Skipping creation.", leases_table + ) + + +def downgrade() -> None: + """Downgrade schema: drop the task_leases table.""" + leases_table = context.config.get_main_option( + 'task_leases_table', 'task_leases' + ) + + if context.is_offline_mode() or table_exists(leases_table): + op.drop_index( + f'ix_{leases_table}_expires_at_ms', table_name=leases_table + ) + op.drop_table(leases_table) + else: + logging.info("Table '%s' does not exist. Skipping drop.", leases_table) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 62b665d6b..2605bd074 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -51,6 +51,7 @@ from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.event_stream import TaskEventStream + from a2a.server.cluster.lease import TaskLease, TaskLeaseStore from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_sender import ( PushNotificationSender, @@ -58,6 +59,10 @@ from a2a.server.tasks.task_manager import TaskManager from a2a.server.cluster.event_stream import VersionedEvent +from a2a.server.cluster.lease import ( + DEFAULT_HEARTBEAT_INTERVAL, + DEFAULT_LEASE_TTL, +) from a2a.server.cluster.task_store import ConcurrentTaskModificationError from a2a.server.events.event_queue_v2 import ( AsyncQueue, @@ -96,6 +101,8 @@ TaskState.TASK_STATE_AUTH_REQUIRED, TaskState.TASK_STATE_INPUT_REQUIRED, } +# Seconds between attempts to take a lease that is about to be released. +_LEASE_RETRY_INTERVAL = 0.05 class _RequestStarted: @@ -109,6 +116,14 @@ def __init__(self, request_id: uuid.UUID): self.request_id = request_id +class _RequestRejected: + """A request the producer refused without running the executor.""" + + def __init__(self, request_id: uuid.UUID, error: Exception): + self.request_id = request_id + self.error = error + + async def send_push_notification( push_sender: PushNotificationSender | None, task_id: str, @@ -273,7 +288,7 @@ def _drop_fenced_event(self, event: Event) -> None: self.active_task._task_id, type(event).__name__, ) - if isinstance(event, _RequestCompleted): + if isinstance(event, _RequestCompleted | _RequestRejected): self.active_task._request_lock.release() self.active_task._event_queue_agent.task_done() @@ -291,12 +306,20 @@ async def _process_event_inner(self, event: Event) -> None: logger.debug( 'Consumer[%s]: Request completed', self.active_task._task_id ) + # The agent queue is FIFO, so every event of this turn has been + # saved by now: the turn has durably ended. + await self.active_task._release_lease() + self.active_task._request_lock.release() + self.active_task.schedule_idle_release() + elif isinstance(event, _RequestRejected): + self.active_task._running_elsewhere = True self.active_task._request_lock.release() self.active_task.schedule_idle_release() elif isinstance(event, _RequestStarted): logger.debug( 'Consumer[%s]: Request started', self.active_task._task_id ) + self.active_task._running_elsewhere = False self.message_to_save = event.request_context.message elif isinstance(event, BaseException): # Only the producer enqueues exceptions, when a request fails. @@ -503,6 +526,9 @@ def __init__( # noqa: PLR0913 on_cleanup: Callable[[ActiveTask], None] | None = None, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + lease_store: TaskLeaseStore | None = None, + lease_ttl: float = DEFAULT_LEASE_TTL, + heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL, ) -> None: """Initializes the ActiveTask. @@ -518,6 +544,11 @@ def __init__( # noqa: PLR0913 published to it so other replicas can observe them. interrupted_task_timeout: Optional seconds after which an idle interrupted task is finished; ``None`` disables it. + lease_store: Optional store of turn leases. When set, each turn + holds the task's lease, renewed every + `heartbeat_interval` seconds for `lease_ttl` seconds. + lease_ttl: Seconds a lease lasts without being renewed. + heartbeat_interval: Seconds between lease renewals. """ # --- Core Dependencies --- self._agent_executor = agent_executor @@ -533,6 +564,14 @@ def __init__( # noqa: PLR0913 self._interrupted_task_timeout = interrupted_task_timeout # Pending timer that releases a task parked in an interrupted state. self._idle_release_task: asyncio.Task[None] | None = None + # Set when the last request was rejected because another replica runs + # the task; this replica then has no use for it. + self._running_elsewhere = False + self._lease_store = lease_store + self._lease_ttl = lease_ttl + self._heartbeat_interval = heartbeat_interval + self._lease: TaskLease | None = None + self._heartbeat_task: asyncio.Task[None] | None = None # --- Synchronization Primitives --- # `_lock` protects structural lifecycle changes: start(), subscribe() counting, @@ -659,6 +698,99 @@ async def start( 'ActiveTask[%s]: Background tasks created', self._task_id ) + @property + def request_in_flight(self) -> bool: + """Whether this replica is running a turn for the task right now.""" + return self._request_lock.locked() + + @property + def call_context(self) -> ServerCallContext | None: + """The call context of the request most recently taken up.""" + return self._task_manager._call_context + + async def _acquire_lease(self, request_context: RequestContext) -> bool: + """Takes the task's turn lease; False if another replica holds it. + + A task waiting for input has ended its turn, and the replica that ran + it releases the lease just after saving that state. Such a lease is + waited for, up to one heartbeat interval, so that a follow-up sent as + soon as the previous turn returned is not rejected. + """ + if self._lease_store is None: + return True + task = request_context.current_task + wait = ( + self._heartbeat_interval + if task is not None and task.status.state in INTERRUPTED_TASK_STATES + else 0.0 + ) + deadline = asyncio.get_running_loop().time() + wait + while True: + lease = await self._lease_store.acquire( + self._task_id, request_context.call_context, self._lease_ttl + ) + if lease is not None: + break + if asyncio.get_running_loop().time() >= deadline: + return False + await asyncio.sleep(_LEASE_RETRY_INTERVAL) + self._lease = lease + self._heartbeat_task = asyncio.create_task( + self._heartbeat(self._lease_store, lease), + name=f'lease:{self._task_id}', + ) + return True + + async def _heartbeat( + self, lease_store: TaskLeaseStore, lease: TaskLease + ) -> None: + """Renews the lease, and stops the executor once the lease is lost.""" + while True: + await asyncio.sleep(self._heartbeat_interval) + try: + renewed = await lease_store.renew(lease, self._lease_ttl) + except Exception: + # Keep trying: if the store stays unreachable past the TTL, + # the next renewal that gets through reports the lease lost. + logger.exception('Lease[%s]: renewal failed', self._task_id) + continue + if not renewed: + logger.warning( + 'Lease[%s]: lease lost, stopping execution', self._task_id + ) + if self._producer_task and not self._producer_task.done(): + self._producer_task.cancel() + return + + async def _release_lease(self) -> None: + heartbeat, self._heartbeat_task = self._heartbeat_task, None + lease, self._lease = self._lease, None + if heartbeat is not None: + heartbeat.cancel() + if lease is not None and self._lease_store is not None: + try: + await self._lease_store.release(lease) + except Exception: + logger.exception('Lease[%s]: release failed', self._task_id) + + async def _end_lease_on_exit(self) -> None: + """Ends the lease when the consumer stops outside a normal turn end. + + The lease is released only if the task is terminal. Otherwise the + turn was cut short without a final state being written, so renewal + stops and the lease is left to expire, letting the task be recovered. + `cancel` releases it once it has written the terminal state. + """ + if self._lease is None: + return + task = self._task_manager._current_task + if task is not None and task.status.state in TERMINAL_TASK_STATES: + await self._release_lease() + return + heartbeat, self._heartbeat_task = self._heartbeat_task, None + if heartbeat is not None: + heartbeat.cancel() + async def _run_producer(self) -> None: """Executes the agent logic. @@ -705,6 +837,21 @@ async def _run_producer(self) -> None: ) try: + if not await self._acquire_lease(request_context): + # Another replica holds the lease, so it is running a + # turn for this task. Refuse rather than run a second + # executor; nothing is written to the task. + rejection = UnsupportedOperationError( + message=f'Task {self._task_id} is being processed ' + 'by another replica' + ) + await self._event_queue_agent.enqueue_event( + cast( + 'Event', + _RequestRejected(request_id, rejection), + ) + ) + continue await self._event_queue_agent.enqueue_event( cast( 'Event', @@ -755,6 +902,7 @@ async def _run_consumer(self) -> None: finally: self._is_finished.set() self.cancel_idle_release() + await self._end_lease_on_exit() self._request_queue.shutdown(immediate=True) await self._event_queue_agent.close(immediate=True) @@ -775,7 +923,7 @@ async def _run_consumer(self) -> None: logger.debug('Consumer[%s]: Finishing', self._task_id) await self._maybe_cleanup() - async def subscribe( # noqa: PLR0915 + async def subscribe( self, *, request: RequestContext | None = None, @@ -850,25 +998,23 @@ async def subscribe( # noqa: PLR0915 ) event = updated_task # Acknowledge internal lifecycle events before skipping them. - if isinstance(event, _RequestStarted | _RequestCompleted): + if isinstance( + event, + _RequestStarted | _RequestCompleted | _RequestRejected, + ): tapped_queue.task_done() - if isinstance(event, _RequestCompleted): if ( - request_id is not None - and event.request_id == request_id + isinstance(event, _RequestStarted) + or event.request_id != request_id ): - logger.debug( - 'Subscriber[%s]: Request completed', - self._task_id, - ) - return - continue - elif isinstance(event, _RequestStarted): + continue + if isinstance(event, _RequestRejected): + raise event.error logger.debug( - 'Subscriber[%s]: Request started', + 'Subscriber[%s]: Request completed', self._task_id, ) - continue + return try: yield event finally: @@ -981,6 +1127,7 @@ async def cancel(self, call_context: ServerCallContext) -> Task: await self._task_manager.save_task_event(updated) await send_push_notification(self._push_sender, updated.id, updated) task = updated + await self._release_lease() return task async def aclose(self) -> None: @@ -1027,12 +1174,19 @@ async def aclose(self) -> None: logger.error('Error during aclose', exc_info=result) def schedule_idle_release(self) -> None: - """(Re)starts the timer that releases a parked interrupted task.""" - if self._interrupted_task_timeout is None or self._is_finished.is_set(): + """(Re)starts the timer that releases a parked interrupted task. + + A task running on another replica is released as soon as it is idle, + whatever its state and the timeout. + """ + timeout = ( + 0.0 if self._running_elsewhere else self._interrupted_task_timeout + ) + if timeout is None or self._is_finished.is_set(): return self.cancel_idle_release() self._idle_release_task = asyncio.create_task( - self._release_when_idle(self._interrupted_task_timeout), + self._release_when_idle(timeout), name=f'idle-release:{self._task_id}', ) @@ -1058,7 +1212,10 @@ async def _release_when_idle(self, timeout: float) -> None: async with self._lock: if ( task is None - or task.status.state not in INTERRUPTED_TASK_STATES + or ( + task.status.state not in INTERRUPTED_TASK_STATES + and not self._running_elsewhere + ) or self._is_finished.is_set() # The consumer holds one reference; more means subscribers. or self._reference_count != 1 diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index e5daaaa2e..d661a4fd9 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.event_stream import TaskEventStream + from a2a.server.cluster.lease import TaskLeaseStore from a2a.server.cluster.task_store import VersionedTaskStore from a2a.server.context import ServerCallContext from a2a.server.tasks.push_notification_sender import PushNotificationSender @@ -17,6 +18,10 @@ from a2a.types.a2a_pb2 import Message from a2a.server.agent_execution.active_task import ActiveTask +from a2a.server.cluster.lease import ( + DEFAULT_HEARTBEAT_INTERVAL, + DEFAULT_LEASE_TTL, +) from a2a.server.cluster.task_store import StoredTask from a2a.server.tasks.task_manager import TaskManager from a2a.utils.errors import TaskNotFoundError @@ -28,19 +33,25 @@ class ActiveTaskRegistry: """A registry for active ActiveTask instances.""" - def __init__( + def __init__( # noqa: PLR0913 self, agent_executor: AgentExecutor, task_store: TaskStore | VersionedTaskStore, push_sender: PushNotificationSender | None = None, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + lease_store: TaskLeaseStore | None = None, + lease_ttl: float = DEFAULT_LEASE_TTL, + heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL, ): self._agent_executor = agent_executor self._task_store = task_store self._push_sender = push_sender self._interrupted_task_timeout = interrupted_task_timeout self._event_stream = event_stream + self._lease_store = lease_store + self._lease_ttl = lease_ttl + self._heartbeat_interval = heartbeat_interval self._active_tasks: dict[str, ActiveTask] = {} self._lock = threading.RLock() self._cleanup_tasks: set[asyncio.Task[None]] = set() @@ -82,6 +93,9 @@ async def get_or_create( on_cleanup=self._on_active_task_cleanup, event_stream=self._event_stream, interrupted_task_timeout=self._interrupted_task_timeout, + lease_store=self._lease_store, + lease_ttl=self._lease_ttl, + heartbeat_interval=self._heartbeat_interval, ) self._active_tasks[task_id] = active_task @@ -135,6 +149,11 @@ async def _remove_task(self, active_task: ActiveTask) -> None: 'Removed active task for %s from registry', task_id ) + def snapshot(self) -> list[ActiveTask]: + """The currently registered tasks.""" + with self._lock: + return list(self._active_tasks.values()) + async def get(self, task_id: str) -> ActiveTask | None: """Retrieves an existing task.""" with self._lock: diff --git a/src/a2a/server/cluster/__init__.py b/src/a2a/server/cluster/__init__.py index f7c4f1ac9..e1c5e0664 100644 --- a/src/a2a/server/cluster/__init__.py +++ b/src/a2a/server/cluster/__init__.py @@ -1,6 +1,7 @@ import logging from a2a.server.cluster.event_stream import TaskEventStream, VersionedEvent +from a2a.server.cluster.lease import TaskLease, TaskLeaseStore from a2a.server.cluster.task_store import ( ConcurrentTaskModificationError, LegacyTaskStoreAdapter, @@ -14,6 +15,7 @@ try: from a2a.server.cluster.database_event_stream import DatabaseTaskEventStream + from a2a.server.cluster.database_lease_store import DatabaseTaskLeaseStore from a2a.server.cluster.database_task_store import ( VersionedDatabaseTaskStore, ) @@ -43,13 +45,25 @@ def __init__(self, *args: object, **kwargs: object) -> None: "installed. Install with 'pip install a2a-sdk[sql]'." ) from _original_error + class DatabaseTaskLeaseStore: # type: ignore[no-redef] + """Placeholder when database dependencies are not installed.""" + + def __init__(self, *args: object, **kwargs: object) -> None: + raise ImportError( + 'To use DatabaseTaskLeaseStore, its dependencies must be ' + "installed. Install with 'pip install a2a-sdk[sql]'." + ) from _original_error + __all__ = [ 'ConcurrentTaskModificationError', 'DatabaseTaskEventStream', + 'DatabaseTaskLeaseStore', 'LegacyTaskStoreAdapter', 'StoredTask', 'TaskEventStream', + 'TaskLease', + 'TaskLeaseStore', 'TaskVersion', 'VersionedDatabaseTaskStore', 'VersionedEvent', diff --git a/src/a2a/server/cluster/database_lease_store.py b/src/a2a/server/cluster/database_lease_store.py new file mode 100644 index 000000000..b58c39724 --- /dev/null +++ b/src/a2a/server/cluster/database_lease_store.py @@ -0,0 +1,174 @@ +import time +import uuid + +from collections.abc import Callable + + +try: + from sqlalchemy import Table, and_, delete, insert, select, update + from sqlalchemy.exc import IntegrityError + from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker + from sqlalchemy.orm import class_mapper +except ImportError as e: + raise ImportError( + 'DatabaseTaskLeaseStore requires SQLAlchemy and a database driver. ' + 'Install with one of: ' + "'pip install a2a-sdk[postgresql]', " + "'pip install a2a-sdk[mysql]', " + "'pip install a2a-sdk[sqlite]', " + "or 'pip install a2a-sdk[sql]'" + ) from e + +from a2a.server.cluster.lease import TaskLease, TaskLeaseStore +from a2a.server.context import ServerCallContext +from a2a.server.models import Base, TaskLeaseModel, create_task_lease_model +from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope + + +class DatabaseTaskLeaseStore(TaskLeaseStore): + """`TaskLeaseStore` backed by a ``task_leases`` table. + + Uses conditional UPDATEs and the primary key for mutual exclusion, so it + works the same on PostgreSQL, MySQL and SQLite. Expiry times come from + each replica's `clock`, so replica clocks must agree to well within the + lease TTL minus the heartbeat interval. + """ + + _lease_model: type[TaskLeaseModel] + + def __init__( + self, + engine: AsyncEngine, + create_table: bool = True, + table_name: str = 'task_leases', + owner_resolver: OwnerResolver = resolve_user_scope, + clock: Callable[[], float] = time.time, + ) -> None: + """Initializes the store over an existing SQLAlchemy AsyncEngine. + + `owner_resolver` must match the one used by the task store, so that + a lease records the owner its task is stored under. + """ + self._engine = engine + self._session_maker = async_sessionmaker(engine, expire_on_commit=False) + self._create_table = create_table + self._owner_resolver = owner_resolver + self._clock = clock + self._initialized = False + self._lease_model = ( # ty:ignore[invalid-assignment] + TaskLeaseModel + if table_name == 'task_leases' + else create_task_lease_model(table_name) + ) + + async def initialize(self) -> None: + """Creates the ``task_leases`` table if requested.""" + if self._initialized: + return + if self._create_table: + async with self._engine.begin() as conn: + mapper = class_mapper(self._lease_model) + tables = [t for t in mapper.tables if isinstance(t, Table)] + await conn.run_sync(Base.metadata.create_all, tables=tables) + self._initialized = True + + def _now_ms(self) -> int: + return int(self._clock() * 1000) + + async def acquire( + self, task_id: str, context: ServerCallContext, ttl: float + ) -> TaskLease | None: + """Takes the lease, or returns None while another holder's is live.""" + await self.initialize() + model = self._lease_model + owner = self._owner_resolver(context) + now = self._now_ms() + token = uuid.uuid4().hex + values = { + 'owner': owner, + 'token': token, + 'expires_at_ms': now + int(ttl * 1000), + } + async with self._session_maker.begin() as session: + # Take over a lease whose holder stopped renewing it. + result = await session.execute( + update(model) + .where( + and_(model.task_id == task_id, model.expires_at_ms < now) + ) + .values(**values) + ) + if result.rowcount == 1: # ty:ignore[unresolved-attribute] + return TaskLease(task_id, owner, token) + try: + # The primary key lets only one of two racing acquirers in. + async with self._session_maker.begin() as session: + await session.execute( + insert(model).values(task_id=task_id, **values) + ) + except IntegrityError: + return None + return TaskLease(task_id, owner, token) + + async def renew(self, lease: TaskLease, ttl: float) -> bool: + """Extends the lease; False if it was released or taken over.""" + model = self._lease_model + async with self._session_maker.begin() as session: + result = await session.execute( + update(model) + .where( + and_( + model.task_id == lease.task_id, + model.token == lease.token, + ) + ) + .values(expires_at_ms=self._now_ms() + int(ttl * 1000)) + ) + return result.rowcount == 1 # ty:ignore[unresolved-attribute] + + async def release(self, lease: TaskLease) -> None: + """Drops the lease if it is still held with this token.""" + model = self._lease_model + async with self._session_maker.begin() as session: + await session.execute( + delete(model).where( + and_( + model.task_id == lease.task_id, + model.token == lease.token, + ) + ) + ) + + async def claim_expired(self, ttl: float, limit: int) -> list[TaskLease]: + """Claims expired leases, each for exactly one caller.""" + await self.initialize() + model = self._lease_model + now = self._now_ms() + async with self._session_maker() as session: + rows = ( + await session.execute( + select(model.task_id, model.owner) + .where(model.expires_at_ms < now) + .order_by(model.expires_at_ms) + .limit(limit) + ) + ).all() + claimed: list[TaskLease] = [] + for task_id, owner in rows: + token = uuid.uuid4().hex + async with self._session_maker.begin() as session: + # Re-checked here: another sweeper or a late renewal may have + # won since the SELECT. + result = await session.execute( + update(model) + .where( + and_( + model.task_id == task_id, + model.expires_at_ms < now, + ) + ) + .values(token=token, expires_at_ms=now + int(ttl * 1000)) + ) + if result.rowcount == 1: # ty:ignore[unresolved-attribute] + claimed.append(TaskLease(task_id, owner or '', token)) + return claimed diff --git a/src/a2a/server/cluster/lease.py b/src/a2a/server/cluster/lease.py new file mode 100644 index 000000000..b54e10055 --- /dev/null +++ b/src/a2a/server/cluster/lease.py @@ -0,0 +1,56 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass + +from a2a.server.context import ServerCallContext + + +DEFAULT_LEASE_TTL = 45.0 +DEFAULT_HEARTBEAT_INTERVAL = 10.0 + + +@dataclass(frozen=True) +class TaskLease: + """A held turn lease. Only the holder's `token` can renew or release it.""" + + task_id: str + owner: str + token: str + + +class TaskLeaseStore(ABC): + """Records which replica is running a turn for each task. + + The replica running a turn holds the task's lease and renews it while + the turn runs. A lease that expires belongs to a replica that stopped + renewing it, so its task can be recovered (see `fail_abandoned_tasks`). + """ + + @abstractmethod + async def acquire( + self, task_id: str, context: ServerCallContext, ttl: float + ) -> TaskLease | None: + """Takes the lease for `ttl` seconds. + + Returns None if another holder's lease has not expired. An expired + lease is taken over. + """ + + @abstractmethod + async def renew(self, lease: TaskLease, ttl: float) -> bool: + """Extends `lease` by `ttl` seconds. + + Returns False if the lease is no longer held with this token. Must + never recreate a released or taken-over lease. + """ + + @abstractmethod + async def release(self, lease: TaskLease) -> None: + """Drops `lease` if it is still held with its token.""" + + @abstractmethod + async def claim_expired(self, ttl: float, limit: int) -> list[TaskLease]: + """Claims up to `limit` expired leases for recovery. + + Each claimed lease gets a new token and a new expiry `ttl` seconds + away, so concurrent callers never claim the same lease. + """ diff --git a/src/a2a/server/cluster/recovery.py b/src/a2a/server/cluster/recovery.py new file mode 100644 index 000000000..2c0853740 --- /dev/null +++ b/src/a2a/server/cluster/recovery.py @@ -0,0 +1,218 @@ +import asyncio +import contextlib +import logging +import uuid + +from collections.abc import Callable +from datetime import datetime, timezone + +from a2a.auth.user import User +from a2a.server.agent_execution.active_task import ( + INTERRUPTED_TASK_STATES, + TERMINAL_TASK_STATES, + send_push_notification, +) +from a2a.server.cluster.event_stream import TaskEventStream, VersionedEvent +from a2a.server.cluster.lease import TaskLeaseStore +from a2a.server.cluster.task_store import ( + ConcurrentTaskModificationError, + VersionedTaskStore, +) +from a2a.server.context import ServerCallContext +from a2a.server.tasks.push_notification_sender import PushNotificationSender +from a2a.types.a2a_pb2 import ( + Message, + Part, + Role, + Task, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) + + +logger = logging.getLogger(__name__) + +ABANDONED_TASK_REASON = 'The replica running this task stopped responding.' + + +class _OwnerUser(User): + def __init__(self, owner: str) -> None: + self._owner = owner + + @property + def is_authenticated(self) -> bool: + return bool(self._owner) + + @property + def user_name(self) -> str: + return self._owner + + +def context_for_owner(owner: str) -> ServerCallContext: + """Builds a call context that `resolve_user_scope` maps to `owner`. + + Deployments using a custom `owner_resolver` must pass their own factory + to `fail_abandoned_tasks`. + """ + return ServerCallContext(user=_OwnerUser(owner)) + + +async def fail_task( # noqa: PLR0913 + task_store: VersionedTaskStore, + task_id: str, + context: ServerCallContext, + *, + reason: str, + push_sender: PushNotificationSender | None = None, + event_stream: TaskEventStream | None = None, +) -> Task | None: + """Marks the task FAILED unless it is terminal or waiting for input. + + A task in an interrupted state, e.g. INPUT_REQUIRED, already recorded + the end of its turn, even if the replica died before releasing the + lease, so it is left as it is. + + The write carries its event, which is also published to `event_stream`, + so it ends SubscribeToTask streams on every replica. It bumps the + version, so a replica still running the task stops at its next save. + Returns the failed task, or None if nothing was written. + """ + for _ in range(2): + stored = await task_store.get(task_id, context) + if stored is None or stored.task.status.state in ( + TERMINAL_TASK_STATES | INTERRUPTED_TASK_STATES + ): + return None + failed = Task() + failed.CopyFrom(stored.task) + if failed.status.HasField('message'): + failed.history.append(failed.status.message) + status = TaskStatus( + state=TaskState.TASK_STATE_FAILED, + message=Message( + message_id=str(uuid.uuid4()), + role=Role.ROLE_AGENT, + task_id=task_id, + context_id=failed.context_id, + parts=[Part(text=reason)], + ), + ) + status.timestamp.FromDatetime(datetime.now(timezone.utc)) + failed.status.CopyFrom(status) + event = TaskStatusUpdateEvent( + task_id=task_id, context_id=failed.context_id, status=status + ) + try: + version = await task_store.save( + failed, + event=event, + prev=stored.task, + prev_version=stored.version, + context=context, + ) + except ConcurrentTaskModificationError: + continue + if event_stream is not None: + await event_stream.publish( + task_id, VersionedEvent(event=event, version=version) + ) + await send_push_notification(push_sender, task_id, event) + return failed + return None + + +async def fail_abandoned_tasks( # noqa: PLR0913 + lease_store: TaskLeaseStore, + task_store: VersionedTaskStore, + *, + push_sender: PushNotificationSender | None = None, + event_stream: TaskEventStream | None = None, + context_factory: Callable[[str], ServerCallContext] = context_for_owner, + reason: str = ABANDONED_TASK_REASON, + claim_ttl: float = 60.0, + limit: int = 100, +) -> int: + """Fails the tasks whose turn lease expired. Returns how many it failed. + + Safe to run on every replica at once: each expired lease is claimed by + one caller. A claim that cannot be completed lapses after `claim_ttl` + seconds and is retried by a later sweep. + """ + failed = 0 + for lease in await lease_store.claim_expired(claim_ttl, limit): + try: + if await fail_task( + task_store, + lease.task_id, + context_factory(lease.owner), + reason=reason, + push_sender=push_sender, + event_stream=event_stream, + ): + failed += 1 + except Exception: + logger.exception('Recovering task %s failed', lease.task_id) + continue + await lease_store.release(lease) + return failed + + +class AbandonedTaskSweeper: + """Runs `fail_abandoned_tasks` every `interval` seconds.""" + + def __init__( # noqa: PLR0913 + self, + lease_store: TaskLeaseStore, + task_store: VersionedTaskStore, + *, + interval: float = 15.0, + push_sender: PushNotificationSender | None = None, + event_stream: TaskEventStream | None = None, + context_factory: Callable[[str], ServerCallContext] = context_for_owner, + reason: str = ABANDONED_TASK_REASON, + claim_ttl: float = 60.0, + limit: int = 100, + ) -> None: + self._lease_store = lease_store + self._task_store = task_store + self._interval = interval + self._push_sender = push_sender + self._event_stream = event_stream + self._context_factory = context_factory + self._reason = reason + self._claim_ttl = claim_ttl + self._limit = limit + self._task: asyncio.Task[None] | None = None + + def start(self) -> None: + """Starts sweeping in the background. Calling it twice is a no-op.""" + if self._task is None or self._task.done(): + self._task = asyncio.create_task( + self._run(), name='a2a-abandoned-task-sweeper' + ) + + async def aclose(self) -> None: + """Stops sweeping and waits for the current sweep to end.""" + task, self._task = self._task, None + if task is not None: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _run(self) -> None: + while True: + try: + await fail_abandoned_tasks( + self._lease_store, + self._task_store, + push_sender=self._push_sender, + event_stream=self._event_stream, + context_factory=self._context_factory, + reason=self._reason, + claim_ttl=self._claim_ttl, + limit=self._limit, + ) + except Exception: + logger.exception('Abandoned-task sweep failed') + await asyncio.sleep(self._interval) diff --git a/src/a2a/server/models.py b/src/a2a/server/models.py index b296900db..2e8084ead 100644 --- a/src/a2a/server/models.py +++ b/src/a2a/server/models.py @@ -279,3 +279,40 @@ class TaskVersionModel(TaskVersionMixin, Base): """Default task-version model with standard table name.""" __tablename__ = 'task_versions' + + +class TaskLeaseMixin: + """Mixin providing columns for the turn lease of a task. + + A row exists while a replica runs a turn for the task. Only used by + `a2a.server.cluster`; unused by the default single-process stores. + """ + + task_id: Mapped[str] = mapped_column(String(36), primary_key=True) + owner: Mapped[str] = mapped_column(String(255), nullable=True) + # Identifies the holder; renewing or releasing requires it. + token: Mapped[str] = mapped_column(String(32), nullable=False) + expires_at_ms: Mapped[int] = mapped_column( + BigInteger, nullable=False, index=True + ) + + +def create_task_lease_model( + table_name: str = 'task_leases', + base: type[DeclarativeBase] = Base, +) -> type: + """Create a TaskLeaseModel class with a configurable table name.""" + + class TaskLeaseModel(TaskLeaseMixin, base): # type: ignore + __tablename__ = table_name + + TaskLeaseModel.__name__ = f'TaskLeaseModel_{table_name}' + TaskLeaseModel.__qualname__ = f'TaskLeaseModel_{table_name}' + return TaskLeaseModel + + +# Default TaskLeaseModel with the standard table name. +class TaskLeaseModel(TaskLeaseMixin, Base): + """Default task-lease model with standard table name.""" + + __tablename__ = 'task_leases' diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index c96bf8c9e..761ae0d6c 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -1,6 +1,6 @@ from __future__ import annotations -import asyncio # noqa: TC003 +import asyncio import logging import warnings @@ -20,6 +20,11 @@ ) from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry from a2a.server.cluster.event_stream import VersionedEvent +from a2a.server.cluster.lease import ( + DEFAULT_HEARTBEAT_INTERVAL, + DEFAULT_LEASE_TTL, +) +from a2a.server.cluster.recovery import fail_task from a2a.server.cluster.task_store import ( ConcurrentTaskModificationError, LegacyTaskStoreAdapter, @@ -76,6 +81,7 @@ from a2a.server.agent_execution.active_task import ActiveTask from a2a.server.cluster.event_stream import TaskEventStream + from a2a.server.cluster.lease import TaskLeaseStore from a2a.server.cluster.version import TaskVersion from a2a.server.context import ServerCallContext from a2a.server.events import Event @@ -102,6 +108,12 @@ class DefaultRequestHandlerV2(RequestHandler): ``input-required`` or ``auth-required`` after that many seconds; a follow-up recreates it from the store. ``None`` (default) never releases; any other value must be greater than 0. + + With a ``lease_store``, each turn holds the task's lease and renews it + every ``heartbeat_interval`` seconds. A follow-up arriving while another + replica holds the lease is rejected with `UnsupportedOperationError`, and + a task whose lease expired can be failed with + `a2a.server.cluster.recovery.fail_abandoned_tasks`. """ _background_tasks: set[asyncio.Task] @@ -125,6 +137,9 @@ def __init__( # noqa: PLR0913 validate_input_modes: bool = False, event_stream: TaskEventStream | None = None, interrupted_task_timeout: float | None = None, + lease_store: TaskLeaseStore | None = None, + lease_ttl: float = DEFAULT_LEASE_TTL, + heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL, ) -> None: if ( interrupted_task_timeout is not None @@ -185,6 +200,9 @@ def __init__( # noqa: PLR0913 push_sender=self._push_sender, event_stream=self._event_stream, interrupted_task_timeout=interrupted_task_timeout, + lease_store=lease_store, + lease_ttl=lease_ttl, + heartbeat_interval=heartbeat_interval, ) self._background_tasks = set() @@ -195,15 +213,52 @@ async def _reject_unsafe_push_url(self, url: str) -> None: if not await self._push_url_validator(url): raise InvalidParamsError(message='Invalid push notification URL') - async def aclose(self) -> None: + async def aclose( + self, *, fail_in_flight: bool | None = None, timeout: float = 5.0 + ) -> None: """Shuts down the handler, draining all active tasks. Drains the ``ActiveTaskRegistry`` so a server shutdown leaves no pending ``asyncio.Task``. Intended to be wired into an ASGI ``lifespan`` / ``on_shutdown`` hook. Safe to call multiple times. + + Args: + fail_in_flight: Whether to mark tasks with a turn still running + FAILED first, spending at most `timeout` seconds on it. + Defaults to True when an ``event_stream`` is configured: the + shared store outlives this replica, and the task would + otherwise stay WORKING for every replica. + timeout: Seconds allowed for failing in-flight tasks. """ + if fail_in_flight is None: + fail_in_flight = self._event_stream is not None + if fail_in_flight: + try: + await asyncio.wait_for(self._fail_in_flight_tasks(), timeout) + except asyncio.TimeoutError: + logger.warning('Timed out failing in-flight tasks on close') await self._active_task_registry.aclose() + async def _fail_in_flight_tasks(self) -> None: + for active_task in self._active_task_registry.snapshot(): + context = active_task.call_context + if not active_task.request_in_flight or context is None: + continue + try: + await fail_task( + self._versioned_store, + active_task.task_id, + context, + reason='The server shut down while the task was running.', + push_sender=self._push_sender, + event_stream=self._event_stream, + ) + except Exception: + logger.exception( + 'Failing in-flight task %s on close failed', + active_task.task_id, + ) + @validate_request_params async def on_get_task( # noqa: D102 self, diff --git a/tests/integration/test_multiserver.py b/tests/integration/test_multiserver.py index d376e4cc6..5ef845c06 100644 --- a/tests/integration/test_multiserver.py +++ b/tests/integration/test_multiserver.py @@ -16,7 +16,9 @@ from a2a.helpers.proto_helpers import new_task_from_user_message from a2a.server.agent_execution.agent_executor import AgentExecutor from a2a.server.cluster.database_event_stream import DatabaseTaskEventStream +from a2a.server.cluster.database_lease_store import DatabaseTaskLeaseStore from a2a.server.cluster.database_task_store import VersionedDatabaseTaskStore +from a2a.server.cluster.recovery import fail_abandoned_tasks from a2a.server.context import ServerCallContext from a2a.server.models import Base from a2a.server.request_handlers.default_request_handler_v2 import ( @@ -593,3 +595,140 @@ async def test_absent_and_non_owner_rejected(dsn: str) -> None: ) finally: await cluster.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(60) +async def test_task_of_stopped_replica_is_recovered_after_lease_expiry( + dsn: str, +) -> None: + """Replica A stops mid-turn without writing anything more. Once its lease + expires, a sweep marks the task FAILED, which ends a stream on B.""" + engines = [create_async_engine(dsn) for _ in range(3)] + for engine in engines: + await DatabaseTaskLeaseStore(engine).initialize() + + def replica(engine, agent: AgentExecutor) -> DefaultRequestHandlerV2: # noqa: ANN001 + return DefaultRequestHandlerV2( + agent_executor=agent, + task_store=VersionedDatabaseTaskStore(engine, create_table=False), + agent_card=streaming_agent_card(), + event_stream=DatabaseTaskEventStream( + engine, create_table=False, poll_interval_s=0.05 + ), + lease_store=DatabaseTaskLeaseStore(engine, create_table=False), + lease_ttl=1.0, + heartbeat_interval=0.1, + ) + + agent = LongRunningAgent() + replica_a = replica(engines[0], agent) + replica_b = replica(engines[1], LongRunningAgent()) + ctx = make_context() + try: + stream = replica_a.on_message_send_stream( + build_send_request('go', message_id='m1'), ctx + ) + task_id = (await asyncio.wait_for(anext(stream), timeout=10)).id + await asyncio.wait_for(agent.working.wait(), timeout=10) + await wait_for_state( + replica_b.task_store, task_id, TaskState.TASK_STATE_WORKING, ctx + ) + b_states: list = [] + + async def follow_on_b() -> None: + async for event in replica_b.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), ctx + ): + b_states.append(event.status.state) + + follow = asyncio.create_task(follow_on_b()) + + # Stop A the way a killed process would: no further writes. + active = await replica_a._active_task_registry.get(task_id) # noqa: SLF001 + await engines[0].dispose() + for task in (active._producer_task, active._consumer_task): # noqa: SLF001 + task.cancel() + with contextlib.suppress(BaseException): + await stream.aclose() + + sweeper_store = VersionedDatabaseTaskStore( + engines[2], create_table=False + ) + sweeper_leases = DatabaseTaskLeaseStore(engines[2], create_table=False) + recovered = 0 + for _ in range(60): + recovered = await fail_abandoned_tasks( + sweeper_leases, sweeper_store + ) + if recovered: + break + await asyncio.sleep(0.1) + assert recovered == 1 + + await asyncio.wait_for(follow, timeout=10) + assert b_states[-1] == TaskState.TASK_STATE_FAILED + task = await replica_b.on_get_task(GetTaskRequest(id=task_id), ctx) + assert task is not None + assert task.status.state == TaskState.TASK_STATE_FAILED + finally: + with contextlib.suppress(Exception): + await replica_b.aclose() + for engine in engines: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(60) +async def test_leased_turns_alternate_between_replicas_without_pauses( + dsn: str, +) -> None: + """With leases, each follow-up can go to the other replica as soon as the + previous turn returned INPUT_REQUIRED.""" + engines = [create_async_engine(dsn) for _ in range(2)] + await DatabaseTaskLeaseStore(engines[0]).initialize() + + def replica(engine) -> DefaultRequestHandlerV2: # noqa: ANN001 + return DefaultRequestHandlerV2( + agent_executor=InputRequiredThenCompleteAgent(), + task_store=VersionedDatabaseTaskStore(engine, create_table=False), + agent_card=streaming_agent_card(), + event_stream=DatabaseTaskEventStream( + engine, create_table=False, poll_interval_s=0.05 + ), + lease_store=DatabaseTaskLeaseStore(engine, create_table=False), + lease_ttl=5.0, + heartbeat_interval=1.0, + ) + + replicas = [replica(engine) for engine in engines] + ctx = make_context() + try: + first = await replicas[0].on_message_send( + build_send_request('q1', message_id='m1'), ctx + ) + result = first + for i, message_id in enumerate(('m2', 'm3'), start=1): + result = await replicas[i % 2].on_message_send( + build_send_request( + 'answer', + task_id=first.id, + context_id=first.context_id, + message_id=message_id, + ), + ctx, + ) + assert result.status.state == TaskState.TASK_STATE_COMPLETED + stored = await replicas[0].task_store.get(first.id, ctx) + assert stored is not None + assert [ + m.message_id + for m in stored.task.history + if m.role == Role.ROLE_USER + ] == ['m1', 'm2', 'm3'] + finally: + for r in replicas: + with contextlib.suppress(Exception): + await r.aclose() + for engine in engines: + await engine.dispose() diff --git a/tests/migrations/versions/test_migration_b0f01b04ff26.py b/tests/migrations/versions/test_migration_b0f01b04ff26.py new file mode 100644 index 000000000..b5bde3211 --- /dev/null +++ b/tests/migrations/versions/test_migration_b0f01b04ff26.py @@ -0,0 +1,150 @@ +import importlib +import os +import sqlite3 +import tempfile + +from typing import Generator +from unittest.mock import patch + +import pytest + +from a2a.a2a_db_cli import run_migrations + + +# Explicitly import the migration module so it is tracked when Alembic loads it +# dynamically. Revision id starts with a letter, so this import is valid. +try: + importlib.import_module( + 'a2a.migrations.versions.b0f01b04ff26_add_task_leases' + ) +except (ImportError, AttributeError): + pass + + +REVISION = 'b0f01b04ff26' +PREV_REVISION = 'b5e3d1c8a2f7' + + +@pytest.fixture(autouse=True) +def mock_logging_config(): + """Prevent tests from mutating global logging state.""" + with patch('logging.basicConfig'), patch('logging.config.fileConfig'): + yield + + +@pytest.fixture +def temp_db() -> Generator[str, None, None]: + """Create a temporary SQLite database for testing.""" + fd, path = tempfile.mkstemp(suffix='.db') + os.close(fd) + yield path + if os.path.exists(path): + os.remove(path) + + +def _setup_initial_schema(db_path: str) -> None: + """Create the base tasks/push tables the earlier migrations expect.""" + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute(""" + CREATE TABLE tasks ( + id VARCHAR(36) PRIMARY KEY, + context_id VARCHAR(36) NOT NULL, + kind VARCHAR(16) NOT NULL, + status TEXT, + artifacts TEXT, + history TEXT, + metadata TEXT + ) + """) + cursor.execute(""" + CREATE TABLE push_notification_configs ( + task_id VARCHAR(36), + config_id VARCHAR(255), + config_data BLOB NOT NULL, + PRIMARY KEY (task_id, config_id) + ) + """) + conn.commit() + conn.close() + + +def _upgrade(db_url: str, revision: str = REVISION) -> None: + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, 'upgrade', revision], + ): + run_migrations() + + +def _tables(db_path: str) -> set[str]: + conn = sqlite3.connect(db_path) + try: + cursor = conn.execute( + "SELECT name FROM sqlite_schema WHERE type='table'" + ) + return {row[0] for row in cursor.fetchall()} + finally: + conn.close() + + +def test_migration_b0f01b04ff26_full_cycle(temp_db: str) -> None: + """Upgrade creates task_leases; downgrade removes only that table.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + + _upgrade(db_url) + + conn = sqlite3.connect(temp_db) + cursor = conn.cursor() + cursor.execute('PRAGMA table_info(task_leases)') + assert {row[1] for row in cursor.fetchall()} == { + 'task_id', + 'owner', + 'token', + 'expires_at_ms', + } + cursor.execute('PRAGMA index_list(task_leases)') + assert 'ix_task_leases_expires_at_ms' in { + row[1] for row in cursor.fetchall() + } + conn.close() + + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, 'downgrade', PREV_REVISION], + ): + run_migrations() + + tables = _tables(temp_db) + assert 'task_leases' not in tables + assert {'task_versions', 'task_events'} <= tables + + +def test_migration_b0f01b04ff26_idempotency(temp_db: str) -> None: + """Running the upgrade twice must not fail.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + _upgrade(db_url) + _upgrade(db_url) + + assert 'task_leases' in _tables(temp_db) + + +def test_migration_b0f01b04ff26_offline( + temp_db: str, capsys: pytest.CaptureFixture[str] +) -> None: + """Offline (--sql) mode emits the DDL without touching the database.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, '--sql', 'upgrade', REVISION], + ): + run_migrations() + + out = capsys.readouterr().out + assert 'CREATE TABLE task_leases' in out + assert 'ix_task_leases_expires_at_ms' in out + assert 'task_leases' not in _tables(temp_db) diff --git a/tests/server/cluster/conftest.py b/tests/server/cluster/conftest.py index c26d2a7d8..91b303c7e 100644 --- a/tests/server/cluster/conftest.py +++ b/tests/server/cluster/conftest.py @@ -1,6 +1,9 @@ +from __future__ import annotations + import asyncio import contextlib import threading +import uuid from collections.abc import AsyncGenerator @@ -13,6 +16,8 @@ ConcurrentTaskModificationError, StoredTask, TaskEventStream, + TaskLease, + TaskLeaseStore, TaskVersion, VersionedEvent, VersionedTaskStore, @@ -203,6 +208,93 @@ async def destroy(self, task_id: str) -> None: self._subscribers.pop(task_id, None) +class FakeClock: + """A settable clock for lease expiry.""" + + def __init__(self, now: float = 1_000_000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class InMemoryTaskLeaseStore(TaskLeaseStore): + """`TaskLeaseStore` kept in a dict, with a settable clock. + + Set `unreachable` to make every call raise, as if the database could not + be reached from this replica. + """ + + def __init__( + self, + clock: FakeClock | None = None, + owner_resolver: OwnerResolver = resolve_user_scope, + ) -> None: + self.clock = clock or FakeClock() + self._owner_resolver = owner_resolver + self._leases: dict[str, tuple[str, str, float]] = {} + self.unreachable = False + + def view(self) -> InMemoryTaskLeaseStore: + """Another replica's handle on the same leases.""" + other = InMemoryTaskLeaseStore(self.clock, self._owner_resolver) + other._leases = self._leases + return other + + def held(self, task_id: str) -> bool: + lease = self._leases.get(task_id) + return lease is not None and lease[2] > self.clock() + + def _check_reachable(self) -> None: + if self.unreachable: + raise ConnectionError('lease store unreachable') + + async def acquire( + self, task_id: str, context: ServerCallContext, ttl: float + ) -> TaskLease | None: + self._check_reachable() + now = self.clock() + current = self._leases.get(task_id) + if current is not None and current[2] > now: + return None + owner = self._owner_resolver(context) + token = uuid.uuid4().hex + self._leases[task_id] = (owner, token, now + ttl) + return TaskLease(task_id, owner, token) + + async def renew(self, lease: TaskLease, ttl: float) -> bool: + self._check_reachable() + current = self._leases.get(lease.task_id) + if current is None or current[1] != lease.token: + return False + self._leases[lease.task_id] = ( + current[0], + current[1], + self.clock() + ttl, + ) + return True + + async def release(self, lease: TaskLease) -> None: + self._check_reachable() + current = self._leases.get(lease.task_id) + if current is not None and current[1] == lease.token: + del self._leases[lease.task_id] + + async def claim_expired(self, ttl: float, limit: int) -> list[TaskLease]: + self._check_reachable() + now = self.clock() + claimed = [] + for task_id, (owner, _, expires_at) in list(self._leases.items()): + if expires_at <= now and len(claimed) < limit: + token = uuid.uuid4().hex + self._leases[task_id] = (owner, token, now + ttl) + claimed.append(TaskLease(task_id, owner, token)) + return claimed + + # --- Identity / context helpers --------------------------------------------- @@ -404,11 +496,15 @@ async def cancel(self, context, event_queue) -> None: # noqa: ANN001 # --- Replica / infra factories ---------------------------------------------- -def make_replica( +def make_replica( # noqa: PLR0913 store: VersionedTaskStore, stream: TaskEventStream, agent: AgentExecutor, push_sender: PushNotificationSender | None = None, + lease_store: TaskLeaseStore | None = None, + lease_ttl: float = 1.0, + heartbeat_interval: float = 0.02, + interrupted_task_timeout: float | None = None, ) -> DefaultRequestHandlerV2: """Builds a handler ('replica') wired to a shared store and stream.""" return DefaultRequestHandlerV2( @@ -417,6 +513,10 @@ def make_replica( agent_card=streaming_agent_card(), event_stream=stream, push_sender=push_sender, + lease_store=lease_store, + lease_ttl=lease_ttl, + heartbeat_interval=heartbeat_interval, + interrupted_task_timeout=interrupted_task_timeout, ) diff --git a/tests/server/cluster/test_database_lease_store.py b/tests/server/cluster/test_database_lease_store.py new file mode 100644 index 000000000..e3ed680d8 --- /dev/null +++ b/tests/server/cluster/test_database_lease_store.py @@ -0,0 +1,247 @@ +"""Tests for `DatabaseTaskLeaseStore`.""" + +import asyncio +import os + +from collections.abc import AsyncGenerator +from pathlib import Path + +import pytest +import pytest_asyncio + +from _pytest.mark.structures import ParameterSet + + +pytest.importorskip('sqlalchemy', reason='Database tests require SQLAlchemy') + +from a2a.auth.user import User +from a2a.server.cluster import DatabaseTaskLeaseStore, TaskLease +from a2a.server.context import ServerCallContext +from a2a.server.models import Base +from sqlalchemy.engine import make_url +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine + + +class _User(User): + def __init__(self, user_name: str) -> None: + self._user_name = user_name + + @property + def is_authenticated(self) -> bool: + return True + + @property + def user_name(self) -> str: + return self._user_name + + +ALICE = ServerCallContext(user=_User('alice')) + +SQLITE_TEST_DSN = ( + 'sqlite+aiosqlite:///file:testdb_leases?mode=memory&cache=shared&uri=true' +) +POSTGRES_TEST_DSN = os.environ.get('POSTGRES_TEST_DSN') +MYSQL_TEST_DSN = os.environ.get('MYSQL_TEST_DSN') + +DB_CONFIGS: list[ParameterSet | tuple[str | None, str]] = [ + pytest.param((SQLITE_TEST_DSN, 'sqlite'), id='sqlite') +] +for _dsn, _name, _var in ( + (POSTGRES_TEST_DSN, 'postgresql', 'POSTGRES_TEST_DSN'), + (MYSQL_TEST_DSN, 'mysql', 'MYSQL_TEST_DSN'), +): + if _dsn: + DB_CONFIGS.append(pytest.param((_dsn, _name), id=_name)) + else: + DB_CONFIGS.append( + pytest.param( + (None, _name), + marks=pytest.mark.skip(reason=f'{_var} not set'), + id=f'{_name}_skipped', + ) + ) + + +class _Clock: + def __init__(self) -> None: + self.now = 1_000_000.0 + + def __call__(self) -> float: + return self.now + + +@pytest_asyncio.fixture(params=DB_CONFIGS) +async def engine(request) -> AsyncGenerator[AsyncEngine, None]: # noqa: ANN001 + db_url, dialect = request.param + if db_url is None: + pytest.skip(f'DSN for {dialect} not set.') + engine = create_async_engine(db_url) + # Holding a connection keeps a shared in-memory SQLite database alive. + keeper = await engine.connect() + try: + yield engine + finally: + await keeper.close() + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.drop_all) + await engine.dispose() + + +@pytest.fixture +def clock() -> _Clock: + return _Clock() + + +@pytest_asyncio.fixture +async def leases(engine: AsyncEngine, clock: _Clock) -> DatabaseTaskLeaseStore: + store = DatabaseTaskLeaseStore(engine, clock=clock) + await store.initialize() + return store + + +@pytest.mark.asyncio +async def test_acquire_records_owner_and_blocks_a_second_holder( + leases: DatabaseTaskLeaseStore, +) -> None: + lease = await leases.acquire('task-1', ALICE, ttl=30) + + assert lease is not None + assert lease.task_id == 'task-1' + assert lease.owner == 'alice' + assert await leases.acquire('task-1', ALICE, ttl=30) is None + assert await leases.acquire('task-2', ALICE, ttl=30) is not None + + +@pytest.mark.asyncio +async def test_expired_lease_is_taken_over( + leases: DatabaseTaskLeaseStore, clock: _Clock +) -> None: + first = await leases.acquire('task-1', ALICE, ttl=30) + assert first is not None + clock.now += 31 + + second = await leases.acquire('task-1', ALICE, ttl=30) + + assert second is not None + assert second.token != first.token + assert await leases.renew(first, ttl=30) is False + assert await leases.renew(second, ttl=30) is True + + +@pytest.mark.asyncio +async def test_renew_extends_the_lease( + leases: DatabaseTaskLeaseStore, clock: _Clock +) -> None: + lease = await leases.acquire('task-1', ALICE, ttl=30) + assert lease is not None + clock.now += 20 + assert await leases.renew(lease, ttl=30) is True + clock.now += 20 + + assert await leases.acquire('task-1', ALICE, ttl=30) is None + assert await leases.claim_expired(ttl=60, limit=10) == [] + + +@pytest.mark.asyncio +async def test_renew_after_release_does_not_recreate_the_lease( + leases: DatabaseTaskLeaseStore, clock: _Clock +) -> None: + lease = await leases.acquire('task-1', ALICE, ttl=30) + assert lease is not None + await leases.release(lease) + + assert await leases.renew(lease, ttl=30) is False + clock.now += 31 + assert await leases.claim_expired(ttl=60, limit=10) == [] + + +@pytest.mark.asyncio +async def test_release_with_a_stale_token_keeps_the_current_lease( + leases: DatabaseTaskLeaseStore, +) -> None: + lease = await leases.acquire('task-1', ALICE, ttl=30) + assert lease is not None + + await leases.release(TaskLease('task-1', 'alice', 'not-the-token')) + + assert await leases.renew(lease, ttl=30) is True + + +@pytest.mark.asyncio +async def test_claim_expired_claims_each_expired_lease_once( + leases: DatabaseTaskLeaseStore, clock: _Clock +) -> None: + stale = await leases.acquire('stale', ALICE, ttl=10) + assert stale is not None + assert await leases.acquire('live', ALICE, ttl=60) is not None + clock.now += 11 + + claimed = await leases.claim_expired(ttl=60, limit=10) + + assert [(c.task_id, c.owner) for c in claimed] == [('stale', 'alice')] + assert claimed[0].token != stale.token + assert await leases.renew(stale, ttl=10) is False + assert await leases.claim_expired(ttl=60, limit=10) == [] + + +@pytest.mark.asyncio +async def test_claim_expired_respects_the_limit( + leases: DatabaseTaskLeaseStore, clock: _Clock +) -> None: + for task_id in ('a', 'b', 'c'): + assert await leases.acquire(task_id, ALICE, ttl=10) is not None + clock.now += 11 + + first = await leases.claim_expired(ttl=60, limit=2) + rest = await leases.claim_expired(ttl=60, limit=2) + + assert len(first) == 2 + assert len(rest) == 1 + assert {c.task_id for c in first + rest} == {'a', 'b', 'c'} + + +@pytest.mark.asyncio +async def test_concurrent_acquire_has_one_winner( + engine: AsyncEngine, clock: _Clock, tmp_path: Path +) -> None: + # One engine per replica. In-memory SQLite engines share one connection, + # which would run the racing transactions as one, so use a file instead. + url = engine.url + if url.get_backend_name() == 'sqlite': + url = make_url(f'sqlite+aiosqlite:///{tmp_path / "leases.db"}') + engines = [create_async_engine(url) for _ in range(5)] + replicas = [DatabaseTaskLeaseStore(e, clock=clock) for e in engines] + try: + await replicas[0].initialize() + + results = await asyncio.gather( + *(r.acquire('task-1', ALICE, ttl=30) for r in replicas), + return_exceptions=True, + ) + + assert sum(isinstance(r, TaskLease) for r in results) == 1 + finally: + for e in engines: + await e.dispose() + + +@pytest.mark.asyncio +async def test_custom_table_name_and_owner_resolver( + engine: AsyncEngine, clock: _Clock +) -> None: + store = DatabaseTaskLeaseStore( + engine, + # A model can be defined once per table name in a process. + table_name=f'custom_leases_{engine.url.get_backend_name()}', + owner_resolver=lambda _: 'tenant-7', + clock=clock, + ) + await store.initialize() + + lease = await store.acquire('task-1', ALICE, ttl=10) + + assert lease is not None + assert lease.owner == 'tenant-7' + clock.now += 11 + claimed = await store.claim_expired(ttl=60, limit=10) + assert [c.owner for c in claimed] == ['tenant-7'] diff --git a/tests/server/cluster/test_handler_wiring.py b/tests/server/cluster/test_handler_wiring.py index 8451aad98..69eba2dc6 100644 --- a/tests/server/cluster/test_handler_wiring.py +++ b/tests/server/cluster/test_handler_wiring.py @@ -1,3 +1,8 @@ +import asyncio +import logging + +from unittest.mock import AsyncMock, MagicMock, patch + import pytest from a2a.server.agent_execution.agent_executor import AgentExecutor @@ -14,6 +19,7 @@ from .conftest import ( InMemoryTaskEventStream, + InMemoryTaskLeaseStore, VersionedInMemoryTaskStore, make_context, ) @@ -163,3 +169,94 @@ def test_plain_store_without_stream_does_not_warn(recwarn) -> None: # noqa: ANN for w in recwarn if 'cross-replica streaming is disabled' in str(w.message) ] + + +@pytest.mark.asyncio +async def test_lease_settings_reach_the_active_task() -> None: + leases = InMemoryTaskLeaseStore() + handler = DefaultRequestHandlerV2( + agent_executor=NoopExecutor(), + task_store=InMemoryTaskStore(), + agent_card=AgentCard(), + lease_store=leases, + lease_ttl=12.0, + heartbeat_interval=3.0, + ) + active_task = await handler._active_task_registry.get_or_create( # noqa: SLF001 + 'task-1', call_context=make_context(), create_task_if_missing=True + ) + try: + assert active_task._lease_store is leases # noqa: SLF001 + assert active_task._lease_ttl == 12.0 # noqa: SLF001 + assert active_task._heartbeat_interval == 3.0 # noqa: SLF001 + finally: + await handler.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('event_stream', 'fail_in_flight', 'expected'), + [ + (None, None, False), + (InMemoryTaskEventStream(), None, True), + (InMemoryTaskEventStream(), False, False), + (None, True, True), + ], +) +async def test_aclose_fails_in_flight_tasks_in_cluster_mode_by_default( + event_stream: TaskEventStream | None, + fail_in_flight: bool | None, + expected: bool, +) -> None: + handler = make_handler(event_stream=event_stream) + with patch.object( + handler, '_fail_in_flight_tasks', new_callable=AsyncMock + ) as fail_in_flight_tasks: + await handler.aclose(fail_in_flight=fail_in_flight) + + assert fail_in_flight_tasks.await_count == int(expected) + + +@pytest.mark.asyncio +async def test_aclose_gives_up_on_failing_in_flight_tasks_after_timeout( + caplog: pytest.LogCaptureFixture, +) -> None: + handler = make_handler(event_stream=InMemoryTaskEventStream()) + + async def hang() -> None: + await asyncio.sleep(10) + + with ( + patch.object(handler, '_fail_in_flight_tasks', side_effect=hang), + caplog.at_level(logging.WARNING), + ): + await asyncio.wait_for(handler.aclose(timeout=0.01), timeout=2) + + assert 'Timed out failing in-flight tasks' in caplog.text + + +@pytest.mark.asyncio +async def test_fail_in_flight_tasks_skips_idle_tasks_and_survives_errors() -> ( + None +): + handler = make_handler(event_stream=InMemoryTaskEventStream()) + running = MagicMock(task_id='running', request_in_flight=True) + running.call_context = make_context() + idle = MagicMock(task_id='idle', request_in_flight=False) + idle.call_context = make_context() + no_context = MagicMock(task_id='new', request_in_flight=True) + no_context.call_context = None + target = 'a2a.server.request_handlers.default_request_handler_v2.fail_task' + + with ( + patch.object( + handler._active_task_registry, # noqa: SLF001 + 'snapshot', + return_value=[running, idle, no_context], + ), + patch(target, side_effect=RuntimeError('store down')) as fail_task, + ): + await handler._fail_in_flight_tasks() # noqa: SLF001 + + assert [c.args[1] for c in fail_task.await_args_list] == ['running'] + await handler.aclose(fail_in_flight=False) diff --git a/tests/server/cluster/test_lease_multireplica.py b/tests/server/cluster/test_lease_multireplica.py new file mode 100644 index 000000000..a6ba80779 --- /dev/null +++ b/tests/server/cluster/test_lease_multireplica.py @@ -0,0 +1,742 @@ +import asyncio + +import pytest + +from a2a.server.cluster import StoredTask, TaskVersion, VersionedTaskStore +from a2a.server.cluster.recovery import ( + ABANDONED_TASK_REASON, + fail_abandoned_tasks, +) +from a2a.server.context import ServerCallContext +from a2a.server.events.event_queue import Event +from a2a.server.tasks.push_notification_sender import PushNotificationSender +from a2a.types.a2a_pb2 import ( + CancelTaskRequest, + GetTaskRequest, + ListTasksRequest, + ListTasksResponse, + SubscribeToTaskRequest, + Task, + TaskState, + TaskStatusUpdateEvent, +) +from a2a.utils.errors import UnsupportedOperationError + +from .conftest import ( + ControlledAgent, + InMemoryTaskLeaseStore, + InputRequiredThenCompleteAgent, + build_send_request, + make_replica, + wait_for_state, +) + + +class _SeverableStore(VersionedTaskStore): + """One replica's handle on a shared store, which can lose its connection.""" + + def __init__(self, store: VersionedTaskStore) -> None: + self._store = store + self.severed = False + + def _check(self) -> None: + if self.severed: + raise ConnectionError('database unreachable') + + async def save( + self, + task: Task, + *, + event: Event | None, + prev: Task | None, + prev_version: TaskVersion, + context: ServerCallContext, + ) -> TaskVersion: + self._check() + return await self._store.save( + task, + event=event, + prev=prev, + prev_version=prev_version, + context=context, + ) + + async def get( + self, task_id: str, context: ServerCallContext + ) -> StoredTask | None: + self._check() + return await self._store.get(task_id, context) + + async def list( + self, params: ListTasksRequest, context: ServerCallContext + ) -> ListTasksResponse: + self._check() + return await self._store.list(params, context) + + async def delete(self, task_id: str, context: ServerCallContext) -> None: + self._check() + await self._store.delete(task_id, context) + + +class _TrackedAgent(ControlledAgent): + """`ControlledAgent` that counts its runs and notes being cancelled.""" + + def __init__(self) -> None: + super().__init__() + self.runs = 0 + self.cancelled = asyncio.Event() + + async def execute(self, context, event_queue) -> None: # noqa: ANN001 + self.runs += 1 + try: + await super().execute(context, event_queue) + except asyncio.CancelledError: + self.cancelled.set() + raise + + +class _RecordingPushSender(PushNotificationSender): + def __init__(self) -> None: + self.sent: list[tuple[str, int]] = [] + + async def send_notification(self, task_id, event) -> None: # noqa: ANN001 + self.sent.append((task_id, event.status.state)) + + +async def _start_turn(replica, agent: ControlledAgent, store, context): # noqa: ANN001, ANN202 + """Starts a turn on `replica` and waits until the agent reports WORKING.""" + events: list = [] + + async def collect() -> None: + async for event in replica.on_message_send_stream( + build_send_request('go', message_id='m1'), context + ): + events.append(event) + + stream = asyncio.create_task(collect()) + await asyncio.wait_for(agent.working.wait(), timeout=5) + task_id = next(iter(replica._active_task_registry._active_tasks)) # noqa: SLF001 + await wait_for_state(store, task_id, TaskState.TASK_STATE_WORKING, context) + return stream, task_id, events + + +async def _crash(replica, task_id: str, store: _SeverableStore, leases) -> None: # noqa: ANN001 + """Stops a replica the way a killed process would: nothing more is written.""" + store.severed = True + leases.unreachable = True + active = await replica._active_task_registry.get(task_id) # noqa: SLF001 + tasks = [ + t + for t in ( + active._producer_task, # noqa: SLF001 + active._consumer_task, # noqa: SLF001 + active._heartbeat_task, # noqa: SLF001 + ) + if t is not None + ] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +async def _cancel(task: asyncio.Task) -> None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +async def _wait_until_released( + leases: InMemoryTaskLeaseStore, task_id: str +) -> None: + """A turn's lease is released once the consumer has saved its last event, + which can be just after the response for that turn was returned.""" + for _ in range(200): + if not leases.held(task_id): + return + await asyncio.sleep(0.01) + raise AssertionError(f'lease of {task_id} was not released') + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_follow_up_while_another_replica_runs_the_turn_is_rejected( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + replica_b = make_replica( + shared_store, shared_stream, agent, lease_store=leases.view() + ) + try: + stream, task_id, _ = await _start_turn( + replica_a, agent, shared_store, context + ) + stored = await shared_store.get(task_id, context) + assert stored is not None + + with pytest.raises( + UnsupportedOperationError, match='processed by another replica' + ): + await replica_b.on_message_send( + build_send_request( + 'again', + task_id=task_id, + context_id=stored.task.context_id, + message_id='m2', + ), + context, + ) + + assert agent.runs == 1 + rejected = await shared_store.get(task_id, context) + assert rejected is not None + assert rejected.task.status.state == TaskState.TASK_STATE_WORKING + + agent.release.set() + await asyncio.wait_for(stream, timeout=5) + final = await shared_store.get(task_id, context) + assert final is not None + assert final.task.status.state == TaskState.TASK_STATE_COMPLETED + await _wait_until_released(leases, task_id) + finally: + await replica_a.aclose() + await replica_b.aclose() + + +async def _wait_until_unregistered(replica, task_id: str) -> None: # noqa: ANN001 + for _ in range(200): + if await replica._active_task_registry.get(task_id) is None: # noqa: SLF001 + return + await asyncio.sleep(0.01) + raise AssertionError(f'{task_id} is still registered') + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +@pytest.mark.parametrize('interrupted_task_timeout', [None, 60.0]) +async def test_rejected_follow_up_leaves_no_local_task_behind( + shared_store, shared_stream, context, interrupted_task_timeout +) -> None: + """After rejecting a follow-up, B must not keep a local task: a later + subscription on B would tap it and never see A's events.""" + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + replica_b = make_replica( + shared_store, + shared_stream, + agent, + lease_store=leases.view(), + interrupted_task_timeout=interrupted_task_timeout, + ) + try: + stream, task_id, _ = await _start_turn( + replica_a, agent, shared_store, context + ) + stored = await shared_store.get(task_id, context) + assert stored is not None + with pytest.raises(UnsupportedOperationError): + await replica_b.on_message_send( + build_send_request( + 'again', + task_id=task_id, + context_id=stored.task.context_id, + message_id='m2', + ), + context, + ) + await _wait_until_unregistered(replica_b, task_id) + + states: list = [] + + async def follow() -> None: + async for event in replica_b.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), context + ): + if isinstance(event, Task | TaskStatusUpdateEvent): + states.append(event.status.state) + + follower = asyncio.create_task(follow()) + await asyncio.sleep(0.05) + agent.release.set() + await asyncio.wait_for(stream, timeout=5) + await asyncio.wait_for(follower, timeout=5) + assert states[-1] == TaskState.TASK_STATE_COMPLETED + assert agent.runs == 1 + finally: + await replica_a.aclose() + await replica_b.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_task_released_while_waiting_for_input_holds_no_lease( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + replica = make_replica( + shared_store, + shared_stream, + InputRequiredThenCompleteAgent(), + lease_store=leases, + interrupted_task_timeout=0.05, + ) + try: + first = await replica.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert first.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + await _wait_until_unregistered(replica, first.id) + assert not leases.held(first.id) + + leases.clock.advance(120) + assert await fail_abandoned_tasks(leases, shared_store) == 0 + + resumed = await replica.on_message_send( + build_send_request( + 'a1', + task_id=first.id, + context_id=first.context_id, + message_id='m2', + ), + context, + ) + assert resumed.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + await _wait_until_released(leases, first.id) + finally: + await replica.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_multiturn_across_replicas_releases_each_turns_lease( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + agent = InputRequiredThenCompleteAgent() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + replica_b = make_replica( + shared_store, shared_stream, agent, lease_store=leases.view() + ) + try: + first = await replica_a.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert first.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + await _wait_until_released(leases, first.id) + + for replica, message_id in ((replica_b, 'm2'), (replica_a, 'm3')): + result = await replica.on_message_send( + build_send_request( + 'answer', + task_id=first.id, + context_id=first.context_id, + message_id=message_id, + ), + context, + ) + await _wait_until_released(leases, first.id) + assert result.status.state == TaskState.TASK_STATE_COMPLETED + finally: + await replica_a.aclose() + await replica_b.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_canceled_turn_releases_its_lease( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + replica = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + try: + stream, task_id, _ = await _start_turn( + replica, agent, shared_store, context + ) + assert leases.held(task_id) + + await replica.on_cancel_task(CancelTaskRequest(id=task_id), context) + await _cancel(stream) + + await _wait_until_released(leases, task_id) + finally: + await replica.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_lease_left_behind_by_finished_turn_does_not_fail_it( + shared_store, shared_stream, context +) -> None: + """A lease that could not be released at the end of a turn expires, and + the sweep that claims it leaves the finished task as it is.""" + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + leases_a = leases.view() + replica = make_replica( + shared_store, shared_stream, agent, lease_store=leases_a + ) + try: + stream, task_id, _ = await _start_turn( + replica, agent, shared_store, context + ) + leases_a.unreachable = True + agent.release.set() + await asyncio.wait_for(stream, timeout=5) + await wait_for_state( + shared_store, task_id, TaskState.TASK_STATE_COMPLETED, context + ) + + leases.clock.advance(2) + assert await fail_abandoned_tasks(leases, shared_store) == 0 + + final = await shared_store.get(task_id, context) + assert final is not None + assert final.task.status.state == TaskState.TASK_STATE_COMPLETED + assert not leases.held(task_id) + finally: + leases_a.unreachable = False + await replica.aclose() + + +class _SlowReleaseLeases(InMemoryTaskLeaseStore): + """Releases take a moment, like a database round trip.""" + + async def release(self, lease) -> None: # noqa: ANN001 + await asyncio.sleep(0.05) + await super().release(lease) + + def view(self) -> InMemoryTaskLeaseStore: + other = _SlowReleaseLeases(self.clock, self._owner_resolver) + other._leases = self._leases # noqa: SLF001 + return other + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_follow_up_sent_as_soon_as_input_is_requested_is_accepted( + shared_store, shared_stream, context +) -> None: + """A returns INPUT_REQUIRED a moment before it releases the lease; a + follow-up that reaches B in between waits for the lease.""" + leases = _SlowReleaseLeases() + agent = InputRequiredThenCompleteAgent() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + replica_b = make_replica( + shared_store, shared_stream, agent, lease_store=leases.view() + ) + try: + first = await replica_a.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert first.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + assert leases.held(first.id) + + await replica_b.on_message_send( + build_send_request( + 'a1', + task_id=first.id, + context_id=first.context_id, + message_id='m2', + ), + context, + ) + stored = await shared_store.get(first.id, context) + assert stored is not None + assert 'm2' in {m.message_id for m in stored.task.history} + finally: + await replica_a.aclose() + await replica_b.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_follow_up_to_waiting_task_with_a_stuck_lease_is_rejected( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + agent = InputRequiredThenCompleteAgent() + replica = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + try: + first = await replica.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + await _wait_until_released(leases, first.id) + # Held by a replica that died after saving INPUT_REQUIRED. + assert await leases.acquire(first.id, context, 60) is not None + + with pytest.raises(UnsupportedOperationError): + await replica.on_message_send( + build_send_request( + 'a1', + task_id=first.id, + context_id=first.context_id, + message_id='m2', + ), + context, + ) + finally: + await replica.aclose() + + +class _FailingReleaseLeases(InMemoryTaskLeaseStore): + """A replica's view of `leases` whose releases fail.""" + + def __init__(self, leases: InMemoryTaskLeaseStore) -> None: + super().__init__(leases.clock) + self._leases = leases._leases # noqa: SLF001 + + async def release(self, lease) -> None: # noqa: ANN001 + raise ConnectionError('lease store unreachable') + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_lease_left_behind_by_turn_waiting_for_input_does_not_fail_it( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + replica = make_replica( + shared_store, + shared_stream, + InputRequiredThenCompleteAgent(), + lease_store=_FailingReleaseLeases(leases), + ) + try: + first = await replica.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert first.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + await asyncio.sleep(0.05) + + leases.clock.advance(2) + assert await fail_abandoned_tasks(leases, shared_store) == 0 + + stored = await shared_store.get(first.id, context) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + assert not leases.held(first.id) + finally: + await replica.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_task_waiting_for_input_is_never_recovered( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + replica = make_replica( + shared_store, + shared_stream, + InputRequiredThenCompleteAgent(), + lease_store=leases, + ) + try: + result = await replica.on_message_send( + build_send_request('q1', message_id='m1'), context + ) + assert result.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + await _wait_until_released(leases, result.id) + + leases.clock.advance(120) + assert await fail_abandoned_tasks(leases, shared_store) == 0 + + stored = await shared_store.get(result.id, context) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + finally: + await replica.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_task_of_crashed_replica_is_failed_after_its_lease_expires( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + store_a = _SeverableStore(shared_store) + leases_a = leases.view() + replica_a = make_replica( + store_a, shared_stream, agent, lease_store=leases_a + ) + replica_b = make_replica(shared_store, shared_stream, agent) + push_sender = _RecordingPushSender() + try: + stream, task_id, _ = await _start_turn( + replica_a, agent, shared_store, context + ) + b_states: list = [] + + async def follow_on_b() -> None: + async for event in replica_b.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), context + ): + b_states.append(event.status.state) + + follow = asyncio.create_task(follow_on_b()) + await asyncio.sleep(0.05) + + await _crash(replica_a, task_id, store_a, leases_a) + await _cancel(stream) + + assert await fail_abandoned_tasks(leases, shared_store) == 0 + leases.clock.advance(2) + assert ( + await fail_abandoned_tasks( + leases, + shared_store, + push_sender=push_sender, + event_stream=shared_stream, + ) + == 1 + ) + + await asyncio.wait_for(follow, timeout=5) + assert b_states[-1] == TaskState.TASK_STATE_FAILED + task = await replica_b.on_get_task(GetTaskRequest(id=task_id), context) + assert task is not None + assert task.status.state == TaskState.TASK_STATE_FAILED + assert task.status.message.parts[0].text == ABANDONED_TASK_REASON + assert push_sender.sent == [(task_id, TaskState.TASK_STATE_FAILED)] + assert not leases.held(task_id) + finally: + await replica_b.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_replica_that_lost_its_lease_stops_its_agent( + shared_store, shared_stream, context +) -> None: + """A replica cut off from the lease store is recovered as abandoned. When + it reconnects, its renewal fails and it stops; its late events are + dropped, so the task stays FAILED.""" + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + leases_a = leases.view() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases_a + ) + try: + stream, task_id, _ = await _start_turn( + replica_a, agent, shared_store, context + ) + + leases_a.unreachable = True + await asyncio.sleep(0.05) + leases.clock.advance(2) + assert await fail_abandoned_tasks(leases, shared_store) == 1 + leases_a.unreachable = False + + await asyncio.wait_for(agent.cancelled.wait(), timeout=5) + await asyncio.wait_for(stream, timeout=5) + final = await shared_store.get(task_id, context) + assert final is not None + assert final.task.status.state == TaskState.TASK_STATE_FAILED + finally: + await replica_a.aclose() + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_turn_cut_short_in_process_is_recovered( + shared_store, shared_stream, context +) -> None: + """Closing without failing in-flight turns must not drop the lease of a + task left WORKING, or nothing could recover it.""" + leases = InMemoryTaskLeaseStore() + agent = _TrackedAgent() + replica_a = make_replica( + shared_store, shared_stream, agent, lease_store=leases + ) + stream, task_id, _ = await _start_turn( + replica_a, agent, shared_store, context + ) + + await replica_a.aclose(fail_in_flight=False) + await _cancel(stream) + stored = await shared_store.get(task_id, context) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_WORKING + + leases.clock.advance(2) + assert await fail_abandoned_tasks(leases, shared_store) == 1 + final = await shared_store.get(task_id, context) + assert final is not None + assert final.task.status.state == TaskState.TASK_STATE_FAILED + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_graceful_close_fails_running_turns_only( + shared_store, shared_stream, context +) -> None: + leases = InMemoryTaskLeaseStore() + push_sender = _RecordingPushSender() + waiting = make_replica( + shared_store, + shared_stream, + InputRequiredThenCompleteAgent(), + lease_store=leases, + ) + agent = _TrackedAgent() + running = make_replica( + shared_store, + shared_stream, + agent, + push_sender=push_sender, + lease_store=leases.view(), + ) + observer = make_replica(shared_store, shared_stream, agent) + try: + paused = await waiting.on_message_send( + build_send_request('q1', message_id='p1'), context + ) + stream, task_id, _ = await _start_turn( + running, agent, shared_store, context + ) + states: list = [] + + async def follow() -> None: + async for event in observer.on_subscribe_to_task( + SubscribeToTaskRequest(id=task_id), context + ): + states.append(event.status.state) + + follower = asyncio.create_task(follow()) + await asyncio.sleep(0.05) + + await waiting.aclose() + await running.aclose() + await _cancel(stream) + await asyncio.wait_for(follower, timeout=5) + + failed = await shared_store.get(task_id, context) + assert failed is not None + assert failed.task.status.state == TaskState.TASK_STATE_FAILED + assert 'shut down' in failed.task.status.message.parts[0].text + assert states[-1] == TaskState.TASK_STATE_FAILED + assert push_sender.sent[-1] == (task_id, TaskState.TASK_STATE_FAILED) + still_paused = await shared_store.get(paused.id, context) + assert still_paused is not None + assert ( + still_paused.task.status.state + == TaskState.TASK_STATE_INPUT_REQUIRED + ) + finally: + await observer.aclose() diff --git a/tests/server/cluster/test_recovery.py b/tests/server/cluster/test_recovery.py new file mode 100644 index 000000000..f51bf792c --- /dev/null +++ b/tests/server/cluster/test_recovery.py @@ -0,0 +1,300 @@ +import asyncio + +import pytest + +from a2a.server.cluster import ConcurrentTaskModificationError, TaskVersion +from a2a.server.cluster.recovery import ( + ABANDONED_TASK_REASON, + AbandonedTaskSweeper, + context_for_owner, + fail_abandoned_tasks, + fail_task, +) +from a2a.server.owner_resolver import resolve_user_scope +from a2a.server.tasks.push_notification_sender import PushNotificationSender +from a2a.types.a2a_pb2 import ( + Message, + Part, + Role, + Task, + TaskState, + TaskStatus, + TaskStatusUpdateEvent, +) + +from .conftest import ( + InMemoryTaskEventStream, + InMemoryTaskLeaseStore, + VersionedInMemoryTaskStore, + make_context, +) + + +ALICE = make_context('alice') + + +class _RecordingPushSender(PushNotificationSender): + def __init__(self) -> None: + self.sent: list[tuple[str, int]] = [] + + async def send_notification(self, task_id, event) -> None: # noqa: ANN001 + self.sent.append((task_id, event.status.state)) + + +class _RaisingPushSender(PushNotificationSender): + async def send_notification(self, task_id, event) -> None: # noqa: ANN001 + raise RuntimeError('push failed') + + +class _ConflictingStore(VersionedInMemoryTaskStore): + """Rejects the first `conflicts` saves, as if another writer won them.""" + + def __init__(self, conflicts: int) -> None: + super().__init__() + self.conflicts = conflicts + + async def save(self, task: Task, **kwargs) -> TaskVersion: # noqa: ANN003 + if self.conflicts: + self.conflicts -= 1 + raise ConcurrentTaskModificationError(task.id) + return await super().save(task, **kwargs) + + +async def _store_task( + store: VersionedInMemoryTaskStore, + state: TaskState = TaskState.TASK_STATE_WORKING, + task_id: str = 'task-1', +) -> Task: + task = Task( + id=task_id, + context_id='ctx-1', + status=TaskStatus( + state=state, + message=Message( + message_id='agent-1', + role=Role.ROLE_AGENT, + parts=[Part(text='thinking')], + ), + ), + ) + await store.save( + task, + event=None, + prev=None, + prev_version=TaskVersion.MISSING, + context=ALICE, + ) + return task + + +def test_context_for_owner_maps_back_to_the_owner() -> None: + assert resolve_user_scope(context_for_owner('alice')) == 'alice' + assert context_for_owner('alice').user.is_authenticated + assert not context_for_owner('').user.is_authenticated + + +@pytest.mark.asyncio +async def test_fail_task_writes_failed_and_notifies() -> None: + store = VersionedInMemoryTaskStore() + stream = InMemoryTaskEventStream() + push_sender = _RecordingPushSender() + await _store_task(store) + subscription = stream.subscribe('task-1', after=TaskVersion.MISSING) + next_event = asyncio.ensure_future(anext(subscription)) + await asyncio.sleep(0) + + failed = await fail_task( + store, + 'task-1', + ALICE, + reason='replica gone', + push_sender=push_sender, + event_stream=stream, + ) + + assert failed is not None + stored = await store.get('task-1', ALICE) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_FAILED + assert stored.task.status.message.parts[0].text == 'replica gone' + assert stored.task.status.HasField('timestamp') + assert [m.message_id for m in stored.task.history] == ['agent-1'] + published = await asyncio.wait_for(next_event, timeout=1) + assert isinstance(published.event, TaskStatusUpdateEvent) + assert published.event.status.state == TaskState.TASK_STATE_FAILED + assert published.version == stored.version + assert push_sender.sent == [('task-1', TaskState.TASK_STATE_FAILED)] + await subscription.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'state', + [ + TaskState.TASK_STATE_COMPLETED, + TaskState.TASK_STATE_FAILED, + TaskState.TASK_STATE_INPUT_REQUIRED, + TaskState.TASK_STATE_AUTH_REQUIRED, + ], +) +async def test_fail_task_leaves_a_finished_turn_alone(state: TaskState) -> None: + store = VersionedInMemoryTaskStore() + push_sender = _RecordingPushSender() + await _store_task(store, state) + + result = await fail_task( + store, 'task-1', ALICE, reason='x', push_sender=push_sender + ) + + assert result is None + stored = await store.get('task-1', ALICE) + assert stored is not None + assert stored.task.status.state == state + assert push_sender.sent == [] + + +@pytest.mark.asyncio +async def test_fail_task_ignores_a_missing_task() -> None: + store = VersionedInMemoryTaskStore() + + assert await fail_task(store, 'nope', ALICE, reason='x') is None + + +@pytest.mark.asyncio +async def test_fail_task_rereads_after_one_conflict() -> None: + store = _ConflictingStore(conflicts=0) + await _store_task(store) + store.conflicts = 1 + + assert await fail_task(store, 'task-1', ALICE, reason='x') is not None + stored = await store.get('task-1', ALICE) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_FAILED + + +@pytest.mark.asyncio +async def test_fail_task_gives_up_after_repeated_conflicts() -> None: + store = _ConflictingStore(conflicts=0) + await _store_task(store) + store.conflicts = 2 + + assert await fail_task(store, 'task-1', ALICE, reason='x') is None + + +@pytest.mark.asyncio +async def test_fail_task_survives_a_raising_push_sender() -> None: + store = VersionedInMemoryTaskStore() + await _store_task(store) + + result = await fail_task( + store, 'task-1', ALICE, reason='x', push_sender=_RaisingPushSender() + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_fail_abandoned_tasks_fails_owned_task_and_releases_lease() -> ( + None +): + store = VersionedInMemoryTaskStore() + leases = InMemoryTaskLeaseStore() + await _store_task(store) + await leases.acquire('task-1', ALICE, ttl=10) + leases.clock.advance(11) + + assert await fail_abandoned_tasks(leases, store) == 1 + + stored = await store.get('task-1', ALICE) + assert stored is not None + assert stored.task.status.state == TaskState.TASK_STATE_FAILED + assert stored.task.status.message.parts[0].text == ABANDONED_TASK_REASON + assert not leases.held('task-1') + assert await leases.claim_expired(ttl=60, limit=10) == [] + + +@pytest.mark.asyncio +async def test_fail_abandoned_tasks_skips_live_and_finished_tasks() -> None: + store = VersionedInMemoryTaskStore() + leases = InMemoryTaskLeaseStore() + await _store_task(store, TaskState.TASK_STATE_COMPLETED, 'done') + await _store_task(store, TaskState.TASK_STATE_WORKING, 'live') + await leases.acquire('done', ALICE, ttl=10) + leases.clock.advance(11) + await leases.acquire('live', ALICE, ttl=60) + + assert await fail_abandoned_tasks(leases, store) == 0 + + assert not leases.held('done') + assert leases.held('live') + live = await store.get('live', ALICE) + assert live is not None + assert live.task.status.state == TaskState.TASK_STATE_WORKING + + +@pytest.mark.asyncio +async def test_fail_abandoned_tasks_keeps_the_claim_when_recovery_fails() -> ( + None +): + store = VersionedInMemoryTaskStore() + leases = InMemoryTaskLeaseStore() + await _store_task(store) + await leases.acquire('task-1', ALICE, ttl=10) + leases.clock.advance(11) + + def unreachable(owner: str): # noqa: ANN202 + raise ConnectionError('task store unreachable') + + assert ( + await fail_abandoned_tasks( + leases, store, context_factory=unreachable, claim_ttl=30 + ) + == 0 + ) + assert leases.held('task-1') + leases.clock.advance(31) + assert await fail_abandoned_tasks(leases, store) == 1 + + +@pytest.mark.asyncio +async def test_fail_abandoned_tasks_uses_the_context_factory() -> None: + store = VersionedInMemoryTaskStore() + leases = InMemoryTaskLeaseStore() + await _store_task(store) + await leases.acquire('task-1', ALICE, ttl=10) + leases.clock.advance(11) + owners: list[str] = [] + + def factory(owner: str): # noqa: ANN202 + owners.append(owner) + return ALICE + + assert await fail_abandoned_tasks(leases, store, context_factory=factory) + assert owners == ['alice'] + + +@pytest.mark.asyncio +@pytest.mark.timeout(5) +async def test_sweeper_recovers_in_the_background_and_stops() -> None: + store = VersionedInMemoryTaskStore() + leases = InMemoryTaskLeaseStore() + await _store_task(store) + await leases.acquire('task-1', ALICE, ttl=10) + leases.clock.advance(11) + leases.unreachable = True + sweeper = AbandonedTaskSweeper(leases, store, interval=0.01) + + sweeper.start() + sweeper.start() + await asyncio.sleep(0.05) + leases.unreachable = False + for _ in range(100): + stored = await store.get('task-1', ALICE) + assert stored is not None + if stored.task.status.state == TaskState.TASK_STATE_FAILED: + break + await asyncio.sleep(0.01) + await sweeper.aclose() + await sweeper.aclose() + + assert stored.task.status.state == TaskState.TASK_STATE_FAILED