Skip to content

Commit 17473dd

Browse files
feat(client): expose close() and async context manager support on abstract Client
1 parent d2d3860 commit 17473dd

1 file changed

Lines changed: 20 additions & 0 deletions

File tree

src/a2a/client/client.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33

44
from abc import ABC, abstractmethod
55
from collections.abc import AsyncIterator, Callable, Coroutine
6+
from types import TracebackType
67
from typing import Any
78

9+
from typing_extensions import Self
10+
811
import httpx
912

1013
from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
@@ -107,6 +110,19 @@ def __init__(
107110
self._consumers = consumers
108111
self._middleware = middleware
109112

113+
async def __aenter__(self) -> Self:
114+
"""Enters the async context manager."""
115+
return self
116+
117+
async def __aexit__(
118+
self,
119+
exc_type: type[BaseException] | None,
120+
exc_val: BaseException | None,
121+
exc_tb: TracebackType | None,
122+
) -> None:
123+
"""Exits the async context manager and closes the client."""
124+
await self.close()
125+
110126
@abstractmethod
111127
async def send_message(
112128
self,
@@ -209,3 +225,7 @@ async def consume(
209225
return
210226
for c in self._consumers:
211227
await c(event, card)
228+
229+
@abstractmethod
230+
async def close(self) -> None:
231+
"""Closes the client and releases any underlying resources."""

0 commit comments

Comments
 (0)