feat: move common functions for managing HTTP extension headers to utility.py Add extensions feature to grpc.

Author: sokoliva

src/a2a/client/client.py

@@ -93,6 +93,7 @@ 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
     ):
         """Initializes the client with consumers and middleware.
 
@@ -113,6 +114,8 @@ async def send_message(
         request: Message,
         *,
         context: ClientCallContext | None = None,
+        # iva todo add optional extensions- it can override value from the config, if it is provided
+        # and to the other ones as well
     ) -> AsyncIterator[ClientEvent | Message]:
         """Sends a message to the server.
 

src/a2a/client/transports/grpc.py

@@ -1,6 +1,7 @@
 import logging
 
 from collections.abc import AsyncGenerator
+from typing import Any
 
 
 try:
@@ -16,6 +17,7 @@
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.optionals import Channel
 from a2a.client.transports.base import ClientTransport
+from a2a.client.transports.utils import update_extension_header
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
     AgentCard,
@@ -44,6 +46,7 @@ def __init__(
         self,
         channel: Channel,
         agent_card: AgentCard | None,
+        extensions: list[str] | None = None,
     ):
         """Initializes the GrpcTransport."""
         self.agent_card = agent_card
@@ -54,6 +57,25 @@ def __init__(
             if agent_card
             else True
         )
+        self.extensions = extensions
+
+    def _get_metadata(
+        self, context: ClientCallContext | None
+    ) -> list[tuple[str, str]]:
+        http_kwargs: dict[str, Any] = {}
+        if context and context.state.get('grpc_metadata'):
+            # Convert existing metadata to headers format for update_extension_header
+            http_kwargs['headers'] = {
+                k: v for k, v in context.state['grpc_metadata']
+            }
+
+        updated_kwargs = update_extension_header(http_kwargs, self.extensions)
+
+        metadata = []
+        if 'headers' in updated_kwargs:
+            metadata.extend(updated_kwargs['headers'].items())
+
+        return metadata
 
     @classmethod
     def create(
@@ -66,10 +88,7 @@ def create(
         """Creates a gRPC transport for the A2A client."""
         if config.grpc_channel_factory is None:
             raise ValueError('grpc_channel_factory is required when using gRPC')
-        return cls(
-            config.grpc_channel_factory(url),
-            card,
-        )
+        return cls(config.grpc_channel_factory(url), card, config.extensions)
 
     async def send_message(
         self,
@@ -85,7 +104,8 @@ async def send_message(
                     request.configuration
                 ),
                 metadata=proto_utils.ToProto.metadata(request.metadata),
-            )
+            ),
+            metadata=self._get_metadata(context),
         )
         if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
@@ -107,7 +127,8 @@ async def send_message_streaming(
                     request.configuration
                 ),
                 metadata=proto_utils.ToProto.metadata(request.metadata),
-            )
+            ),
+            metadata=self._get_metadata(context),
         )
         while True:
             response = await stream.read()
@@ -122,7 +143,8 @@ async def resubscribe(
     ]:
         """Reconnects to get task updates."""
         stream = self.stub.TaskSubscription(
-            a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}')
+            a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}'),
+            metadata=self._get_metadata(context),
         )
         while True:
             response = await stream.read()
@@ -141,7 +163,8 @@ async def get_task(
             a2a_pb2.GetTaskRequest(
                 name=f'tasks/{request.id}',
                 history_length=request.history_length,
-            )
+            ),
+            metadata=self._get_metadata(context),
         )
         return proto_utils.FromProto.task(task)
 
@@ -153,7 +176,8 @@ async def cancel_task(
     ) -> Task:
         """Requests the agent to cancel a specific task."""
         task = await self.stub.CancelTask(
-            a2a_pb2.CancelTaskRequest(name=f'tasks/{request.id}')
+            a2a_pb2.CancelTaskRequest(name=f'tasks/{request.id}'),
+            metadata=self._get_metadata(context),
         )
         return proto_utils.FromProto.task(task)
 
@@ -171,7 +195,8 @@ async def set_task_callback(
                 config=proto_utils.ToProto.task_push_notification_config(
                     request
                 ),
-            )
+            ),
+            metadata=self._get_metadata(context),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
@@ -185,7 +210,8 @@ async def get_task_callback(
         config = await self.stub.GetTaskPushNotificationConfig(
             a2a_pb2.GetTaskPushNotificationConfigRequest(
                 name=f'tasks/{request.id}/pushNotificationConfigs/{request.push_notification_config_id}',
-            )
+            ),
+            metadata=self._get_metadata(context),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
@@ -203,6 +229,7 @@ async def get_card(
 
         card_pb = await self.stub.GetAgentCard(
             a2a_pb2.GetAgentCardRequest(),
+            metadata=self._get_metadata(context),  # probaby not needed
         )
         card = proto_utils.FromProto.agent_card(card_pb)
         self.agent_card = card

src/a2a/client/transports/jsonrpc.py

@@ -18,7 +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.client.transports.utils import get_http_args, update_extension_header
 from a2a.types import (
     AgentCard,
     CancelTaskRequest,
@@ -83,25 +83,6 @@ def __init__(
         )
         self.extensions = extensions
 
-    def _update_extension_header(
-        self, http_kwargs: dict[str, Any]
-    ) -> dict[str, Any]:
-        if not self.extensions:
-            return http_kwargs
-
-        headers = http_kwargs.setdefault('headers', {})
-        existing_extensions_str = headers.get(HTTP_EXTENSION_HEADER, '')
-
-        existing_extensions = [
-            e.strip() for e in existing_extensions_str.split(',') if e.strip()
-        ]
-
-        all_extensions = set(self.extensions)
-        all_extensions.update(existing_extensions)
-
-        headers[HTTP_EXTENSION_HEADER] = ','.join(list(all_extensions))
-        return http_kwargs
-
     async def _apply_interceptors(
         self,
         method_name: str,
@@ -125,11 +106,6 @@ async def _apply_interceptors(
             )
         return final_request_payload, final_http_kwargs
 
-    def _get_http_args(
-        self, context: ClientCallContext | None
-    ) -> dict[str, Any] | None:
-        return context.state.get('http_kwargs') if context else None
-
     async def send_message(
         self,
         request: MessageSendParams,
@@ -141,10 +117,12 @@ async def send_message(
         payload, modified_kwargs = await self._apply_interceptors(
             'message/send',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
-        modified_kwargs = self._update_extension_header(modified_kwargs)
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         response_data = await self._send_request(payload, modified_kwargs)
         response = SendMessageResponse.model_validate(response_data)
         if isinstance(response.root, JSONRPCErrorResponse):
@@ -166,11 +144,13 @@ async def send_message_streaming(
         payload, modified_kwargs = await self._apply_interceptors(
             'message/stream',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
 
-        modified_kwargs = self._update_extension_header(modified_kwargs)
+        modified_kwargs = update_extension_header(
+            modified_kwargs, self.extensions
+        )
         modified_kwargs.setdefault(
             'timeout', self.httpx_client.timeout.as_dict().get('read', None)
         )
@@ -237,7 +217,7 @@ async def get_task(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/get',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
         response_data = await self._send_request(payload, modified_kwargs)
@@ -257,7 +237,7 @@ async def cancel_task(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/cancel',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
         response_data = await self._send_request(payload, modified_kwargs)
@@ -279,7 +259,7 @@ async def set_task_callback(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/pushNotificationConfig/set',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
         response_data = await self._send_request(payload, modified_kwargs)
@@ -303,7 +283,7 @@ async def get_task_callback(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/pushNotificationConfig/get',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
         response_data = await self._send_request(payload, modified_kwargs)
@@ -327,7 +307,7 @@ async def resubscribe(
         payload, modified_kwargs = await self._apply_interceptors(
             'tasks/resubscribe',
             rpc_request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )
 
@@ -369,7 +349,7 @@ async def get_card(
         if not card:
             resolver = A2ACardResolver(self.httpx_client, self.url)
             card = await resolver.get_agent_card(
-                http_kwargs=self._get_http_args(context)
+                http_kwargs=get_http_args(context)
             )
             self._needs_extended_card = (
                 card.supports_authenticated_extended_card
@@ -383,7 +363,7 @@ async def get_card(
         payload, modified_kwargs = await self._apply_interceptors(
             request.method,
             request.model_dump(mode='json', exclude_none=True),
-            self._get_http_args(context),
+            get_http_args(context),
             context,
         )