Merge branch 'main' into ishymko/metadata-refactor

Author: ishymko

src/a2a/client/transports/base.py

@@ -1,5 +1,8 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Callable
+from types import TracebackType
+
+from typing_extensions import Self
 
 from a2a.client.middleware import ClientCallContext
 from a2a.types import (
@@ -19,6 +22,19 @@
 class ClientTransport(ABC):
     """Abstract base class for a client transport."""
 
+    async def __aenter__(self) -> Self:
+        """Enters 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: TracebackType | None,
+    ) -> None:
+        """Exits the async context manager, ensuring close() is called."""
+        await self.close()
+
     @abstractmethod
     async def send_message(
         self,

tests/client/test_base_client.py

@@ -1,4 +1,4 @@
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -61,6 +61,32 @@ def base_client(
     )
 
 
+@pytest.mark.asyncio
+async def test_transport_async_context_manager() -> None:
+    with (
+        patch.object(ClientTransport, '__abstractmethods__', set()),
+        patch.object(ClientTransport, 'close', new_callable=AsyncMock),
+    ):
+        transport = ClientTransport()
+        async with transport as t:
+            assert t is transport
+            transport.close.assert_not_awaited()
+        transport.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_transport_async_context_manager_on_exception() -> None:
+    with (
+        patch.object(ClientTransport, '__abstractmethods__', set()),
+        patch.object(ClientTransport, 'close', new_callable=AsyncMock),
+    ):
+        transport = ClientTransport()
+        with pytest.raises(RuntimeError, match='boom'):
+            async with transport:
+                raise RuntimeError('boom')
+        transport.close.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_send_message_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message