feat(client): expose close() and async context manager support on abstract Client (#719)

cchinchilla-dev · ishymko · web-flow · commit e25ba7be57fe · 2026-03-02T12:37:40.000+01:00
# Description Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [x] Follow the [`CONTRIBUTING` Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md). - [x] Make your Pull Request title in the <https://www.conventionalcommits.org/> specification. - Important Prefixes for [release-please](https://github.com/googleapis/release-please): - `fix:` which represents bug fixes, and correlates to a [SemVer](https://semver.org/) patch. - `feat:` represents a new feature, and correlates to a SemVer minor. - `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking change (indicated by the `!`) and will result in a SemVer major. - [x] Ensure the tests and linter pass (Run `bash scripts/format.sh` from the repository root to format) - [x] Appropriate docs were updated (if necessary) Fixes #689 🦕 ### Problem The abstract `Client` class (`src/a2a/client/client.py`) — the public interface returned by `ClientFactory.connect()` — does not declare `close()`, `__aenter__`, or `__aexit__`. Code typed against `Client` cannot use the async context manager pattern or call `close()`: ```python client = await ClientFactory.connect(card) # client is typed as Client, not BaseClient async with client: # TypeError ... await client.close() # AttributeError ``` `ClientTransport` already supports this protocol since #682, and `BaseClient` since #688, but the gap at the abstract interface level means the pattern is unavailable to consumers that depend on the `Client` type. The integration tests reflect this — they rely on `hasattr` checks instead of the type system: ```python if hasattr(transport, 'close'): await transport.close() ``` ### Fix - Add `close()` as an abstract method on `Client`, consistent with the existing 7 abstract methods. - Add `__aenter__` and `__aexit__` as concrete methods on `Client`, delegating to `close()`. `__aenter__` returns `Self` (via `typing_extensions`), matching the convention established in `ClientTransport` (#682). `BaseClient` already implements `close()` (delegating to `self._transport.close()`), so it satisfies the new abstract method with no additional changes. This enables the idiomatic pattern at the `Client` abstraction level: ```python async with await ClientFactory.connect(card) as client: async for event in client.send_message(msg): ... # close() called automatically, even on exceptions ``` ### Tests No new tests needed. The existing `test_base_client_async_context_manager` and `test_base_client_async_context_manager_on_exception` in `test_base_client.py` already exercise the `__aenter__`/`__aexit__`/`close()` chain through `BaseClient`, which inherits these methods from `Client`. Happy to add dedicated tests if maintainers prefer a different approach. --------- Co-authored-by: Ivan Shymko <ishymko@google.com>
diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py
@@ -1,9 +1,6 @@
 from collections.abc import AsyncGenerator, AsyncIterator, Callable
-from types import TracebackType
 from typing import Any
 
-from typing_extensions import Self
-
 from a2a.client.client import (
     Client,
     ClientCallContext,
@@ -51,19 +48,6 @@ 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,
diff --git a/src/a2a/client/client.py b/src/a2a/client/client.py
@@ -3,10 +3,13 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Callable, Coroutine
+from types import TracebackType
 from typing import Any
 
 import httpx
 
+from typing_extensions import Self
+
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.optionals import Channel
 from a2a.types.a2a_pb2 import (
@@ -110,6 +113,19 @@ def __init__(
         self._consumers = consumers
         self._middleware = middleware
 
+    async def __aenter__(self) -> Self:
+        """Enters the async context manager."""
+        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 and closes the client."""
+        await self.close()
+
     @abstractmethod
     async def send_message(
         self,
@@ -242,3 +258,7 @@ async def consume(
             return
         for c in self._consumers:
             await c(event, card)
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Closes the client and releases any underlying resources."""
