feat: add async context manager support to ClientTransport

Author: cchinchilla-dev

src/a2a/client/transports/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Callable
 
@@ -19,6 +21,19 @@
 class ClientTransport(ABC):
     """Abstract base class for a client transport."""
 
+    async def __aenter__(self) -> ClientTransport:
+        """Enter the async context manager, returning the transport itself."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object | None,
+    ) -> None:
+        """Exit the async context manager, ensuring close() is called."""
+        await self.close()
+
     @abstractmethod
     async def send_message(
         self,

tests/client/test_base_client.py

@@ -61,6 +61,18 @@ def base_client(
     )
 
 
+@pytest.mark.asyncio
+async def test_transport_aenter_returns_self(mock_transport: AsyncMock) -> None:
+    result = await ClientTransport.__aenter__(mock_transport)
+    assert result is mock_transport
+
+
+@pytest.mark.asyncio
+async def test_transport_aexit_calls_close(mock_transport: AsyncMock) -> None:
+    await ClientTransport.__aexit__(mock_transport, None, None, None)
+    mock_transport.close.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_send_message_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message