Repository: a2aproject/a2a-python

Issue:
[Feat]: Add async context manager support to ClientTransport
### Is your feature request related to a problem? Please describe.

`ClientTransport` defines an abstract close() method, but doesn't implement` __aenter__/__aexit__`. This means transports can't be used with `async with`, which is the idiomatic Python pattern for managing async resources and preventing resource leaks.

Currently, if an exception occurs between creating a transport and calling `close()`, the underlying connection (e.g., gRPC channel) is never cleaned up:

```python
transport = GrpcTransport(channel=channel, agent_card=agent_card)
result = await transport.send_message(params)  # if this raises, close() is never called
await transport.close()
```

### Describe the solution you'd like

Add `__aenter__` and `__aexit__` methods to ClientTransport in `src/a2a/client/transports/base.py`:

```python
async def __aenter__(self) -> 'ClientTransport':
    return self

async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
    await self.close()
```

This enables the standard async context manager pattern:

```python
async with GrpcTransport(channel=channel, agent_card=agent_card) as transport:
    result = await transport.send_message(params)
    # close() called automatically, even on exceptions
```

This is a non-breaking, additive change. Calling `close()` manually continues to work as before.

### Describe alternatives you've considered

Relying on `try/finally` blocks at each call site, which is what the integration tests currently do. However, this is error-prone and verbose compared to the context manager pattern that is standard across the Python async ecosystem (`httpx.AsyncClient`, `aiohttp.ClientSession`, `grpc.aio.Channel` all support it).

### Additional context

The same pattern could also be applied to `BaseClient`, which wraps `ClientTransport` and also exposes a `close()` method. Happy to include that in the same PR if maintainers prefer.

### Code of Conduct

- [x] I agree to follow this project's Code of Conduct
Type: Class
Name: ClientTransport
Path: src/a2a/client/transports/base.py
Public API: async def __aenter__(self) -> Self; async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None) -> None
Description: Abstract base class for a client transport that now supports the async context manager protocol. Use `async with` to ensure `close()` is automatically called on exit, even if an exception occurs.

Discussion hints:

Solution patch:
diff --git a/src/a2a/client/transports/base.py b/src/a2a/client/transports/base.py
index 0c54a28dc..5bbf89adb 100644
--- a/src/a2a/client/transports/base.py
+++ b/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,


Test patch:
diff --git a/tests/client/test_base_client.py b/tests/client/test_base_client.py
index 7aa47902d..04bebb3bc 100644
--- a/tests/client/test_base_client.py
+++ b/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