diff --git a/src/a2a/server/agent_execution/active_task_registry.py b/src/a2a/server/agent_execution/active_task_registry.py index ab7d6a11c..39f8ff0ef 100644 --- a/src/a2a/server/agent_execution/active_task_registry.py +++ b/src/a2a/server/agent_execution/active_task_registry.py @@ -2,6 +2,7 @@ import asyncio import logging +import threading from typing import TYPE_CHECKING @@ -33,7 +34,7 @@ def __init__( self._task_store = task_store self._push_sender = push_sender self._active_tasks: dict[str, ActiveTask] = {} - self._lock = asyncio.Lock() + self._lock = threading.RLock() self._cleanup_tasks: set[asyncio.Task[None]] = set() self._closed = False @@ -46,7 +47,7 @@ async def get_or_create( initial_message: Message | None = None, ) -> ActiveTask: """Retrieves an existing ActiveTask or creates a new one.""" - async with self._lock: + with self._lock: if self._closed: raise RuntimeError('ActiveTaskRegistry is closed') if task_id in self._active_tasks: @@ -83,13 +84,13 @@ def _on_active_task_cleanup(self, active_task: ActiveTask) -> None: task.add_done_callback(self._cleanup_tasks.discard) async def _remove_task(self, task_id: str) -> None: - async with self._lock: + with self._lock: self._active_tasks.pop(task_id, None) logger.debug('Removed active task for %s from registry', task_id) async def get(self, task_id: str) -> ActiveTask | None: """Retrieves an existing task.""" - async with self._lock: + with self._lock: return self._active_tasks.get(task_id) async def aclose(self) -> None: @@ -102,13 +103,11 @@ async def aclose(self) -> None: multiple times. The close flag is set and the active-task snapshot is taken under - ``_lock``, and the lock is then released before awaiting, because - ``_remove_task`` re-acquires ``_lock``; holding it while draining - would deadlock. Marking closed under the same lock prevents a - concurrent ``get_or_create`` from registering a task that the drain - would miss. + ``_lock``, then the lock is released before awaiting the drain. Marking + closed under the same lock prevents a concurrent ``get_or_create`` from + registering a task that the drain would miss. """ - async with self._lock: + with self._lock: self._closed = True active_tasks = list(self._active_tasks.values()) @@ -125,5 +124,5 @@ async def aclose(self) -> None: if cleanup_tasks: await asyncio.gather(*cleanup_tasks, return_exceptions=True) - async with self._lock: + with self._lock: self._active_tasks.clear() diff --git a/src/a2a/server/events/in_memory_queue_manager.py b/src/a2a/server/events/in_memory_queue_manager.py index 0beb354f9..493ac18ea 100644 --- a/src/a2a/server/events/in_memory_queue_manager.py +++ b/src/a2a/server/events/in_memory_queue_manager.py @@ -1,4 +1,4 @@ -import asyncio +import threading from a2a.server.events.event_queue import EventQueueLegacy from a2a.server.events.queue_manager import ( @@ -24,7 +24,7 @@ class InMemoryQueueManager(QueueManager): def __init__(self) -> None: """Initializes the InMemoryQueueManager.""" self._task_queue: dict[str, EventQueueLegacy] = {} - self._lock = asyncio.Lock() + self._lock = threading.RLock() async def add(self, task_id: str, queue: EventQueueLegacy) -> None: """Adds a new event queue for a task ID. @@ -32,7 +32,7 @@ async def add(self, task_id: str, queue: EventQueueLegacy) -> None: Raises: TaskQueueExists: If a queue for the given `task_id` already exists. """ - async with self._lock: + with self._lock: if task_id in self._task_queue: raise TaskQueueExists self._task_queue[task_id] = queue @@ -43,10 +43,8 @@ async def get(self, task_id: str) -> EventQueueLegacy | None: Returns: The `EventQueueLegacy` instance for the `task_id`, or `None` if not found. """ - async with self._lock: - if task_id not in self._task_queue: - return None - return self._task_queue[task_id] + with self._lock: + return self._task_queue.get(task_id) async def tap(self, task_id: str) -> EventQueueLegacy | None: """Taps the event queue for a task ID to create a child queue. @@ -54,10 +52,11 @@ async def tap(self, task_id: str) -> EventQueueLegacy | None: Returns: A new child `EventQueueLegacy` instance, or `None` if the task ID is not found. """ - async with self._lock: - if task_id not in self._task_queue: - return None - return await self._task_queue[task_id].tap() + with self._lock: + queue = self._task_queue.get(task_id) + if queue is None: + return None + return await queue.tap() async def close(self, task_id: str) -> None: """Closes and removes the event queue for a task ID. @@ -65,11 +64,11 @@ async def close(self, task_id: str) -> None: Raises: NoTaskQueue: If no queue exists for the given `task_id`. """ - async with self._lock: - if task_id not in self._task_queue: - raise NoTaskQueue - queue = self._task_queue.pop(task_id) - await queue.close() + with self._lock: + queue = self._task_queue.pop(task_id, None) + if queue is None: + raise NoTaskQueue + await queue.close() async def create_or_tap(self, task_id: str) -> EventQueueLegacy: """Creates a new event queue for a task ID if one doesn't exist, otherwise taps the existing one. @@ -77,9 +76,10 @@ async def create_or_tap(self, task_id: str) -> EventQueueLegacy: Returns: A new or child `EventQueueLegacy` instance for the `task_id`. """ - async with self._lock: - if task_id not in self._task_queue: + with self._lock: + queue = self._task_queue.get(task_id) + if queue is None: queue = EventQueueLegacy() self._task_queue[task_id] = queue return queue - return await self._task_queue[task_id].tap() + return await queue.tap() diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index ef61dcca7..fbd050a41 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/src/a2a/server/request_handlers/default_request_handler.py @@ -1,5 +1,6 @@ import asyncio import logging +import threading from collections.abc import AsyncGenerator, Awaitable, Callable from typing import cast @@ -133,7 +134,7 @@ def __init__( # noqa: PLR0913 ) # TODO: Likely want an interface for managing this, like AgentExecutionManager. self._running_agents = {} - self._running_agents_lock = asyncio.Lock() + self._running_agents_lock = threading.RLock() # Tracks background tasks (e.g., deferred cleanups) to avoid orphaning # asyncio tasks and to surface unexpected exceptions. self._background_tasks = set() @@ -462,7 +463,7 @@ async def _register_producer( self, task_id: str, producer_task: asyncio.Task ) -> None: """Registers the agent execution task with the handler.""" - async with self._running_agents_lock: + with self._running_agents_lock: self._running_agents[task_id] = producer_task def _track_background_task(self, task: asyncio.Task) -> None: @@ -501,7 +502,7 @@ async def _cleanup_producer( 'Producer task %s was cancelled during cleanup', task_id ) await self._queue_manager.close(task_id) - async with self._running_agents_lock: + with self._running_agents_lock: self._running_agents.pop(task_id, None) @validate_request_params diff --git a/src/a2a/server/tasks/inmemory_push_notification_config_store.py b/src/a2a/server/tasks/inmemory_push_notification_config_store.py index 19e35074a..3fdccf53e 100644 --- a/src/a2a/server/tasks/inmemory_push_notification_config_store.py +++ b/src/a2a/server/tasks/inmemory_push_notification_config_store.py @@ -1,5 +1,5 @@ -import asyncio import logging +import threading from a2a.server.context import ServerCallContext from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope @@ -24,7 +24,7 @@ def __init__( owner_resolver: OwnerResolver = resolve_user_scope, ) -> None: """Initializes the InMemoryPushNotificationConfigStore.""" - self.lock = asyncio.Lock() + self.lock = threading.RLock() self._push_notification_infos: dict[ str, dict[str, list[TaskPushNotificationConfig]] ] = {} @@ -45,7 +45,7 @@ async def set_info( owner = self.owner_resolver(context) if owner not in self._push_notification_infos: self._push_notification_infos[owner] = {} - async with self.lock: + with self.lock: owner_infos = self._push_notification_infos[owner] if task_id not in owner_infos: owner_infos[task_id] = [] @@ -77,7 +77,7 @@ async def get_info( Used by the user-callable read endpoints. """ owner = self.owner_resolver(context) - async with self.lock: + with self.lock: owner_infos = self._get_owner_push_notification_infos(owner) return list(owner_infos.get(task_id, [])) @@ -89,7 +89,7 @@ async def get_info_for_dispatch( Used by the push-notification dispatch path. """ - async with self.lock: + with self.lock: results: list[TaskPushNotificationConfig] = [] for all_configs in self._push_notification_infos.values(): results.extend(all_configs.get(task_id, [])) @@ -107,7 +107,7 @@ async def delete_info( If config_id is None, all configurations for the task for the owner are deleted. """ owner = self.owner_resolver(context) - async with self.lock: + with self.lock: owner_infos = self._get_owner_push_notification_infos(owner) if task_id not in owner_infos: logger.warning( diff --git a/src/a2a/server/tasks/inmemory_task_store.py b/src/a2a/server/tasks/inmemory_task_store.py index 75d2269bc..2e1328ba2 100644 --- a/src/a2a/server/tasks/inmemory_task_store.py +++ b/src/a2a/server/tasks/inmemory_task_store.py @@ -1,5 +1,5 @@ -import asyncio import logging +import threading from a2a.server.context import ServerCallContext from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope @@ -29,7 +29,7 @@ def __init__( """Initializes the internal _InMemoryTaskStoreImpl.""" logger.debug('Initializing _InMemoryTaskStoreImpl') self.tasks: dict[str, dict[str, Task]] = {} - self.lock = asyncio.Lock() + self.lock = threading.RLock() self.owner_resolver = owner_resolver def _get_owner_tasks(self, owner: str) -> dict[str, Task]: @@ -41,7 +41,7 @@ async def save(self, task: Task, context: ServerCallContext) -> None: if owner not in self.tasks: self.tasks[owner] = {} - async with self.lock: + with self.lock: self.tasks[owner][task.id] = task logger.debug( 'Task %s for owner %s saved successfully.', task.id, owner @@ -52,7 +52,7 @@ async def get( ) -> Task | None: """Retrieves a task from the in-memory store by ID, for the given owner.""" owner = self.owner_resolver(context) - async with self.lock: + with self.lock: logger.debug( 'Attempting to get task with id: %s for owner: %s', task_id, @@ -81,7 +81,7 @@ async def list( owner = self.owner_resolver(context) logger.debug('Listing tasks for owner %s with params %s', owner, params) - async with self.lock: + with self.lock: owner_tasks = self._get_owner_tasks(owner) tasks = list(owner_tasks.values()) @@ -157,7 +157,7 @@ async def list( async def delete(self, task_id: str, context: ServerCallContext) -> None: """Deletes a task from the in-memory store by ID, for the given owner.""" owner = self.owner_resolver(context) - async with self.lock: + with self.lock: logger.debug( 'Attempting to delete task with id: %s for owner %s', task_id, diff --git a/tests/server/events/test_inmemory_queue_manager.py b/tests/server/events/test_inmemory_queue_manager.py index 9716b13bf..c24755d68 100644 --- a/tests/server/events/test_inmemory_queue_manager.py +++ b/tests/server/events/test_inmemory_queue_manager.py @@ -1,4 +1,5 @@ import asyncio +import threading from unittest.mock import MagicMock @@ -31,7 +32,8 @@ def event_queue(self) -> MagicMock: async def test_init(self, queue_manager: InMemoryQueueManager) -> None: """Test that the InMemoryQueueManager initializes with empty task queue and a lock.""" assert queue_manager._task_queue == {} - assert isinstance(queue_manager._lock, asyncio.Lock) + assert not isinstance(queue_manager._lock, asyncio.Lock) + assert isinstance(queue_manager._lock, type(threading.RLock())) @pytest.mark.asyncio async def test_add_new_queue( diff --git a/tests/server/test_cross_event_loop_locks.py b/tests/server/test_cross_event_loop_locks.py new file mode 100644 index 000000000..edbac80d7 --- /dev/null +++ b/tests/server/test_cross_event_loop_locks.py @@ -0,0 +1,260 @@ +"""Regression tests: in-memory server singletons are safe across event loops.""" + +import asyncio +import threading + +from unittest.mock import MagicMock + +import pytest + +from a2a.server.agent_execution.active_task_registry import ActiveTaskRegistry +from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager +from a2a.server.request_handlers.default_request_handler import ( + LegacyRequestHandler, +) +from a2a.server.tasks.inmemory_push_notification_config_store import ( + InMemoryPushNotificationConfigStore, +) +from a2a.server.tasks.inmemory_task_store import _InMemoryTaskStoreImpl + + +_RLOCK_TYPE = type(threading.RLock()) + + +class _PersistentLoop: + """An asyncio event loop running forever on its own daemon thread.""" + + def __init__(self) -> None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def _run(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def run(self, coro, timeout: float = 10): + """Runs coro to completion on this loop and returns its result.""" + return asyncio.run_coroutine_threadsafe(coro, self._loop).result( + timeout=timeout + ) + + def close(self) -> None: + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(timeout=10) + + +@pytest.fixture +def two_loops(): + loop_a = _PersistentLoop() + loop_b = _PersistentLoop() + try: + yield loop_a, loop_b + finally: + loop_a.close() + loop_b.close() + + +# --- deterministic invariant: the guard is a threading lock, not asyncio --- + + +def test_queue_manager_lock_is_threading_lock() -> None: + manager = InMemoryQueueManager() + assert not isinstance(manager._lock, asyncio.Lock) + assert isinstance(manager._lock, _RLOCK_TYPE) + + +def test_active_task_registry_lock_is_threading_lock() -> None: + registry = ActiveTaskRegistry( + agent_executor=MagicMock(), task_store=MagicMock() + ) + assert not isinstance(registry._lock, asyncio.Lock) + assert isinstance(registry._lock, _RLOCK_TYPE) + + +def test_inmemory_task_store_lock_is_threading_lock() -> None: + store = _InMemoryTaskStoreImpl() + assert not isinstance(store.lock, asyncio.Lock) + assert isinstance(store.lock, _RLOCK_TYPE) + + +def test_push_config_store_lock_is_threading_lock() -> None: + store = InMemoryPushNotificationConfigStore() + assert not isinstance(store.lock, asyncio.Lock) + assert isinstance(store.lock, _RLOCK_TYPE) + + +def test_request_handler_running_agents_lock_is_threading_lock() -> None: + handler = LegacyRequestHandler(MagicMock(), MagicMock(), MagicMock()) + assert not isinstance(handler._running_agents_lock, asyncio.Lock) + assert isinstance(handler._running_agents_lock, _RLOCK_TYPE) + + +# --- functional: the real objects work when driven from two event loops --- + + +def test_queue_manager_across_event_loops(two_loops) -> None: + loop_a, loop_b = two_loops + manager = InMemoryQueueManager() + + # Task 1 is created on loop A (the first loop to touch the manager's guard). + loop_a.run(manager.create_or_tap('task-1')) + + # Task 2 is created on loop B, and task 1 is tapped/closed from loop B -- + # exercising the shared guard from a different loop than created it. + assert loop_b.run(manager.create_or_tap('task-2')) is not None + assert loop_b.run(manager.tap('task-1')) is not None + assert loop_a.run(manager.get('task-2')) is not None + loop_b.run(manager.close('task-1')) + assert loop_a.run(manager.get('task-1')) is None + + +def test_push_config_store_across_event_loops(two_loops) -> None: + loop_a, loop_b = two_loops + store = InMemoryPushNotificationConfigStore() + + # get_info_for_dispatch takes no ServerCallContext, so it exercises the + # guard directly from each loop. + assert loop_a.run(store.get_info_for_dispatch('task-1')) == [] + assert loop_b.run(store.get_info_for_dispatch('task-1')) == [] + + +def test_queue_manager_interleaved_across_loops(two_loops) -> None: + loop_a, loop_b = two_loops + manager = InMemoryQueueManager() + for i in range(10): + loop = loop_a if i % 2 == 0 else loop_b + loop.run(manager.create_or_tap(f'task-{i}')) + for i in range(10): + # Read each task from the opposite loop that created it. + loop = loop_b if i % 2 == 0 else loop_a + assert loop.run(manager.get(f'task-{i}')) is not None + + +# --- contended acquisition across loops (deterministic; the discriminator) --- + + +def _contended_acquire(lock, two_loops) -> None: + """Holds ``lock`` on loop A while loop B acquires it (a genuinely contended + cross-loop acquisition). With an ``asyncio.Lock`` this raises + ``RuntimeError: bound to a different event loop`` (or deadlocks); with a + ``threading.RLock`` it serializes and completes.""" + loop_a, loop_b = two_loops + a_holds = threading.Event() + b_queued = threading.Event() + + async def holder() -> None: + lock.acquire() + try: + a_holds.set() + b_queued.wait(3) + finally: + lock.release() + + async def contender() -> None: + a_holds.wait(3) + b_queued.set() + # Acquire on loop B's thread without freezing the loop. + await asyncio.get_event_loop().run_in_executor( + None, _acquire_release, lock + ) + + fut_a = asyncio.run_coroutine_threadsafe(holder(), loop_a._loop) + fut_b = asyncio.run_coroutine_threadsafe(contender(), loop_b._loop) + fut_b.result(timeout=5) + fut_a.result(timeout=5) + + +def _acquire_release(lock) -> None: + lock.acquire() + lock.release() + + +def test_queue_manager_lock_contended_across_loops(two_loops) -> None: + _contended_acquire(InMemoryQueueManager()._lock, two_loops) + + +def test_task_store_lock_contended_across_loops(two_loops) -> None: + _contended_acquire(_InMemoryTaskStoreImpl().lock, two_loops) + + +def test_active_task_registry_lock_contended_across_loops(two_loops) -> None: + registry = ActiveTaskRegistry( + agent_executor=MagicMock(), task_store=MagicMock() + ) + _contended_acquire(registry._lock, two_loops) + + +# --- end-to-end: on_message_send driven from two loops through v2 handler --- + + +def test_on_message_send_across_event_loops(two_loops) -> None: + """Drive the real request handler's ``on_message_send`` for two task ids + concurrently on two event loops sharing one handler + one task store.""" + import uuid + + from a2a.auth.user import UnauthenticatedUser + from a2a.helpers.proto_helpers import new_task_from_user_message + from a2a.server.agent_execution import AgentExecutor + from a2a.server.context import ServerCallContext + from a2a.server.request_handlers.default_request_handler_v2 import ( + DefaultRequestHandlerV2, + ) + from a2a.server.tasks import InMemoryTaskStore, TaskUpdater + from a2a.types import ( + AgentCapabilities, + AgentCard, + Message, + Part, + Role, + SendMessageRequest, + TaskState, + ) + + class _CompletingAgent(AgentExecutor): + async def execute(self, context, event_queue) -> None: + if context.message: + await event_queue.enqueue_event( + new_task_from_user_message(context.message) + ) + updater = TaskUpdater( + event_queue, + task_id=context.task_id or str(uuid.uuid4()), + context_id=context.context_id or str(uuid.uuid4()), + ) + await updater.update_status(TaskState.TASK_STATE_WORKING) + await updater.complete() + + async def cancel(self, context, event_queue) -> None: + raise NotImplementedError + + loop_a, loop_b = two_loops + handler = DefaultRequestHandlerV2( + _CompletingAgent(), + InMemoryTaskStore(), + AgentCard( + name='test_agent', + version='1.0', + capabilities=AgentCapabilities(streaming=True), + ), + ) + ctx = ServerCallContext(user=UnauthenticatedUser()) + + def req(mid: str) -> SendMessageRequest: + return SendMessageRequest( + message=Message( + role=Role.ROLE_USER, message_id=mid, parts=[Part(text='hi')] + ) + ) + + # Two requests genuinely in-flight on two loops against the same handler. + fut_a = asyncio.run_coroutine_threadsafe( + handler.on_message_send(req('a1'), ctx), loop_a._loop + ) + fut_b = asyncio.run_coroutine_threadsafe( + handler.on_message_send(req('b1'), ctx), loop_b._loop + ) + task_a = fut_a.result(timeout=15) + task_b = fut_b.result(timeout=15) + assert task_a.status.state == TaskState.TASK_STATE_COMPLETED + assert task_b.status.state == TaskState.TASK_STATE_COMPLETED