fix: use threading.RLock for in-memory server singletons

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()

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,15 +24,15 @@ 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.
 
         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,43 +43,43 @@ 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.
 
         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.
 
         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.
 
         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()

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

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(

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,

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(