feat: add async context manager support to BaseClient

Author: cchinchilla-dev

src/a2a/client/base_client.py

@@ -1,6 +1,9 @@
 from collections.abc import AsyncIterator, Callable
+from types import TracebackType
 from typing import Any
 
+from typing_extensions import Self
+
 from a2a.client.client import (
     Client,
     ClientCallContext,
@@ -43,6 +46,19 @@ def __init__(
         self._config = config
         self._transport = transport
 
+    async def __aenter__(self) -> Self:
+        """Enters the async context manager, returning the client 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 send_message(
         self,
         request: Message,

tests/client/test_base_client.py

@@ -61,6 +61,26 @@ def base_client(
     )
 
 
+@pytest.mark.asyncio
+async def test_base_client_async_context_manager(
+    base_client: BaseClient, mock_transport: AsyncMock
+) -> None:
+    async with base_client as client:
+        assert client is base_client
+        mock_transport.close.assert_not_awaited()
+    mock_transport.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_base_client_async_context_manager_on_exception(
+    base_client: BaseClient, mock_transport: AsyncMock
+) -> None:
+    with pytest.raises(RuntimeError, match='boom'):
+        async with base_client:
+            raise RuntimeError('boom')
+    mock_transport.close.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_send_message_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message