feat(server): add async context manager support to EventQueue

Author: cchinchilla-dev

src/a2a/server/events/event_queue.py

@@ -2,6 +2,10 @@
 import logging
 import sys
 
+from types import TracebackType
+
+from typing_extensions import Self
+
 from a2a.types import (
     Message,
     Task,
@@ -43,6 +47,19 @@ def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
         self._lock = asyncio.Lock()
         logger.debug('EventQueue initialized.')
 
+    async def __aenter__(self) -> Self:
+        """Enters the async context manager, returning the queue itself."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exits the async context manager, ensuring close() is called."""
+        await self.close()
+
     async def enqueue_event(self, event: Event) -> None:
         """Enqueues an event to this queue and all its children.
 

tests/server/events/test_event_queue.py

@@ -69,6 +69,27 @@ def test_constructor_invalid_max_queue_size() -> None:
     ):
         EventQueue(max_queue_size=-10)
 
+@pytest.mark.asyncio
+async def test_event_queue_async_context_manager(
+    event_queue: EventQueue,
+) -> None:
+    """Test that EventQueue can be used as an async context manager."""
+    async with event_queue as q:
+        assert q is event_queue
+        assert event_queue.is_closed() is False
+    assert event_queue.is_closed() is True
+
+
+@pytest.mark.asyncio
+async def test_event_queue_async_context_manager_on_exception(
+    event_queue: EventQueue,
+) -> None:
+    """Test that close() is called even when an exception occurs inside the context."""
+    with pytest.raises(RuntimeError, match='boom'):
+        async with event_queue:
+            raise RuntimeError('boom')
+    assert event_queue.is_closed() is True
+
 
 @pytest.mark.asyncio
 async def test_enqueue_and_dequeue_event(event_queue: EventQueue) -> None: