feat: Add client-side extension support

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,
@@ -56,14 +57,15 @@
 class JsonRpcTransport(ClientTransport):
     """A JSON-RPC transport for the A2A client."""
 
     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,
     ):
         """Initializes the JsonRpcTransport."""
         if url:
             self.url = url
         elif agent_card:
@@ -71,26 +73,41 @@
         else:
             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 = (
             agent_card.supports_authenticated_extended_card
             if agent_card
             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,
         request_payload: dict[str, Any],
         http_kwargs: dict[str, Any] | None,
         context: ClientCallContext | None,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         final_http_kwargs = http_kwargs or {}
         final_request_payload = request_payload
 
         for interceptor in self.interceptors:
             (
                 final_request_payload,
                 final_http_kwargs,
@@ -122,6 +139,7 @@
             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 @@
             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,38 +13,40 @@
 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,
     GetTaskPushNotificationConfigParams,
     Message,
     MessageSendParams,
     Task,
     TaskArtifactUpdateEvent,
     TaskIdParams,
     TaskPushNotificationConfig,
     TaskQueryParams,
     TaskStatusUpdateEvent,
 )
 from a2a.utils import proto_utils
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
 @trace_class(kind=SpanKind.CLIENT)
 class RestTransport(ClientTransport):
     """A REST transport for the A2A client."""
 
     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,
     ):
         """Initializes the RestTransport."""
         if url:
             self.url = url
         elif agent_card:
@@ -53,24 +55,39 @@
             raise ValueError('Must provide either agent_card or url')
         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 = (
             agent_card.supports_authenticated_extended_card
             if agent_card
             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],
         http_kwargs: dict[str, Any] | None,
         context: ClientCallContext | None,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         final_http_kwargs = http_kwargs or {}
         final_request_payload = request_payload
         # TODO: Implement interceptors for other transports
         return final_request_payload, final_http_kwargs
 
     def _get_http_args(
@@ -98,6 +115,7 @@
             self._get_http_args(context),
             context,
         )
+        modified_kwargs = self._update_extension_header(modified_kwargs)
         return payload, modified_kwargs
 
     async def send_message(

tests/README.md

@@ -2,7 +2,7 @@
 
 1. Run the tests
     ```bash
-    uv run pytest -v -s client/test_client.py
+    uv run pytest -v -s client/test_client_factory.py
     ```
 
 In case of failures, you can cleanup the cache:

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'