Extract common code for HTTP from JSON-RPC and REST

ishymko · ishymko · commit f114f3d5a5ec · 2026-03-04T12:11:48.000Z
diff --git a/src/a2a/client/__init__.py b/src/a2a/client/__init__.py
@@ -13,6 +13,7 @@
 from a2a.client.client_factory import ClientFactory, minimal_agent_card
 from a2a.client.errors import (
     A2AClientError,
+    A2AClientTimeoutError,
     AgentCardResolutionError,
 )
 from a2a.client.helpers import create_text_message_object
@@ -25,6 +26,7 @@
 __all__ = [
     'A2ACardResolver',
     'A2AClientError',
+    'A2AClientTimeoutError',
     'AgentCardResolutionError',
     'AuthInterceptor',
     'BaseClient',
diff --git a/src/a2a/client/errors.py b/src/a2a/client/errors.py
@@ -13,3 +13,7 @@ class AgentCardResolutionError(A2AClientError):
     def __init__(self, message: str, status_code: int | None = None) -> None:
         super().__init__(message)
         self.status_code = status_code
+
+
+class A2AClientTimeoutError(A2AClientError):
+    """Exception for timeout errors during a request."""
diff --git a/src/a2a/client/transports/grpc.py b/src/a2a/client/transports/grpc.py
@@ -4,7 +4,7 @@
 from functools import wraps
 from typing import Any, NoReturn
 
-from a2a.client.errors import A2AClientError
+from a2a.client.errors import A2AClientError, A2AClientTimeoutError
 from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 
 
@@ -53,6 +53,9 @@
 
 
 def _map_grpc_error(e: grpc.aio.AioRpcError) -> NoReturn:
+    if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+        raise A2AClientTimeoutError('Client Request timed out') from e
+
     details = e.details()
     if isinstance(details, str) and ': ' in details:
         error_type_name, error_message = details.split(': ', 1)
diff --git a/src/a2a/client/transports/http_helpers.py b/src/a2a/client/transports/http_helpers.py
@@ -0,0 +1,74 @@
+import json
+
+from collections.abc import AsyncGenerator, Callable, Iterator
+from contextlib import contextmanager
+from typing import Any, NoReturn
+
+import httpx
+
+from httpx_sse import SSEError, aconnect_sse
+
+from a2a.client.errors import A2AClientError, A2AClientTimeoutError
+
+
+@contextmanager
+def handle_http_exceptions(
+    status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
+    | None = None,
+) -> Iterator[None]:
+    """Handles common HTTP exceptions for REST and JSON-RPC transports.
+
+    Args:
+        status_error_handler: Optional handler for `httpx.HTTPStatusError`.
+            If provided, this handler should raise an appropriate domain-specific exception.
+            If not provided, a default `A2AClientError` will be raised.
+    """
+    try:
+        yield
+    except httpx.TimeoutException as e:
+        raise A2AClientTimeoutError('Client Request timed out') from e
+    except httpx.HTTPStatusError as e:
+        if status_error_handler:
+            status_error_handler(e)
+        raise A2AClientError(f'HTTP Error {e.response.status_code}: {e}') from e
+    except SSEError as e:
+        raise A2AClientError(
+            f'Invalid SSE response or protocol error: {e}'
+        ) from e
+    except httpx.RequestError as e:
+        raise A2AClientError(f'Network communication error: {e}') from e
+    except json.JSONDecodeError as e:
+        raise A2AClientError(f'JSON Decode Error: {e}') from e
+
+
+async def send_http_request(
+    httpx_client: httpx.AsyncClient,
+    request: httpx.Request,
+    status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
+    | None = None,
+) -> dict[str, Any]:
+    """Sends an HTTP request and parses the JSON response, handling common exceptions."""
+    with handle_http_exceptions(status_error_handler):
+        response = await httpx_client.send(request)
+        response.raise_for_status()
+        return response.json()
+
+
+async def send_http_stream_request(
+    httpx_client: httpx.AsyncClient,
+    method: str,
+    url: str,
+    status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
+    | None = None,
+    **kwargs: Any,
+) -> AsyncGenerator[str]:
+    """Sends a streaming HTTP request, yielding SSE data strings and handling exceptions."""
+    with handle_http_exceptions(status_error_handler):
+        async with aconnect_sse(
+            httpx_client, method, url, **kwargs
+        ) as event_source:
+            event_source.response.raise_for_status()
+            async for sse in event_source.aiter_sse():
+                if not sse.data:
+                    continue
+                yield sse.data
diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py
@@ -1,4 +1,3 @@
-import json
 import logging
 
 from collections.abc import AsyncGenerator, Callable
@@ -8,12 +7,15 @@
 import httpx
 
 from google.protobuf import json_format
-from httpx_sse import SSEError, aconnect_sse
 from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
 
 from a2a.client.errors import A2AClientError
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
+from a2a.client.transports.http_helpers import (
+    send_http_request,
+    send_http_stream_request,
+)
 from a2a.extensions.common import update_extension_header
 from a2a.types.a2a_pb2 import (
     AgentCard,
@@ -470,22 +472,10 @@ async def _send_request(
         rpc_request_payload: dict[str, Any],
         http_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
-        try:
-            response = await self.httpx_client.post(
-                self.url, json=rpc_request_payload, **(http_kwargs or {})
-            )
-            response.raise_for_status()
-            return response.json()
-        except httpx.TimeoutException as e:
-            raise A2AClientError('Client Request timed out') from e
-        except httpx.HTTPStatusError as e:
-            raise A2AClientError(
-                f'HTTP Error {e.response.status_code}: {e}'
-            ) from e
-        except json.JSONDecodeError as e:
-            raise A2AClientError(str(e)) from e
-        except httpx.RequestError as e:
-            raise A2AClientError(f'Network communication error: {e}') from e
+        request = self.httpx_client.build_request(
+            'POST', self.url, json=rpc_request_payload, **(http_kwargs or {})
+        )
+        return await send_http_request(self.httpx_client, request)
 
     async def _send_stream_request(
         self,
@@ -500,39 +490,18 @@ async def _send_stream_request(
         headers.update(final_kwargs.get('headers', {}))
         final_kwargs['headers'] = headers
 
-        try:
-            async with aconnect_sse(
-                self.httpx_client,
-                'POST',
-                self.url,
-                json=rpc_request_payload,
-                **final_kwargs,
-            ) as event_source:
-                try:
-                    event_source.response.raise_for_status()
-                    async for sse in event_source.aiter_sse():
-                        if not sse.data:
-                            continue
-                        json_rpc_response = JSONRPC20Response.from_json(
-                            sse.data
-                        )
-                        if json_rpc_response.error:
-                            self._handle_jsonrpc_error(json_rpc_response.error)
-                        response: StreamResponse = json_format.ParseDict(
-                            json_rpc_response.result, StreamResponse()
-                        )
-                        yield response
-                except httpx.HTTPStatusError as e:
-                    raise A2AClientError(
-                        f'HTTP Error {e.response.status_code}: {e}'
-                    ) from e
-                except SSEError as e:
-                    raise A2AClientError(
-                        f'Invalid SSE response or protocol error: {e}'
-                    ) from e
-        except httpx.TimeoutException as e:
-            raise A2AClientError('Client Request timed out') from e
-        except httpx.RequestError as e:
-            raise A2AClientError(f'Network communication error: {e}') from e
-        except json.JSONDecodeError as e:
-            raise A2AClientError(str(e)) from e
+        async for sse_data in send_http_stream_request(
+            self.httpx_client,
+            'POST',
+            self.url,
+            None,
+            json=rpc_request_payload,
+            **final_kwargs,
+        ):
+            json_rpc_response = JSONRPC20Response.from_json(sse_data)
+            if json_rpc_response.error:
+                self._handle_jsonrpc_error(json_rpc_response.error)
+            response: StreamResponse = json_format.ParseDict(
+                json_rpc_response.result, StreamResponse()
+            )
+            yield response
diff --git a/src/a2a/client/transports/rest.py b/src/a2a/client/transports/rest.py
@@ -8,11 +8,14 @@
 
 from google.protobuf.json_format import MessageToDict, Parse, ParseDict
 from google.protobuf.message import Message
-from httpx_sse import SSEError, aconnect_sse
 
 from a2a.client.errors import A2AClientError
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
+from a2a.client.transports.http_helpers import (
+    send_http_request,
+    send_http_stream_request,
+)
 from a2a.extensions.common import update_extension_header
 from a2a.types.a2a_pb2 import (
     AgentCard,
@@ -427,48 +430,20 @@ async def _send_stream_request(
         final_kwargs.update(kwargs)
         final_kwargs.setdefault('timeout', None)
 
-        try:
-            async with aconnect_sse(
-                self.httpx_client,
-                method,
-                f'{self.url}{target}',
-                **final_kwargs,
-            ) as event_source:
-                try:
-                    event_source.response.raise_for_status()
-                    async for sse in event_source.aiter_sse():
-                        if not sse.data:
-                            continue
-                        event: StreamResponse = Parse(
-                            sse.data, StreamResponse()
-                        )
-                        yield event
-                except httpx.HTTPStatusError as e:
-                    self._handle_http_error(e)
-                except SSEError as e:
-                    raise A2AClientError(
-                        f'Invalid SSE response or protocol error: {e}'
-                    ) from e
-        except httpx.TimeoutException as e:
-            raise A2AClientError('Client Request timed out') from e
-        except httpx.RequestError as e:
-            raise A2AClientError(f'Network communication error: {e}') from e
-        except json.JSONDecodeError as e:
-            raise A2AClientError(f'JSON Decode Error: {e}') from e
+        async for sse_data in send_http_stream_request(
+            self.httpx_client,
+            method,
+            f'{self.url}{target}',
+            self._handle_http_error,
+            **final_kwargs,
+        ):
+            event: StreamResponse = Parse(sse_data, StreamResponse())
+            yield event
 
     async def _send_request(self, request: httpx.Request) -> dict[str, Any]:
-        try:
-            response = await self.httpx_client.send(request)
-            response.raise_for_status()
-            return response.json()
-        except httpx.TimeoutException as e:
-            raise A2AClientError('Client Request timed out') from e
-        except httpx.HTTPStatusError as e:
-            self._handle_http_error(e)
-        except json.JSONDecodeError as e:
-            raise A2AClientError(f'JSON Decode Error: {e}') from e
-        except httpx.RequestError as e:
-            raise A2AClientError(f'Network communication error: {e}') from e
+        return await send_http_request(
+            self.httpx_client, request, self._handle_http_error
+        )
 
     async def _send_post_request(
         self,
diff --git a/tests/client/transports/test_jsonrpc_client.py b/tests/client/transports/test_jsonrpc_client.py
diff --git a/tests/client/transports/test_rest_client.py b/tests/client/transports/test_rest_client.py
