refactor: streamline extension handling in BaseClient and GrpcTransport. Add helper __merge_extensions method in utils.py

Author: sokoliva

src/a2a/client/base_client.py

@@ -42,8 +42,7 @@ def __init__(
         self._card = card
         self._config = config
         self._transport = transport
-        if self._extensions:
-            self._config.extensions = self._extensions
+        self._config.extensions = self._extensions
 
     async def send_message(
         self,

src/a2a/client/transports/grpc.py

@@ -13,11 +13,13 @@
         "'pip install a2a-sdk[grpc]'"
     ) from e
 
+from google.protobuf import struct_pb2
+
 from a2a.client.client import ClientConfig
 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.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
     AgentCard,
@@ -59,24 +61,6 @@ def __init__(
         )
         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(
         cls,
@@ -105,7 +89,7 @@ async def send_message(
                 ),
                 metadata=proto_utils.ToProto.metadata(request.metadata),
             ),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
@@ -128,14 +112,15 @@ async def send_message_streaming(
                 ),
                 metadata=proto_utils.ToProto.metadata(request.metadata),
             ),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(request.metadata),
         )
         while True:
             response = await stream.read()
             if response == grpc.aio.EOF:  # pyright: ignore[reportAttributeAccessIssue]
                 break
             yield proto_utils.FromProto.stream_response(response)
 
+    # iva todo TaskIdParams has metadata
     async def resubscribe(
         self, request: TaskIdParams, *, context: ClientCallContext | None = None
     ) -> AsyncGenerator[
@@ -144,14 +129,15 @@ async def resubscribe(
         """Reconnects to get task updates."""
         stream = self.stub.TaskSubscription(
             a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}'),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         while True:
             response = await stream.read()
             if response == grpc.aio.EOF:  # pyright: ignore[reportAttributeAccessIssue]
                 break
             yield proto_utils.FromProto.stream_response(response)
 
+    # iva todo TaskQueryParams has metadata
     async def get_task(
         self,
         request: TaskQueryParams,
@@ -164,7 +150,7 @@ async def get_task(
                 name=f'tasks/{request.id}',
                 history_length=request.history_length,
             ),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         return proto_utils.FromProto.task(task)
 
@@ -177,7 +163,7 @@ async def cancel_task(
         """Requests the agent to cancel a specific task."""
         task = await self.stub.CancelTask(
             a2a_pb2.CancelTaskRequest(name=f'tasks/{request.id}'),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         return proto_utils.FromProto.task(task)
 
@@ -196,10 +182,11 @@ async def set_task_callback(
                     request
                 ),
             ),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
+    # iva todo GetTaskPushNotificationConfigParams has metadata
     async def get_task_callback(
         self,
         request: GetTaskPushNotificationConfigParams,
@@ -211,7 +198,7 @@ async def get_task_callback(
             a2a_pb2.GetTaskPushNotificationConfigRequest(
                 name=f'tasks/{request.id}/pushNotificationConfigs/{request.push_notification_config_id}',
             ),
-            metadata=self._get_metadata(context),
+            metadata=self._update_extension_metadata(),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
@@ -235,6 +222,33 @@ async def get_card(
         self._needs_extended_card = False
         return card
 
+    def _update_extension_metadata(
+        self, metadata: dict[str, Any] | None = None
+    ) -> struct_pb2.Struct | None:
+        """Gets the metadata for the gRPC call."""
+        if metadata is None:
+            metadata = {}
+
+        if self.extensions:
+            existing_extensions_str = str(
+                metadata.get(HTTP_EXTENSION_HEADER, '')
+            )
+            existing_extensions = {
+                e.strip()
+                for e in existing_extensions_str.split(',')
+                if e.strip()
+            }
+
+            all_extensions = set(existing_extensions)
+            all_extensions.update(self.extensions)
+
+            if all_extensions:
+                metadata[HTTP_EXTENSION_HEADER] = ','.join(all_extensions)
+            elif HTTP_EXTENSION_HEADER in metadata:
+                del metadata[HTTP_EXTENSION_HEADER]
+
+        return proto_utils.ToProto.metadata(metadata if metadata else None)
+
     async def close(self) -> None:
         """Closes the gRPC channel."""
         await self.channel.close()

src/a2a/client/transports/utils.py

@@ -8,17 +8,33 @@ def get_http_args(context: ClientCallContext | None) -> dict[str, Any] | None:
     return context.state.get('http_kwargs') if context else None
 
 
+def __merge_extensions(
+    existing_extensions: str, new_extensions: list[str]
+) -> str:
+    existing_extensions_list = [
+        e.strip() for e in existing_extensions.split(',') if e.strip()
+    ]
+    existing_extensions_set = set(existing_extensions_list)
+    new_extensions = [
+        ext for ext in new_extensions if ext not in existing_extensions_set
+    ]
+
+    return ','.join(existing_extensions_list + new_extensions)
+
+
 def update_extension_header(
     http_kwargs: dict[str, Any], extensions: list[str] | None
 ) -> dict[str, Any]:
     if not extensions:
         return http_kwargs
     headers = http_kwargs.setdefault('headers', {})
     existing_extensions_str = headers.get(HTTP_EXTENSION_HEADER, '')
-    existing_extensions = [
+    """existing_extensions = [
         e.strip() for e in existing_extensions_str.split(',') if e.strip()
     ]
     all_extensions = set(existing_extensions)
-    all_extensions.update(extensions)
-    headers[HTTP_EXTENSION_HEADER] = ','.join(all_extensions)
+    all_extensions.update(extensions)"""
+    headers[HTTP_EXTENSION_HEADER] = __merge_extensions(
+        existing_extensions_str, extensions
+    )
     return http_kwargs