feat: Add client-side extension support

src/a2a/client/base_client.py

@@ -36,9 +36,11 @@ def __init__(
         transport: ClientTransport,
         consumers: list[Consumer],
         middleware: list[ClientCallInterceptor],
+        extensions: list[str],
     ):
-        super().__init__(consumers, middleware)
+        super().__init__(consumers, middleware, extensions)
         self._card = card
+        config.extensions = extensions
         self._config = config
         self._transport = transport
 

src/a2a/client/client.py

@@ -67,6 +67,9 @@ class ClientConfig:
     )
     """Push notification callbacks to use for every request."""
 
+    extensions: list[str] = dataclasses.field(default_factory=list)
+    """A list of extension URIs the client supports."""
+
 
 UpdateEvent = TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None
 # Alias for emitted events from client
@@ -90,6 +93,7 @@ def __init__(
         self,
         consumers: list[Consumer] | None = None,
         middleware: list[ClientCallInterceptor] | None = None,
+        extensions: list[str] | None = None,
     ):
         """Initializes the client with consumers and middleware.
 
@@ -101,8 +105,11 @@ def __init__(
             middleware = []
         if consumers is None:
             consumers = []
+        if extensions is None:
+            extensions = []
         self._consumers = consumers
         self._middleware = middleware
+        self._extensions = extensions
 
     @abstractmethod
     async def send_message(

src/a2a/client/client_factory.py

@@ -80,6 +80,7 @@ def _register_defaults(
                     card,
                     url,
                     interceptors,
+                    config.extensions or None,
                 ),
             )
         if TransportProtocol.http_json in supported:
@@ -90,6 +91,7 @@ def _register_defaults(
                     card,
                     url,
                     interceptors,
+                    config.extensions or None,
                 ),
             )
         if TransportProtocol.grpc in supported:
@@ -113,6 +115,7 @@ async def connect(  # noqa: PLR0913
         relative_card_path: str | None = None,
         resolver_http_kwargs: dict[str, Any] | None = None,
         extra_transports: dict[str, TransportProducer] | None = None,
+        extensions: list[str] | None = None,
     ) -> Client:
         """Convenience method for constructing a client.
 
@@ -166,7 +169,7 @@ async def connect(  # noqa: PLR0913
         factory = cls(client_config)
         for label, generator in (extra_transports or {}).items():
             factory.register(label, generator)
-        return factory.create(card, consumers, interceptors)
+        return factory.create(card, consumers, interceptors, extensions)
 
     def register(self, label: str, generator: TransportProducer) -> None:
         """Register a new transport producer for a given transport label."""
@@ -177,6 +180,7 @@ def create(
         card: AgentCard,
         consumers: list[Consumer] | None = None,
         interceptors: list[ClientCallInterceptor] | None = None,
+        extensions: list[str] | None = None,
     ) -> Client:
         """Create a new `Client` for the provided `AgentCard`.
 
@@ -226,12 +230,22 @@ def create(
         if consumers:
             all_consumers.extend(consumers)
 
+        all_extensions = self._config.extensions.copy()
+        if extensions:
+            all_extensions.extend(extensions)
+            self._config.extensions = all_extensions
+
         transport = self._registry[transport_protocol](
             card, transport_url, self._config, interceptors or []
         )
 
         return BaseClient(
-            card, self._config, transport, all_consumers, interceptors or []
+            card,
+            self._config,
+            transport,
+            all_consumers,
+            interceptors or [],
+            all_extensions,
         )
 
 

src/a2a/client/transports/grpc.py

@@ -12,10 +12,12 @@
         "'pip install a2a-sdk[grpc]'"
     ) from e
 
+
 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.client.transports.utils import update_extension_metadata
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
     AgentCard,
@@ -44,6 +46,7 @@ def __init__(
         self,
         channel: Channel,
         agent_card: AgentCard | None,
+        extensions: list[str] | None = None,
     ):
         """Initializes the GrpcTransport."""
         self.agent_card = agent_card
@@ -54,6 +57,7 @@ def __init__(
             if agent_card
             else True
         )
+        self.extensions = extensions
 
     @classmethod
     def create(
@@ -66,10 +70,7 @@ def create(
         """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,
-        )
+        return cls(config.grpc_channel_factory(url), card, config.extensions)
 
     async def send_message(
         self,
@@ -84,8 +85,10 @@ async def send_message(
                 configuration=proto_utils.ToProto.message_send_configuration(
                     request.configuration
                 ),
-                metadata=proto_utils.ToProto.metadata(request.metadata),
-            )
+                metadata=update_extension_metadata(
+                    request.metadata, self.extensions
+                ),
+            ),
         )
         if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
@@ -106,8 +109,10 @@ async def send_message_streaming(
                 configuration=proto_utils.ToProto.message_send_configuration(
                     request.configuration
                 ),
-                metadata=proto_utils.ToProto.metadata(request.metadata),
-            )
+                metadata=update_extension_metadata(
+                    request.metadata, self.extensions
+                ),
+            ),
         )
         while True:
             response = await stream.read()
@@ -122,7 +127,10 @@ async def resubscribe(
     ]:
         """Reconnects to get task updates."""
         stream = self.stub.TaskSubscription(
-            a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}')
+            a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}'),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         while True:
             response = await stream.read()
@@ -141,7 +149,10 @@ async def get_task(
             a2a_pb2.GetTaskRequest(
                 name=f'tasks/{request.id}',
                 history_length=request.history_length,
-            )
+            ),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task(task)
 
@@ -153,7 +164,7 @@ async def cancel_task(
     ) -> Task:
         """Requests the agent to cancel a specific task."""
         task = await self.stub.CancelTask(
-            a2a_pb2.CancelTaskRequest(name=f'tasks/{request.id}')
+            a2a_pb2.CancelTaskRequest(name=f'tasks/{request.id}'),
         )
         return proto_utils.FromProto.task(task)
 
@@ -171,7 +182,10 @@ async def set_task_callback(
                 config=proto_utils.ToProto.task_push_notification_config(
                     request
                 ),
-            )
+            ),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
@@ -185,7 +199,10 @@ async def get_task_callback(
         config = await self.stub.GetTaskPushNotificationConfig(
             a2a_pb2.GetTaskPushNotificationConfigRequest(
                 name=f'tasks/{request.id}/pushNotificationConfigs/{request.push_notification_config_id}',
-            )
+            ),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 

src/a2a/client/transports/jsonrpc.py

@@ -18,6 +18,7 @@
 )
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
+from a2a.client.transports.utils import get_http_args, update_extension_header
 from a2a.types import (
     AgentCard,
     CancelTaskRequest,
@@ -62,6 +63,7 @@ def __init__(
         agent_card: AgentCard | None = None,
         url: str | None = None,
         interceptors: list[ClientCallInterceptor] | None = None,
+        extensions: list[str] | None = None,
     ):
         """Initializes the JsonRpcTransport."""
         if url:
@@ -79,6 +81,7 @@ def __init__(
             if agent_card
             else True
         )
+        self.extensions = extensions
 
     async def _apply_interceptors(
         self,
@@ -103,11 +106,6 @@ async def _apply_interceptors(
             )
         return final_request_payload, final_http_kwargs
 
-    def _get_http_args(
-        self, context: ClientCallContext | None
-    ) -> dict[str, Any] | None:
-        return context.state.get('http_kwargs') if context else None
-
     async def send_message(
         self,
         request: MessageSendParams,
@@ -119,9 +117,12 @@ async def send_message(
         payload, modified_kwargs = await self._apply_interceptors(
             'message/send',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = SendMessageResponse.model_validate(response_data)
         if isinstance(response.root, JSONRPCErrorResponse):
@@ -143,10 +144,13 @@ async def send_message_streaming(
         payload, modified_kwargs = await self._apply_interceptors(
             'message/stream',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
 
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         modified_kwargs.setdefault(
             'timeout', self.httpx_client.timeout.as_dict().get('read', None)
         )
@@ -213,9 +217,12 @@ async def get_task(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/get',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = GetTaskResponse.model_validate(response_data)
         if isinstance(response.root, JSONRPCErrorResponse):
@@ -233,9 +240,12 @@ async def cancel_task(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/cancel',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = CancelTaskResponse.model_validate(response_data)
         if isinstance(response.root, JSONRPCErrorResponse):
@@ -255,9 +265,12 @@ async def set_task_callback(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/pushNotificationConfig/set',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = SetTaskPushNotificationConfigResponse.model_validate(
             response_data
@@ -279,9 +292,12 @@ async def get_task_callback(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/pushNotificationConfig/get',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = GetTaskPushNotificationConfigResponse.model_validate(
             response_data
@@ -303,10 +319,12 @@ async def resubscribe(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/resubscribe',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
-
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         modified_kwargs.setdefault('timeout', None)
 
         async with aconnect_sse(
@@ -345,7 +363,7 @@ async def get_card(
         if not card:
             resolver = A2ACardResolver(self.httpx_client, self.url)
             card = await resolver.get_agent_card(
-                http_kwargs=self._get_http_args(context)
+                http_kwargs=get_http_args(context)
             )
             self._needs_extended_card = (
                 card.supports_authenticated_extended_card
@@ -359,9 +377,12 @@ async def get_card(
         payload, modified_kwargs = await self._apply_interceptors(
             request.method,
             request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
 
         response_data = await self._send_request(
             payload,