feat(server): implement RedisQueueManager for distributed event handling

- Add RedisQueueManager class to manage event queues across distributed services
- Implement local and proxy queue management using Redis pub/sub
- Add logging for better visibility and debugging
- Update tests to cover new functionality

Author: long.qul

src/a2a/server/events/redis_queue_manager.py

@@ -1,8 +1,11 @@
 import asyncio
+import logging
 
 from asyncio import Task
 from functools import partial
+from typing import Any, Dict, Optional
 
+from pydantic import ValidationError, TypeAdapter
 from redis.asyncio import Redis
 
 from a2a.server.events import (
@@ -14,6 +17,8 @@
     TaskQueueExists,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class RedisQueueManager(QueueManager):
     """This implements the `QueueManager` interface using Redis for event.
@@ -40,14 +45,7 @@ def __init__(
         self._relay_channel_name = relay_channel_key_prefix
         self._background_tasks: dict[str, Task] = {}
         self._task_registry_name = task_registry_key
-
-    async def _listen_and_relay(self, task_id: str) -> None:
-        c = EventConsumer(self._local_queue[task_id])
-        async for event in c.consume_all():
-            await self._redis.publish(
-                self._task_channel_name(task_id),
-                event.model_dump_json(exclude_none=True),
-            )
+        self._pubsub_listener_task: Optional[Task] = None
 
     def _task_channel_name(self, task_id: str) -> str:
         return self._relay_channel_name + task_id
@@ -57,10 +55,23 @@ async def _has_task_id(self, task_id: str) -> bool:
         return ret == 1
 
     async def _register_task_id(self, task_id: str) -> None:
-        await self._redis.sadd(self._task_registry_name, task_id)
+        task_started_event = asyncio.Event()
+        async def _wrapped_listen_and_relay() -> None:
+            task_started_event.set()
+            c = EventConsumer(self._local_queue[task_id].tap())
+            async for event in c.consume_all():
+                logger.debug(f'Publishing event for task {task_id} in QM {self}: {event}')
+                await self._redis.publish(
+                    self._task_channel_name(task_id),
+                    event.model_dump_json(exclude_none=True),
+                )
+
         self._background_tasks[task_id] = asyncio.create_task(
-            self._listen_and_relay(task_id)
+            _wrapped_listen_and_relay()
         )
+        await task_started_event.wait()
+        await self._redis.sadd(self._task_registry_name, task_id)
+        logger.debug(f'Started to listen and relay events for task {task_id}')
 
     async def _remove_task_id(self, task_id: str) -> bool:
         if task_id in self._background_tasks:
@@ -70,21 +81,51 @@ async def _remove_task_id(self, task_id: str) -> bool:
         return await self._redis.srem(self._task_registry_name, task_id) == 1
 
     async def _subscribe_remote_task_events(self, task_id: str) -> None:
-        await self._pubsub.subscribe(
-            **{
-                self._task_channel_name(task_id): partial(
-                    self._relay_remote_events, task_id
-                )
-            }
-        )
-
-    def _unsubscribe_remote_task_events(self, task_id: str) -> None:
-        self._pubsub.unsubscribe(self._task_channel_name(task_id))
-
-    def _relay_remote_events(self, task_id: str, event_json: str) -> None:
-        if task_id in self._proxy_queue:
-            event = Event.model_validate_json(event_json)
-            self._proxy_queue[task_id].enqueue_event(event)
+        channel_id = self._task_channel_name(task_id)
+        await self._pubsub.subscribe(**{channel_id: self._relay_remote_events})
+
+        # this is a global listener to handle incoming pubsub events
+        if not self._pubsub_listener_task:
+            logger.debug('Creating pubsub listener task.')
+            self._pubsub_listener_task = asyncio.create_task(self._consume_pubsub_messages())
+
+        logger.debug(f"Subscribed for remote events for task {task_id}")
+
+    async def _consume_pubsub_messages(self):
+        async for _ in self._pubsub.listen():
+            pass
+
+    async def _relay_remote_events(self, subscription_event) -> None:
+        if 'channel' not in subscription_event or 'data' not in subscription_event:
+            logger.warning(f"channel or data is absent in subscription event: {subscription_event}")
+            return
+
+        channel_id: str = subscription_event['channel'].decode('utf-8')
+        data_string: str = subscription_event['data'].decode('utf-8')
+        task_id = channel_id.split('.')[-1]
+        if task_id not in self._proxy_queue:
+            logger.warning(f"task_id {task_id} not found in proxy queue")
+            return
+
+        try:
+            logger.debug(f"Received event for task_id {task_id} in QM {self}: {data_string}")
+            event = TypeAdapter(Event).validate_json(data_string)
+        except Exception as e:
+            logger.warning(f"Failed to parse event from subscription event: {subscription_event}: {e}")
+            return
+
+        logger.debug(f"Enqueuing event for task_id {task_id} in QM {self}: {event}")
+        await self._proxy_queue[task_id].enqueue_event(event)
+
+
+    async def _unsubscribe_remote_task_events(self, task_id: str) -> None:
+        # unsubscribe channel for given task_id
+        await self._pubsub.unsubscribe(self._task_channel_name(task_id))
+        # release global listener if not channel is subscribed
+        async with self._lock:
+            if not self._pubsub.subscribed and self._pubsub_listener_task:
+                self._pubsub_listener_task.cancel()
+                self._pubsub_listener_task = None
 
     async def add(self, task_id: str, queue: EventQueue) -> None:
         """Add a new local event queue for the specified task.
@@ -96,11 +137,13 @@ async def add(self, task_id: str, queue: EventQueue) -> None:
         Raises:
             TaskQueueExists: If a queue for the task already exists.
         """
+        logger.debug(f"add {task_id}")
         async with self._lock:
             if await self._has_task_id(task_id):
                 raise TaskQueueExists()
             self._local_queue[task_id] = queue
             await self._register_task_id(task_id)
+            logger.debug(f"Local queue is created for task {task_id}")
 
     async def get(self, task_id: str) -> EventQueue | None:
         """Get the event queue associated with the given task ID.
@@ -115,17 +158,22 @@ async def get(self, task_id: str) -> EventQueue | None:
         Returns:
             EventQueue | None: The event queue if found, otherwise None.
         """
+        logger.debug(f"get {task_id}")
         async with self._lock:
             # lookup locally
             if task_id in self._local_queue:
+                logger.debug(f"Got local queue for task_id {task_id}")
                 return self._local_queue[task_id]
             # lookup globally
             if await self._has_task_id(task_id):
                 if task_id not in self._proxy_queue:
+                    logger.debug(f"Creating proxy queue for {task_id}")
                     queue = EventQueue()
                     self._proxy_queue[task_id] = queue
                     await self._subscribe_remote_task_events(task_id)
+                logger.debug(f"Got proxy queue for task_id {task_id}")
                 return self._proxy_queue[task_id]
+            logger.warning(f"Attempted to get non-existing queue for task {task_id}")
             return None
 
     async def tap(self, task_id: str) -> EventQueue | None:
@@ -137,8 +185,10 @@ async def tap(self, task_id: str) -> EventQueue | None:
         Returns:
             EventQueue | None: A new reference to the event queue if it exists, otherwise None.
         """
+        logger.debug(f"tap {task_id}")
         event_queue = await self.get(task_id)
         if event_queue:
+            logger.debug(f'Tapping event queue for task: {task_id}')
             return event_queue.tap()
         return None
 
@@ -155,23 +205,27 @@ async def close(self, task_id: str) -> None:
         Raises:
             NoTaskQueue: If no queue exists for the given task ID.
         """
+        logger.debug(f"close {task_id}")
         async with self._lock:
             if task_id in self._local_queue:
                 # close locally
                 queue = self._local_queue.pop(task_id)
                 await queue.close()
                 # remove from global registry if a local queue is closed
                 await self._remove_task_id(task_id)
+                logger.debug(f"Closing local queue for task {task_id}")
                 return
 
             if task_id in self._proxy_queue:
                 # close proxy queue
                 queue = self._proxy_queue.pop(task_id)
                 await queue.close()
                 # unsubscribe from remote, but don't remove from global registry
-                self._unsubscribe_remote_task_events(task_id)
+                await self._unsubscribe_remote_task_events(task_id)
+                logger.debug(f"Closing proxy queue for task {task_id}")
                 return
 
+            logger.warning(f"Attempted to close non-existing queue found for task {task_id}")
             raise NoTaskQueue()
 
     async def create_or_tap(self, task_id: str) -> EventQueue:
@@ -186,23 +240,28 @@ async def create_or_tap(self, task_id: str) -> EventQueue:
         Returns:
             EventQueue: An event queue associated with the given task ID.
         """
+        logger.debug(f"create_or_tap {task_id}")
         async with self._lock:
             if await self._has_task_id(task_id):
                 # if it's a local queue, tap directly
                 if task_id in self._local_queue:
+                    logger.debug(f"Tapping a local queue for task {task_id}")
                     return self._local_queue[task_id].tap()
 
                 # if it's a proxy queue, tap the proxy
                 if task_id in self._proxy_queue:
+                    logger.debug(f"Tapping a proxy queue for task {task_id}")
                     return self._proxy_queue[task_id].tap()
 
                 # if the proxy is not created, create the proxy and return
                 queue = EventQueue()
                 self._proxy_queue[task_id] = queue
                 await self._subscribe_remote_task_events(task_id)
+                logger.debug(f"Creating a proxy queue for task {task_id}")
                 return self._proxy_queue[task_id]
             # the task doesn't exist before, create a local queue
             queue = EventQueue()
             self._local_queue[task_id] = queue
             await self._register_task_id(task_id)
+            logger.debug(f"Creating a local queue for task {task_id}")
             return queue

tests/server/test_integration.py

@@ -1,9 +1,13 @@
 import asyncio
+import logging
+import os
 
 from typing import Any
 from unittest import mock
 
 import pytest
+import pytest_asyncio
+from redis.asyncio import Redis
 
 from starlette.authentication import (
     AuthCredentials,
@@ -22,6 +26,8 @@
     A2AFastAPIApplication,
     A2AStarletteApplication,
 )
+from a2a.server.events import EventQueue
+from a2a.server.events.redis_queue_manager import RedisQueueManager
 from a2a.types import (
     AgentCapabilities,
     AgentCard,
@@ -44,6 +50,7 @@
     TextPart,
     UnsupportedOperationError,
 )
+from a2a.utils import new_agent_text_message
 from a2a.utils.errors import MethodNotImplementedError
 
 
@@ -884,3 +891,88 @@ def test_non_dict_json(client: TestClient):
     data = response.json()
     assert 'error' in data
     assert data['error']['code'] == InvalidRequestError().code
+
+
+# === RedisQueueManager ===
+@pytest.mark.asyncio
+@pytest_asyncio.fixture(scope="function")
+async def asyncio_redis():
+    redis_server_url = os.getenv("REDIS_SERVER_URL")
+    if redis_server_url:
+        redis = Redis.from_url(redis_server_url)
+    else:
+        # use fake redis instead if no redis server url is given
+        from fakeredis import FakeAsyncRedis
+        redis = FakeAsyncRedis()
+    logging.info("flush redis for next test case")
+    await redis.flushall(asynchronous=False)
+    yield redis
+    await redis.close()
+
+
+@pytest.mark.asyncio
+async def test_redis_queue_local_only_queue(asyncio_redis):
+    queue_manager = RedisQueueManager(asyncio_redis)
+
+    # setup local queues
+    q1 = EventQueue()
+    await queue_manager.add('task_1', q1)
+    q2 = EventQueue()
+    await queue_manager.add('task_2', q2)
+    q3 = await queue_manager.tap("task_1")
+    assert await queue_manager.get('task_1') == q1
+    assert await queue_manager.get('task_2') == q2
+
+    # send and receive locally
+    msg1 = new_agent_text_message('hello')
+    await q1.enqueue_event(msg1)
+    assert await q1.dequeue_event(no_wait=True) == msg1
+    assert await q3.dequeue_event(no_wait=True) == msg1
+    # raise error if queue is empty
+    with pytest.raises(asyncio.QueueEmpty):
+        await q1.dequeue_event(no_wait=True)
+    # q2 is empty
+    with pytest.raises(asyncio.QueueEmpty):
+        await q2.dequeue_event(no_wait=True)
+
+
+@pytest.mark.asyncio
+async def test_redis_queue_mixed_queue(asyncio_redis):
+    qm1 = RedisQueueManager(asyncio_redis)
+    qm2 = RedisQueueManager(asyncio_redis)
+    qm3 = RedisQueueManager(asyncio_redis)
+
+    # create local queue in qm1
+    q1 = EventQueue()
+    await qm1.add('task_1', q1)
+    assert 'task_1' in qm1._local_queue
+    assert await qm1.get('task_1') == q1
+
+    # create proxy queue in qm2 through `get` method
+    q1_1 = await qm2.get('task_1')
+    assert 'task_1' in qm2._proxy_queue and 'task_1' not in qm2._local_queue
+    assert q1_1 != q1
+
+    # create proxy queue in qm3 through `tap` method
+    q1_2 = await qm3.tap("task_1")
+    assert 'task_1' in qm3._proxy_queue and 'task_1' not in qm3._local_queue
+
+    # enqueue and dequeue in q1
+    msg1 = new_agent_text_message('hello')
+    await q1.enqueue_event(msg1)
+    assert await q1.dequeue_event() == msg1
+    with pytest.raises(asyncio.QueueEmpty):
+        await q1.dequeue_event(no_wait=True)
+
+    # dequeue in q1_1
+    msg1_1: Message = await q1_1.dequeue_event()
+    assert msg1_1.parts[0].root.text == msg1.parts[0].root.text
+    with pytest.raises(asyncio.QueueEmpty):
+        await q1_1.dequeue_event(no_wait=True)
+
+    # dequeue in q1_2
+    msg1_2: Message = await q1_2.dequeue_event()
+    assert msg1_2.parts[0].root.text == msg1.parts[0].root.text
+    with pytest.raises(asyncio.QueueEmpty):
+        await q1_2.dequeue_event(no_wait=True)
+