feat: Add client-side extension support (#525)

feat(rest, jsonrpc): Add client-side extension support
This commit introduces support for clients to declare the extensions
they support.

- Adds an `extensions` list to `ClientConfig`.
- Updates `ClientFactory` to pass `client_extensions` to
`JsonRpcTransport` and `RestTransport`.
- Adds `_update_extension_header` method to both transports to update
the `X-A2A-Extensions` header.
- Modifies `send_message` and `send_message_streaming` in
`JsonRpcTransport` to include the extension headers.
- Modifies `_prepare_send_message` in `RestTransport` to include the
extension headers.
- Adds tests for the extension header logic in both JSON-RPC and REST
transports, including a new test file `test_rest_client.py`.
-
Fixes #504 🦕

Author: sokoliva

src/a2a/client/base_client.py

@@ -49,6 +49,7 @@ async def send_message(
         *,
         context: ClientCallContext | None = None,
         request_metadata: dict[str, Any] | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncIterator[ClientEvent | Message]:
         """Sends a message to the agent.
 
@@ -60,6 +61,7 @@ async def send_message(
             request: The message to send to the agent.
             context: The client call context.
             request_metadata: Extensions Metadata attached to the request.
+            extensions: List of extensions to be activated.
 
         Yields:
             An async iterator of `ClientEvent` or a final `Message` response.
@@ -79,7 +81,7 @@ async def send_message(
 
         if not self._config.streaming or not self._card.capabilities.streaming:
             response = await self._transport.send_message(
-                params, context=context
+                params, context=context, extensions=extensions
             )
             result = (
                 (response, None) if isinstance(response, Task) else response
@@ -89,7 +91,9 @@ async def send_message(
             return
 
         tracker = ClientTaskManager()
-        stream = self._transport.send_message_streaming(params, context=context)
+        stream = self._transport.send_message_streaming(
+            params, context=context, extensions=extensions
+        )
 
         first_event = await anext(stream)
         # The response from a server may be either exactly one Message or a
@@ -126,74 +130,91 @@ async def get_task(
         request: TaskQueryParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Retrieves the current state and history of a specific task.
 
         Args:
             request: The `TaskQueryParams` object specifying the task ID.
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Returns:
             A `Task` object representing the current state of the task.
         """
-        return await self._transport.get_task(request, context=context)
+        return await self._transport.get_task(
+            request, context=context, extensions=extensions
+        )
 
     async def cancel_task(
         self,
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Requests the agent to cancel a specific task.
 
         Args:
             request: The `TaskIdParams` object specifying the task ID.
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Returns:
             A `Task` object containing the updated task status.
         """
-        return await self._transport.cancel_task(request, context=context)
+        return await self._transport.cancel_task(
+            request, context=context, extensions=extensions
+        )
 
     async def set_task_callback(
         self,
         request: TaskPushNotificationConfig,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Sets or updates the push notification configuration for a specific task.
 
         Args:
             request: The `TaskPushNotificationConfig` object with the new configuration.
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Returns:
             The created or updated `TaskPushNotificationConfig` object.
         """
-        return await self._transport.set_task_callback(request, context=context)
+        return await self._transport.set_task_callback(
+            request, context=context, extensions=extensions
+        )
 
     async def get_task_callback(
         self,
         request: GetTaskPushNotificationConfigParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Retrieves the push notification configuration for a specific task.
 
         Args:
             request: The `GetTaskPushNotificationConfigParams` object specifying the task.
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Returns:
             A `TaskPushNotificationConfig` object containing the configuration.
         """
-        return await self._transport.get_task_callback(request, context=context)
+        return await self._transport.get_task_callback(
+            request, context=context, extensions=extensions
+        )
 
     async def resubscribe(
         self,
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncIterator[ClientEvent]:
         """Resubscribes to a task's event stream.
 
@@ -202,6 +223,7 @@ async def resubscribe(
         Args:
             request: Parameters to identify the task to resubscribe to.
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Yields:
             An async iterator of `ClientEvent` objects.
@@ -219,12 +241,15 @@ async def resubscribe(
         # we should never see Message updates, despite the typing of the service
         # definition indicating it may be possible.
         async for event in self._transport.resubscribe(
-            request, context=context
+            request, context=context, extensions=extensions
         ):
             yield await self._process_response(tracker, event)
 
     async def get_card(
-        self, *, context: ClientCallContext | None = None
+        self,
+        *,
+        context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AgentCard:
         """Retrieves the agent's card.
 
@@ -233,11 +258,14 @@ async def get_card(
 
         Args:
             context: The client call context.
+            extensions: List of extensions to be activated.
 
         Returns:
             The `AgentCard` for the agent.
         """
-        card = await self._transport.get_card(context=context)
+        card = await self._transport.get_card(
+            context=context, extensions=extensions
+        )
         self._card = card
         return card
 

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
@@ -111,6 +114,7 @@ async def send_message(
         *,
         context: ClientCallContext | None = None,
         request_metadata: dict[str, Any] | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncIterator[ClientEvent | Message]:
         """Sends a message to the server.
 
@@ -129,6 +133,7 @@ async def get_task(
         request: TaskQueryParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Retrieves the current state and history of a specific task."""
 
@@ -138,6 +143,7 @@ async def cancel_task(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Requests the agent to cancel a specific task."""
 
@@ -147,6 +153,7 @@ async def set_task_callback(
         request: TaskPushNotificationConfig,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Sets or updates the push notification configuration for a specific task."""
 
@@ -156,6 +163,7 @@ async def get_task_callback(
         request: GetTaskPushNotificationConfigParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Retrieves the push notification configuration for a specific task."""
 
@@ -165,14 +173,18 @@ async def resubscribe(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncIterator[ClientEvent]:
         """Resubscribes to a task's event stream."""
         return
         yield
 
     @abstractmethod
     async def get_card(
-        self, *, context: ClientCallContext | None = None
+        self,
+        *,
+        context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AgentCard:
         """Retrieves the agent's card."""
 

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.
 
@@ -142,6 +145,7 @@ async def connect(  # noqa: PLR0913
             A2AAgentCardResolver.get_agent_card as the http_kwargs parameter.
           extra_transports: Additional transport protocols to enable when
             constructing the client.
+          extensions: List of extensions to be activated.
 
         Returns:
           A `Client` object.
@@ -166,7 +170,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 +181,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`.
 
@@ -186,6 +191,7 @@ def create(
           interceptors: A list of interceptors to use for each request. These
             are used for things like attaching credentials or http headers
             to all outbound requests.
+          extensions: List of extensions to be activated.
 
         Returns:
           A `Client` object.
@@ -226,12 +232,21 @@ 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 [],
         )
 
 

src/a2a/client/transports/base.py

@@ -25,6 +25,7 @@ async def send_message(
         request: MessageSendParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task | Message:
         """Sends a non-streaming message request to the agent."""
 
@@ -34,6 +35,7 @@ async def send_message_streaming(
         request: MessageSendParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncGenerator[
         Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
     ]:
@@ -47,6 +49,7 @@ async def get_task(
         request: TaskQueryParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Retrieves the current state and history of a specific task."""
 
@@ -56,6 +59,7 @@ async def cancel_task(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> Task:
         """Requests the agent to cancel a specific task."""
 
@@ -65,6 +69,7 @@ async def set_task_callback(
         request: TaskPushNotificationConfig,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Sets or updates the push notification configuration for a specific task."""
 
@@ -74,6 +79,7 @@ async def get_task_callback(
         request: GetTaskPushNotificationConfigParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> TaskPushNotificationConfig:
         """Retrieves the push notification configuration for a specific task."""
 
@@ -83,6 +89,7 @@ async def resubscribe(
         request: TaskIdParams,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AsyncGenerator[
         Task | Message | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
     ]:
@@ -95,6 +102,7 @@ async def get_card(
         self,
         *,
         context: ClientCallContext | None = None,
+        extensions: list[str] | None = None,
     ) -> AgentCard:
         """Retrieves the AgentCard."""