feat: enhance extension handling across client and transport layers

Author: sokoliva

src/a2a/client/base_client.py

@@ -37,9 +37,8 @@ def __init__(
         transport: ClientTransport,
         consumers: list[Consumer],
         middleware: list[ClientCallInterceptor],
-        extensions: list[str],
     ):
-        super().__init__(consumers, middleware, extensions)
+        super().__init__(consumers, middleware)
         self._card = card
         self._config = config
         self._transport = transport
@@ -50,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.
 
@@ -61,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.
@@ -80,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
@@ -90,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
@@ -127,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.
 
@@ -203,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.
@@ -220,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.
 
@@ -234,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

@@ -93,24 +93,19 @@ 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.
 
         Args:
             consumers: A list of callables to process events from the agent.
             middleware: A list of interceptors to process requests and responses.
-            extensions: A list of extension URIs the client supports.
         """
         if middleware is None:
             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(
@@ -119,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.
 
@@ -137,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."""
 
@@ -146,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."""
 
@@ -155,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."""
 
@@ -164,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."""
 
@@ -173,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

@@ -245,7 +245,6 @@ def create(
             transport,
             all_consumers,
             interceptors or [],
-            all_extensions,
         )
 
 

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."""