feat: add support for extensions in Client and BaseClient, update transport methods to handle extensions

sokoliva · sokoliva · commit a9aa9ee846e6 · 2025-11-03T15:42:15.000Z
diff --git a/src/a2a/client/base_client.py b/src/a2a/client/base_client.py
@@ -36,11 +36,14 @@ 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
         self._config = config
         self._transport = transport
+        if self._extensions:
+            self._config.extensions = self._extensions
 
     async def send_message(
         self,
diff --git a/src/a2a/client/client.py b/src/a2a/client/client.py
@@ -93,7 +93,8 @@ def __init__(
         self,
         consumers: list[Consumer] | None = None,
         middleware: list[ClientCallInterceptor] | None = None,
-        # iva todo add optional extensions- it can override value from the config, if it is provided
+        # iva todo - it can override value from the config, if it is provided
+        extensions: list[str] | None = None,
     ):
         """Initializes the client with consumers and middleware.
 
@@ -105,8 +106,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(
diff --git a/src/a2a/client/client_factory.py b/src/a2a/client/client_factory.py
@@ -115,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.
 
@@ -168,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."""
@@ -179,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`.
 
@@ -228,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,
         )
 
 
diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py
@@ -220,6 +220,9 @@ async def get_task(
             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):
@@ -240,6 +243,9 @@ async def cancel_task(
             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):
@@ -262,6 +268,9 @@ async def set_task_callback(
             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
@@ -286,6 +295,9 @@ async def get_task_callback(
             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
@@ -310,7 +322,9 @@ async def resubscribe(
             get_http_args(context),
             context,
         )
-
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         modified_kwargs.setdefault('timeout', None)
 
         async with aconnect_sse(
@@ -366,6 +380,9 @@ async def get_card(
             get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
 
         response_data = await self._send_request(
             payload,
diff --git a/src/a2a/client/transports/rest.py b/src/a2a/client/transports/rest.py
@@ -212,6 +212,9 @@ async def get_task(
             get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_get_request(
             f'/v1/tasks/{request.id}',
             {'historyLength': str(request.history_length)}
@@ -237,6 +240,9 @@ async def cancel_task(
             get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_post_request(
             f'/v1/tasks/{request.id}:cancel', payload, modified_kwargs
         )
@@ -260,6 +266,9 @@ async def set_task_callback(
         payload, modified_kwargs = await self._apply_interceptors(
             payload, get_http_args(context), context
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_post_request(
             f'/v1/tasks/{request.task_id}/pushNotificationConfigs',
             payload,
@@ -285,6 +294,9 @@ async def get_task_callback(
             get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_get_request(
             f'/v1/tasks/{request.id}/pushNotificationConfigs/{request.push_notification_config_id}',
             {},
@@ -305,12 +317,13 @@ async def resubscribe(
         """Reconnects to get task updates."""
         http_kwargs = get_http_args(context) or {}
         http_kwargs.setdefault('timeout', None)
+        modified_kwargs = update_extension_header(http_kwargs, self.extensions)
 
         async with aconnect_sse(
             self.httpx_client,
             'GET',
             f'{self.url}/v1/tasks/{request.id}:subscribe',
-            **http_kwargs,
+            **modified_kwargs,
         ) as event_source:
             try:
                 async for sse in event_source.aiter_sse():
@@ -353,6 +366,9 @@ async def get_card(
             get_http_args(context),
             context,
         )
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_get_request(
             '/v1/card', {}, modified_kwargs
         )
diff --git a/tests/client/test_base_client.py b/tests/client/test_base_client.py
@@ -55,6 +55,7 @@ def base_client(sample_agent_card, mock_transport):
         transport=mock_transport,
         consumers=[],
         middleware=[],
+        extensions=[],
     )
 
 
diff --git a/tests/client/transports/test_utils.py b/tests/client/transports/test_utils.py
@@ -5,70 +5,66 @@
 
 
 class TestUtils:
-    def test_update_extension_header_no_initial_headers(self):
-        extensions = ['test_extension_1', 'test_extension_2']
-
-        http_kwargs = {}
-        result_kwargs = update_extension_header(http_kwargs, extensions)
-        header_value = result_kwargs['headers'][HTTP_EXTENSION_HEADER]
-        actual_extensions_list = [e.strip() for e in header_value.split(',')]
-        actual_extensions = set(actual_extensions_list)
-
-        expected_extensions = {
-            'test_extension_1',
-            'test_extension_2',
-        }
-        assert len(actual_extensions_list) == 2
-        assert actual_extensions == expected_extensions
-
     @pytest.mark.parametrize(
-        'existing_header, expected_count',
+        'extensions, existing_header, expected_extensions, expected_count',
         [
-            ('test_extension_1, test_extension_2', 3),
-            ('test_extension_1,test_extension_2', 3),
-            ('test_extension_1', 3),
+            (
+                ['test_extension_1', 'test_extension_2'],
+                '',
+                {
+                    'test_extension_1',
+                    'test_extension_2',
+                },
+                2,
+            ),
+            (
+                ['test_extension_1', 'test_extension_2'],
+                'test_extension_2, test_extension_3',
+                {
+                    'test_extension_1',
+                    'test_extension_2',
+                    'test_extension_3',
+                },
+                3,
+            ),
+            (
+                ['test_extension_1', 'test_extension_2'],
+                'test_extension_3',
+                {
+                    'test_extension_1',
+                    'test_extension_2',
+                    'test_extension_3',
+                },
+                3,
+            ),
         ],
     )
     def test_update_extension_header_merge_with_existing_extensions(
         self,
+        extensions: list[str],
         existing_header: str,
+        expected_extensions: set[str],
         expected_count: int,
     ):
-        extensions = ['test_extension_2', 'test_extension_3']
         http_kwargs = {'headers': {HTTP_EXTENSION_HEADER: existing_header}}
         result_kwargs = update_extension_header(http_kwargs, extensions)
 
         header_value = result_kwargs['headers'][HTTP_EXTENSION_HEADER]
         actual_extensions_list = [e.strip() for e in header_value.split(',')]
         actual_extensions = set(actual_extensions_list)
 
-        expected_extensions = {
-            'test_extension_1',
-            'test_extension_2',
-            'test_extension_3',
-        }
         assert len(actual_extensions_list) == expected_count
         assert actual_extensions == expected_extensions
 
     def test_update_extension_header_with_other_headers(self):
-        extensions = ['test_extension_1']
+        extensions = ['test_extension']
         http_kwargs = {'headers': {'X_Other': 'Test'}}
         result_kwargs = update_extension_header(http_kwargs, extensions)
         headers = result_kwargs.get('headers', {})
         assert HTTP_EXTENSION_HEADER in headers
-        assert headers[HTTP_EXTENSION_HEADER] == 'test_extension_1'
+        assert headers[HTTP_EXTENSION_HEADER] == 'test_extension'
         assert headers['X_Other'] == 'Test'
 
-    def test_update_extension_header_with_existing_other_headers(self):
-        extensions = ['test_extension_1']
-        http_kwargs = {'headers': {'X_Other': 'Test'}}
-        result_kwargs = update_extension_header(http_kwargs, extensions)
-        assert (
-            result_kwargs['headers'][HTTP_EXTENSION_HEADER]
-            == 'test_extension_1'
-        )
-        assert result_kwargs['headers']['X_Other'] == 'Test'
-
     def test_update_extension_header_no_extensions(self):
         http_kwargs = {'headers': {'X_Other': 'Test'}}
         result_kwargs = update_extension_header(http_kwargs, None)

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@ def base_client(sample_agent_card, mock_transport):`
`55`	`55`	`transport=mock_transport,`
`56`	`56`	`consumers=[],`
`57`	`57`	`middleware=[],`
	`58`	`+ extensions=[],`
`58`	`59`	`)`
`59`	`60`
`60`	`61`