Refactor rest transport to reduce duplication, fix pyright errors

Author: mikeas1

src/a2a/client/legacy_grpc.py

@@ -15,7 +15,7 @@
 class A2AGrpcClient(GrpcTransport):
     """[DEPRECATED] Backwards compatibility wrapper for the gRPC client."""
 
-    def __init__(
+    def __init__(  # pylint: disable=super-init-not-called
         self,
         grpc_stub: 'A2AServiceStub',
         agent_card: AgentCard,
@@ -26,4 +26,19 @@ def __init__(
             DeprecationWarning,
             stacklevel=2,
         )
-        super().__init__(grpc_stub, agent_card)
+        # The old gRPC client accepted a stub directly. The new one accepts a
+        # channel and builds the stub itself. We just have a stub here, so we
+        # need to handle initialization ourselves.
+        self.stub = grpc_stub
+        self.agent_card = agent_card
+        self._needs_extended_card = (
+            agent_card.supports_authenticated_extended_card
+            if agent_card
+            else True
+        )
+
+        class _NopChannel:
+            async def close(self):
+                pass
+
+        self.channel = _NopChannel()

src/a2a/client/optionals.py

@@ -3,7 +3,7 @@
 
 # Attempt to import the optional module
 try:
-    from grpc.aio import Channel
+    from grpc.aio import Channel  # pyright: ignore[reportAssignmentType]
 except ImportError:
     # If grpc.aio is not available, define a dummy type for type checking.
     # This dummy type will only be used by type checkers.

src/a2a/client/transports/grpc.py

@@ -5,8 +5,6 @@
 
 try:
     import grpc
-
-    from grpc.aio import Channel
 except ImportError as e:
     raise ImportError(
         'A2AGrpcClient requires grpcio and grpcio-tools to be installed. '
@@ -16,6 +14,7 @@
 
 from a2a.client.client import ClientConfig
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
+from a2a.client.optionals import Channel
 from a2a.client.transports.base import ClientTransport
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
@@ -112,7 +111,7 @@ async def send_message_streaming(
         )
         while True:
             response = await stream.read()
-            if response == grpc.aio.EOF:
+            if response == grpc.aio.EOF:  # pyright: ignore[reportAttributeAccessIssue]
                 break
             yield proto_utils.FromProto.stream_response(response)
 
@@ -127,7 +126,7 @@ async def resubscribe(
         )
         while True:
             response = await stream.read()
-            if response == grpc.aio.EOF:
+            if response == grpc.aio.EOF:  # pyright: ignore[reportAttributeAccessIssue]
                 break
             yield proto_utils.FromProto.stream_response(response)
 

src/a2a/client/transports/rest.py

@@ -78,13 +78,9 @@ def _get_http_args(
     ) -> dict[str, Any] | None:
         return context.state.get('http_kwargs') if context else None
 
-    async def send_message(
-        self,
-        request: MessageSendParams,
-        *,
-        context: ClientCallContext | None = None,
-    ) -> Task | Message:
-        """Sends a non-streaming message request to the agent."""
+    async def _prepare_send_message(
+        self, request: MessageSendParams, context: ClientCallContext | None
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
         pb = a2a_pb2.SendMessageRequest(
             request=proto_utils.ToProto.message(request.message),
             configuration=proto_utils.ToProto.message_send_configuration(
@@ -102,6 +98,18 @@ async def send_message(
             self._get_http_args(context),
             context,
         )
+        return payload, modified_kwargs
+
+    async def send_message(
+        self,
+        request: MessageSendParams,
+        *,
+        context: ClientCallContext | None = None,
+    ) -> Task | Message:
+        """Sends a non-streaming message request to the agent."""
+        payload, modified_kwargs = await self._prepare_send_message(
+            request, context
+        )
         response_data = await self._send_post_request(
             '/v1/message:send', payload, modified_kwargs
         )
@@ -118,22 +126,8 @@ async def send_message_streaming(
         Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent | Message
     ]:
         """Sends a streaming message request to the agent and yields responses as they arrive."""
-        pb = a2a_pb2.SendMessageRequest(
-            request=proto_utils.ToProto.message(request.message),
-            configuration=proto_utils.ToProto.message_send_configuration(
-                request.configuration
-            ),
-            metadata=(
-                proto_utils.ToProto.metadata(request.metadata)
-                if request.metadata
-                else None
-            ),
-        )
-        payload = MessageToDict(pb)
-        payload, modified_kwargs = await self._apply_interceptors(
-            payload,
-            self._get_http_args(context),
-            context,
+        payload, modified_kwargs = await self._prepare_send_message(
+            request, context
         )
 
         modified_kwargs.setdefault('timeout', None)
@@ -161,18 +155,9 @@ async def send_message_streaming(
                     503, f'Network communication error: {e}'
                 ) from e
 
-    async def _send_post_request(
-        self,
-        target: str,
-        rpc_request_payload: dict[str, Any],
-        http_kwargs: dict[str, Any] | None = None,
-    ) -> dict[str, Any]:
+    async def _send_request(self, request: httpx.Request) -> dict[str, Any]:
         try:
-            response = await self.httpx_client.post(
-                f'{self.url}{target}',
-                json=rpc_request_payload,
-                **(http_kwargs or {}),
-            )
+            response = await self.httpx_client.send(request)
             response.raise_for_status()
             return response.json()
         except httpx.HTTPStatusError as e:
@@ -184,28 +169,35 @@ async def _send_post_request(
                 503, f'Network communication error: {e}'
             ) from e
 
+    async def _send_post_request(
+        self,
+        target: str,
+        rpc_request_payload: dict[str, Any],
+        http_kwargs: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        return await self._send_request(
+            self.httpx_client.build_request(
+                'POST',
+                f'{self.url}{target}',
+                json=rpc_request_payload,
+                **(http_kwargs or {}),
+            )
+        )
+
     async def _send_get_request(
         self,
         target: str,
         query_params: dict[str, str],
         http_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
-        try:
-            response = await self.httpx_client.get(
+        return await self._send_request(
+            self.httpx_client.build_request(
+                'GET',
                 f'{self.url}{target}',
                 params=query_params,
                 **(http_kwargs or {}),
             )
-            response.raise_for_status()
-            return response.json()
-        except httpx.HTTPStatusError as e:
-            raise A2AClientHTTPError(e.response.status_code, str(e)) from e
-        except json.JSONDecodeError as e:
-            raise A2AClientJSONError(str(e)) from e
-        except httpx.RequestError as e:
-            raise A2AClientHTTPError(
-                503, f'Network communication error: {e}'
-            ) from e
+        )
 
     async def get_task(
         self,

src/a2a/utils/telemetry.py

@@ -83,7 +83,7 @@ def internal_method(self):
     class _NoOp:
         """A no-op object that absorbs all tracing calls when OpenTelemetry is not installed."""
 
-        def __call__(self, *args: Any, **kwargs: Any) -> '_NoOp':
+        def __call__(self, *args: Any, **kwargs: Any) -> Any:
             return self
 
         def __enter__(self) -> '_NoOp':
@@ -92,7 +92,7 @@ def __enter__(self) -> '_NoOp':
         def __exit__(self, *args: object, **kwargs: Any) -> None:
             pass
 
-        def __getattr__(self, name: str) -> '_NoOp':
+        def __getattr__(self, name: str) -> Any:
             return self
 
     trace = _NoOp()  # type: ignore