Address mypy errors

Author: mikeas1

src/a2a/client/__init__.py

@@ -9,10 +9,7 @@
 )
 from a2a.client.card_resolver import A2ACardResolver
 from a2a.client.client import Client, ClientConfig, ClientEvent, Consumer
-from a2a.client.client_factory import (
-    ClientFactory,
-    minimal_agent_card,
-)
+from a2a.client.client_factory import ClientFactory, minimal_agent_card
 from a2a.client.errors import (
     A2AClientError,
     A2AClientHTTPError,
@@ -27,7 +24,7 @@
 logger = logging.getLogger(__name__)
 
 try:
-    from a2a.client.legacy_grpc import A2AGrpcClient
+    from a2a.client.legacy_grpc import A2AGrpcClient  # type: ignore
 except ImportError as e:
     _original_error = e
     logger.debug(

src/a2a/client/client_factory.py

@@ -22,10 +22,8 @@
 
 try:
     from a2a.client.transports.grpc import GrpcTransport
-    from a2a.grpc import a2a_pb2_grpc
 except ImportError:
-    GrpcTransport = None
-    a2a_pb2_grpc = None
+    GrpcTransport = None  # type: ignore
 
 
 logger = logging.getLogger(__name__)
@@ -63,37 +61,42 @@ def __init__(
         self._config = config
         self._consumers = consumers
         self._registry: dict[str, TransportProducer] = {}
-        self._register_defaults()
-
-    def _register_defaults(self) -> None:
-        self.register(
-            TransportProtocol.jsonrpc,
-            lambda card, url, config, interceptors: JsonRpcTransport(
-                config.httpx_client or httpx.AsyncClient(),
-                card,
-                url,
-                interceptors,
-            ),
-        )
-        self.register(
-            TransportProtocol.http_json,
-            lambda card, url, config, interceptors: RestTransport(
-                config.httpx_client or httpx.AsyncClient(),
-                card,
-                url,
-                interceptors,
-            ),
-        )
-        if GrpcTransport:
+        self._register_defaults(config.supported_transports)
+
+    def _register_defaults(
+        self, supported: list[str | TransportProtocol]
+    ) -> None:
+        # Empty support list implies JSON-RPC only.
+        if TransportProtocol.jsonrpc in supported or not supported:
             self.register(
-                TransportProtocol.grpc,
-                lambda card, url, config, interceptors: GrpcTransport(
-                    a2a_pb2_grpc.A2AServiceStub(
-                        config.grpc_channel_factory(url)
-                    ),
+                TransportProtocol.jsonrpc,
+                lambda card, url, config, interceptors: JsonRpcTransport(
+                    config.httpx_client or httpx.AsyncClient(),
                     card,
+                    url,
+                    interceptors,
                 ),
             )
+        if TransportProtocol.http_json in supported:
+            self.register(
+                TransportProtocol.http_json,
+                lambda card, url, config, interceptors: RestTransport(
+                    config.httpx_client or httpx.AsyncClient(),
+                    card,
+                    url,
+                    interceptors,
+                ),
+            )
+        if TransportProtocol.grpc in supported:
+            if GrpcTransport is None:
+                raise ImportError(
+                    'To use GrpcClient, its dependencies must be installed. '
+                    'You can install them with \'pip install "a2a-sdk[grpc]"\''
+                )
+            self.register(
+                TransportProtocol.grpc,
+                GrpcTransport.create,
+            )
 
     def register(self, label: str, generator: TransportProducer) -> None:
         """Register a new transport producer for a given transport label."""

src/a2a/client/legacy.py

@@ -268,7 +268,7 @@ async def get_task_callback(
         params = request.params
         if isinstance(params, TaskIdParams):
             params = GetTaskPushNotificationConfigParams(
-                id=request.params.task_id
+                id=request.params.id
             )
         try:
             result = await self._transport.get_task_callback(

src/a2a/client/transports/__init__.py

@@ -8,7 +8,7 @@
 try:
     from a2a.client.transports.grpc import GrpcTransport
 except ImportError:
-    GrpcTransport = None
+    GrpcTransport = None  # type: ignore
 
 
 __all__ = [

src/a2a/client/transports/grpc.py

@@ -5,14 +5,17 @@
 
 try:
     import grpc
+
+    from grpc.aio import Channel
 except ImportError as e:
     raise ImportError(
         'A2AGrpcClient requires grpcio and grpcio-tools to be installed. '
         'Install with: '
         "'pip install a2a-sdk[grpc]'"
     ) from e
 
-from a2a.client.middleware import ClientCallContext
+from a2a.client.client import ClientConfig
+from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
@@ -40,18 +43,35 @@ class GrpcTransport(ClientTransport):
 
     def __init__(
         self,
-        grpc_stub: a2a_pb2_grpc.A2AServiceStub,
+        channel: Channel,
         agent_card: AgentCard | None,
     ):
         """Initializes the GrpcTransport."""
         self.agent_card = agent_card
-        self.stub = grpc_stub
+        self.channel = channel
+        self.stub = a2a_pb2_grpc.A2AServiceStub(channel)
         self._needs_extended_card = (
             agent_card.supports_authenticated_extended_card
             if agent_card
             else True
         )
 
+    @classmethod
+    def create(
+        cls,
+        card: AgentCard,
+        url: str,
+        config: ClientConfig,
+        interceptors: list[ClientCallInterceptor],
+    ) -> 'GrpcTransport':
+        """Creates a gRPC transport for the A2A client."""
+        if config.grpc_channel_factory is None:
+            raise ValueError('grpc_channel_factory is required when using gRPC')
+        return cls(
+            config.grpc_channel_factory(url),
+            card,
+        )
+
     async def send_message(
         self,
         request: MessageSendParams,
@@ -189,5 +209,4 @@ async def get_card(
 
     async def close(self) -> None:
         """Closes the gRPC channel."""
-        if hasattr(self.stub, 'close'):
-            await self.stub.close()
+        await self.channel.close()

src/a2a/server/apps/rest/rest_adapter.py

@@ -135,7 +135,7 @@ def routes(self) -> dict[tuple[str, str], Callable[[Request], Any]]:
             A dictionary where each key is a tuple of (path, http_method) and
             the value is the callable handler for that route.
         """
-        routes = {
+        routes: dict[tuple[str, str], Callable[[Request], Any]] = {
             ('/v1/message:send', 'POST'): functools.partial(
                 self._handle_request, self.handler.on_message_send
             ),

src/a2a/server/request_handlers/rest_handler.py

@@ -176,12 +176,9 @@ async def get_push_notification(
         """
         task_id = request.path_params['id']
         push_id = request.path_params['push_id']
-        if push_id:
-            params = GetTaskPushNotificationConfigParams(
-                id=task_id, push_notification_config_id=push_id
-            )
-        else:
-            params = TaskIdParams(id=task_id)
+        params = GetTaskPushNotificationConfigParams(
+            id=task_id, push_notification_config_id=push_id
+        )
         config = (
             await self.request_handler.on_get_task_push_notification_config(
                 params, context