feat: 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`.

Author: sokoliva

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

src/a2a/client/client_factory.py

@@ -77,6 +77,7 @@ def _register_defaults(
                 TransportProtocol.jsonrpc,
                 lambda card, url, config, interceptors: JsonRpcTransport(
                     config.httpx_client or httpx.AsyncClient(),
+                    config.extensions or None,
                     card,
                     url,
                     interceptors,
@@ -87,6 +88,7 @@ def _register_defaults(
                 TransportProtocol.http_json,
                 lambda card, url, config, interceptors: RestTransport(
                     config.httpx_client or httpx.AsyncClient(),
+                    config.extensions or None,
                     card,
                     url,
                     interceptors,

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.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.types import (
     AgentCard,
     CancelTaskRequest,
@@ -59,6 +60,7 @@ class JsonRpcTransport(ClientTransport):
     def __init__(
         self,
         httpx_client: httpx.AsyncClient,
+        client_extensions: list[str] | None = None,
         agent_card: AgentCard | None = None,
         url: str | None = None,
         interceptors: list[ClientCallInterceptor] | None = None,
@@ -72,6 +74,7 @@ def __init__(
             raise ValueError('Must provide either agent_card or url')
 
         self.httpx_client = httpx_client
+        self.client_extensions = client_extensions
         self.agent_card = agent_card
         self.interceptors = interceptors or []
         self._needs_extended_card = (
@@ -80,6 +83,20 @@ def __init__(
             else True
         )
 
+    def _update_extension_header(
+        self, http_kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        if self.client_extensions:
+            headers = http_kwargs.get('headers', {})
+            existing_extensions = headers.get(HTTP_EXTENSION_HEADER, '')
+            split = (
+                existing_extensions.split(', ') if existing_extensions else []
+            )
+            updated_extensions = list(set(self.client_extensions + split))
+            headers[HTTP_EXTENSION_HEADER] = ', '.join(updated_extensions)
+            http_kwargs['headers'] = headers
+        return http_kwargs
+
     async def _apply_interceptors(
         self,
         method_name: str,
@@ -122,6 +139,7 @@ async def send_message(
             self._get_http_args(context),
             context,
         )
+        modified_kwargs = self._update_extension_header(modified_kwargs)
         response_data = await self._send_request(payload, modified_kwargs)
         response = SendMessageResponse.model_validate(response_data)
         if isinstance(response.root, JSONRPCErrorResponse):
@@ -147,6 +165,7 @@ async def send_message_streaming(
             context,
         )
 
+        modified_kwargs = self._update_extension_header(modified_kwargs)
         modified_kwargs.setdefault(
             'timeout', self.httpx_client.timeout.as_dict().get('read', None)
         )

src/a2a/client/transports/rest.py

@@ -13,6 +13,7 @@
 from a2a.client.errors import A2AClientHTTPError, A2AClientJSONError
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.transports.base import ClientTransport
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.grpc import a2a_pb2
 from a2a.types import (
     AgentCard,
@@ -40,6 +41,7 @@ class RestTransport(ClientTransport):
     def __init__(
         self,
         httpx_client: httpx.AsyncClient,
+        client_extensions: list[str] | None = None,
         agent_card: AgentCard | None = None,
         url: str | None = None,
         interceptors: list[ClientCallInterceptor] | None = None,
@@ -54,6 +56,7 @@ def __init__(
         if self.url.endswith('/'):
             self.url = self.url[:-1]
         self.httpx_client = httpx_client
+        self.client_extensions = client_extensions
         self.agent_card = agent_card
         self.interceptors = interceptors or []
         self._needs_extended_card = (
@@ -62,6 +65,20 @@ def __init__(
             else True
         )
 
+    def _update_extension_header(
+        self, http_kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        if self.client_extensions:
+            headers = http_kwargs.get('headers', {})
+            existing_extensions = headers.get(HTTP_EXTENSION_HEADER, '')
+            split = (
+                existing_extensions.split(', ') if existing_extensions else []
+            )
+            updated_extensions = list(set(self.client_extensions + split))
+            headers[HTTP_EXTENSION_HEADER] = ', '.join(updated_extensions)
+            http_kwargs['headers'] = headers
+        return http_kwargs
+
     async def _apply_interceptors(
         self,
         request_payload: dict[str, Any],
@@ -98,6 +115,7 @@ async def _prepare_send_message(
             self._get_http_args(context),
             context,
         )
+        modified_kwargs = self._update_extension_header(modified_kwargs)
         return payload, modified_kwargs
 
     async def send_message(

tests/client/test_jsonrpc_client.py

@@ -17,6 +17,7 @@
     create_text_message_object,
 )
 from a2a.client.transports.jsonrpc import JsonRpcTransport
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.types import (
     AgentCapabilities,
     AgentCard,
@@ -785,3 +786,181 @@ async def test_close(self, mock_httpx_client: AsyncMock):
         )
         await client.close()
         mock_httpx_client.aclose.assert_called_once()
+
+
+class TestJsonRpcTransportExtensions:
+    def test_update_extension_header_no_initial_headers(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        extensions = ['test_extension_1', 'test_extension_2']
+        client = JsonRpcTransport(
+            mock_httpx_client, extensions, mock_agent_card
+        )
+        http_kwargs = {}
+        result_kwargs = client._update_extension_header(http_kwargs)
+        actual_extensions = set(
+            result_kwargs['headers'][HTTP_EXTENSION_HEADER].split(', ')
+        )
+        expected_extensions = {'test_extension_1', 'test_extension_2'}
+        assert actual_extensions == expected_extensions
+
+    def test_update_extension_header_with_existing_other_headers(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        extensions = ['test_extension_1']
+        client = JsonRpcTransport(
+            mock_httpx_client, extensions, mock_agent_card
+        )
+        http_kwargs = {'headers': {'X_Other': 'Test'}}
+        result_kwargs = client._update_extension_header(http_kwargs)
+        assert (
+            result_kwargs['headers'][HTTP_EXTENSION_HEADER]
+            == 'test_extension_1'
+        )
+        assert result_kwargs['headers']['X_Other'] == 'Test'
+
+    def test_update_extension_header_merge_with_existing_extensions(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        extensions = ['test_extension_1', 'test_extension_2']
+        client = JsonRpcTransport(
+            mock_httpx_client, extensions, mock_agent_card
+        )
+        http_kwargs = {
+            'headers': {
+                HTTP_EXTENSION_HEADER: 'test_extension_2, test_extension_3'
+            }
+        }
+        result_kwargs = client._update_extension_header(http_kwargs)
+        actual_extensions_list = result_kwargs['headers'][
+            HTTP_EXTENSION_HEADER
+        ].split(', ')
+        actual_extensions = set(actual_extensions_list)
+        expected_extensions = {
+            'test_extension_1',
+            'test_extension_2',
+            'test_extension_3',
+        }
+        assert len(actual_extensions_list) == 3
+        assert actual_extensions == expected_extensions
+
+    def test_update_extension_header_no_client_extensions(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        client = JsonRpcTransport(mock_httpx_client, None, mock_agent_card)
+        http_kwargs = {'headers': {'X_Other': 'Test'}}
+        result_kwargs = client._update_extension_header(http_kwargs)
+        assert HTTP_EXTENSION_HEADER not in result_kwargs['headers']
+        assert result_kwargs['headers']['X_Other'] == 'Test'
+
+    def test_update_extension_header_empty_client_extensions(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        client = JsonRpcTransport(mock_httpx_client, [], mock_agent_card)
+        http_kwargs = {'headers': {'X_Other': 'Test'}}
+        result_kwargs = client._update_extension_header(http_kwargs)
+        assert HTTP_EXTENSION_HEADER not in result_kwargs['headers']
+        assert result_kwargs['headers']['X_Other'] == 'Test'
+
+    @pytest.mark.asyncio
+    async def test_send_message_with_extensions(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        """Test that send_message adds extension headers when client_extensions are provided."""
+        extensions = ['test_extension_1', 'test_extension_2']
+        client = JsonRpcTransport(
+            httpx_client=mock_httpx_client,
+            client_extensions=extensions,
+            agent_card=mock_agent_card,
+        )
+        params = MessageSendParams(
+            message=create_text_message_object(content='Hello')
+        )
+        success_response = create_text_message_object(
+            role=Role.agent, content='Hi there!'
+        )
+        rpc_response = SendMessageSuccessResponse(
+            id='123', jsonrpc='2.0', result=success_response
+        )
+        # Mock the response from httpx_client.post
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = rpc_response.model_dump(mode='json')
+        mock_httpx_client.post.return_value = mock_response
+
+        await client.send_message(request=params)
+
+        mock_httpx_client.post.assert_called_once()
+        _, mock_kwargs = mock_httpx_client.post.call_args
+        headers = mock_kwargs.get('headers', {})
+        assert HTTP_EXTENSION_HEADER in headers
+        actual_extensions = set(headers[HTTP_EXTENSION_HEADER].split(', '))
+        expected_extensions = {'test_extension_1', 'test_extension_2'}
+        assert actual_extensions == expected_extensions
+
+    @pytest.mark.asyncio
+    async def test_send_message_no_extensions(
+        self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock
+    ):
+        """Test that send_message does not add extension headers when client_extensions is None."""
+        client = JsonRpcTransport(
+            httpx_client=mock_httpx_client,
+            client_extensions=None,
+            agent_card=mock_agent_card,
+        )
+        params = MessageSendParams(
+            message=create_text_message_object(content='Hello')
+        )
+        success_response = create_text_message_object(
+            role=Role.agent, content='Hi there!'
+        )
+        rpc_response = SendMessageSuccessResponse(
+            id='123', jsonrpc='2.0', result=success_response
+        )
+        # Mock the response from httpx_client.post
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = rpc_response.model_dump(mode='json')
+        mock_httpx_client.post.return_value = mock_response
+
+        await client.send_message(request=params)
+
+        mock_httpx_client.post.assert_called_once()
+        _, mock_kwargs = mock_httpx_client.post.call_args
+        headers = mock_kwargs.get('headers', {})
+        assert HTTP_EXTENSION_HEADER not in headers
+
+    @pytest.mark.asyncio
+    @patch('a2a.client.transports.jsonrpc.aconnect_sse')
+    async def test_send_message_streaming_with_extensions(
+        self,
+        mock_aconnect_sse: AsyncMock,
+        mock_httpx_client: AsyncMock,
+        mock_agent_card: MagicMock,
+    ):
+        """Test X-A2A-Extensions header in send_message_streaming."""
+        extensions = ['test_extension']
+        client = JsonRpcTransport(
+            httpx_client=mock_httpx_client,
+            client_extensions=extensions,
+            agent_card=mock_agent_card,
+        )
+        params = MessageSendParams(
+            message=create_text_message_object(content='Hello stream')
+        )
+
+        mock_event_source = AsyncMock(spec=EventSource)
+        mock_event_source.aiter_sse.return_value = async_iterable_from_list([])
+        mock_aconnect_sse.return_value.__aenter__.return_value = (
+            mock_event_source
+        )
+
+        async for _ in client.send_message_streaming(request=params):
+            pass
+
+        mock_aconnect_sse.assert_called_once()
+        _, kwargs = mock_aconnect_sse.call_args
+
+        headers = kwargs.get('headers', {})
+        assert HTTP_EXTENSION_HEADER in headers
+        assert headers[HTTP_EXTENSION_HEADER] == 'test_extension'