Merge remote-tracking branch 'origin/main' into 1.0-dev

Author: ishymko

src/a2a/client/base_client.py

@@ -1,6 +1,9 @@
 from collections.abc import AsyncIterator, Callable
+from types import TracebackType
 from typing import Any
 
+from typing_extensions import Self
+
 from a2a.client.client import (
     Client,
     ClientCallContext,
@@ -45,6 +48,19 @@ 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,

src/a2a/client/transports/jsonrpc.py

@@ -188,6 +188,8 @@ async def send_message_streaming(
                     if isinstance(response.root, JSONRPCErrorResponse):
                         raise A2AClientJSONRPCError(response.root)
                     yield response.root.result
+            except httpx.TimeoutException as e:
+                raise A2AClientTimeoutError('Client Request timed out') from e
             except httpx.HTTPStatusError as e:
                 raise A2AClientHTTPError(e.response.status_code, str(e)) from e
             except SSEError as e:
@@ -212,7 +214,7 @@ async def _send_request(
             )
             response.raise_for_status()
             return response.json()
-        except httpx.ReadTimeout as e:
+        except httpx.TimeoutException as e:
             raise A2AClientTimeoutError('Client Request timed out') from e
         except httpx.HTTPStatusError as e:
             raise A2AClientHTTPError(e.response.status_code, str(e)) from e
@@ -389,6 +391,8 @@ async def resubscribe(
                     if isinstance(response.root, JSONRPCErrorResponse):
                         raise A2AClientJSONRPCError(response.root)
                     yield response.root.result
+            except httpx.TimeoutException as e:
+                raise A2AClientTimeoutError('Client Request timed out') from e
             except SSEError as e:
                 raise A2AClientHTTPError(
                     400, f'Invalid SSE response or protocol error: {e}'

src/a2a/client/transports/rest.py

@@ -11,7 +11,11 @@
 from pydantic import BaseModel
 
 from a2a.client.card_resolver import A2ACardResolver
-from a2a.client.errors import A2AClientHTTPError, A2AClientJSONError
+from a2a.client.errors import (
+    A2AClientHTTPError,
+    A2AClientJSONError,
+    A2AClientTimeoutError,
+)
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
 from a2a.extensions.common import update_extension_header
@@ -163,6 +167,8 @@ async def send_message_streaming(
                     event = a2a_pb2.StreamResponse()
                     Parse(sse.data, event)
                     yield proto_utils.FromProto.stream_response(event)
+            except httpx.TimeoutException as e:
+                raise A2AClientTimeoutError('Client Request timed out') from e
             except httpx.HTTPStatusError as e:
                 raise A2AClientHTTPError(e.response.status_code, str(e)) from e
             except SSEError as e:
@@ -181,6 +187,8 @@ async def _send_request(self, request: httpx.Request) -> dict[str, Any]:
             response = await self.httpx_client.send(request)
             response.raise_for_status()
             return response.json()
+        except httpx.TimeoutException as e:
+            raise A2AClientTimeoutError('Client Request timed out') from e
         except httpx.HTTPStatusError as e:
             raise A2AClientHTTPError(e.response.status_code, str(e)) from e
         except json.JSONDecodeError as e:
@@ -383,6 +391,8 @@ async def resubscribe(
                     event = a2a_pb2.StreamResponse()
                     Parse(sse.data, event)
                     yield proto_utils.FromProto.stream_response(event)
+            except httpx.TimeoutException as e:
+                raise A2AClientTimeoutError('Client Request timed out') from e
             except SSEError as e:
                 raise A2AClientHTTPError(
                     400, f'Invalid SSE response or protocol error: {e}'

tests/client/test_base_client.py

@@ -87,6 +87,26 @@ async def test_transport_async_context_manager_on_exception() -> None:
         transport.close.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_base_client_async_context_manager(
+    base_client: BaseClient, mock_transport: AsyncMock
+) -> None:
+    async with base_client as client:
+        assert client is base_client
+        mock_transport.close.assert_not_awaited()
+    mock_transport.close.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_base_client_async_context_manager_on_exception(
+    base_client: BaseClient, mock_transport: AsyncMock
+) -> None:
+    with pytest.raises(RuntimeError, match='boom'):
+        async with base_client:
+            raise RuntimeError('boom')
+    mock_transport.close.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_send_message_streaming(
     base_client: BaseClient, mock_transport: MagicMock, sample_message: Message

tests/client/transports/test_jsonrpc_client.py

@@ -601,6 +601,38 @@ async def test_send_message_client_timeout(
 
         assert 'Client Request timed out' in str(exc_info.value)
 
+    @pytest.mark.asyncio
+    @patch('a2a.client.transports.jsonrpc.aconnect_sse')
+    async def test_send_message_streaming_timeout(
+        self,
+        mock_aconnect_sse: AsyncMock,
+        mock_httpx_client: AsyncMock,
+        mock_agent_card: MagicMock,
+    ):
+        client = JsonRpcTransport(
+            httpx_client=mock_httpx_client, agent_card=mock_agent_card
+        )
+        params = MessageSendParams(
+            message=create_text_message_object(content='Hello stream')
+        )
+        mock_event_source = AsyncMock(spec=EventSource)
+        mock_event_source.response = MagicMock(spec=httpx.Response)
+        mock_event_source.response.raise_for_status.return_value = None
+        mock_event_source.aiter_sse.side_effect = httpx.TimeoutException(
+            'Read timed out'
+        )
+        mock_aconnect_sse.return_value.__aenter__.return_value = (
+            mock_event_source
+        )
+
+        with pytest.raises(A2AClientTimeoutError) as exc_info:
+            _ = [
+                item
+                async for item in client.send_message_streaming(request=params)
+            ]
+
+        assert 'Client Request timed out' in str(exc_info.value)
+
     @pytest.mark.asyncio
     async def test_get_task_success(
         self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock

tests/client/transports/test_rest_client.py

@@ -9,7 +9,7 @@
 from httpx_sse import EventSource, ServerSentEvent
 
 from a2a.client import create_text_message_object
-from a2a.client.errors import A2AClientHTTPError
+from a2a.client.errors import A2AClientHTTPError, A2AClientTimeoutError
 from a2a.client.transports.rest import RestTransport
 from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.grpc import a2a_pb2
@@ -50,6 +50,40 @@ def _assert_extensions_header(mock_kwargs: dict, expected_extensions: set[str]):
     assert actual_extensions == expected_extensions
 
 
+class TestRestTransport:
+    @pytest.mark.asyncio
+    @patch('a2a.client.transports.rest.aconnect_sse')
+    async def test_send_message_streaming_timeout(
+        self,
+        mock_aconnect_sse: AsyncMock,
+        mock_httpx_client: AsyncMock,
+        mock_agent_card: MagicMock,
+    ):
+        client = RestTransport(
+            httpx_client=mock_httpx_client, agent_card=mock_agent_card
+        )
+        params = MessageSendParams(
+            message=create_text_message_object(content='Hello stream')
+        )
+        mock_event_source = AsyncMock(spec=EventSource)
+        mock_event_source.response = MagicMock(spec=httpx.Response)
+        mock_event_source.response.raise_for_status.return_value = None
+        mock_event_source.aiter_sse.side_effect = httpx.TimeoutException(
+            'Read timed out'
+        )
+        mock_aconnect_sse.return_value.__aenter__.return_value = (
+            mock_event_source
+        )
+
+        with pytest.raises(A2AClientTimeoutError) as exc_info:
+            _ = [
+                item
+                async for item in client.send_message_streaming(request=params)
+            ]
+
+        assert 'Client Request timed out' in str(exc_info.value)
+
+
 class TestRestTransportExtensions:
     @pytest.mark.asyncio
     async def test_send_message_with_default_extensions(