Move transport tests from tests/client to tests/client/transport. Add method "update_extension_metadata" to tansports/utils.py.

Author: sokoliva

src/a2a/client/base_client.py

@@ -40,9 +40,9 @@ def __init__(
     ):
         super().__init__(consumers, middleware, extensions)
         self._card = card
+        config.extensions = extensions
         self._config = config
         self._transport = transport
-        self._config.extensions = self._extensions
 
     async def send_message(
         self,

src/a2a/client/client.py

@@ -93,7 +93,6 @@ def __init__(
         self,
         consumers: list[Consumer] | None = None,
         middleware: list[ClientCallInterceptor] | None = None,
-        # 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.
@@ -118,8 +117,6 @@ 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,7 +1,6 @@
 import logging
 
 from collections.abc import AsyncGenerator
-from typing import Any
 
 
 try:
@@ -13,13 +12,12 @@
         "'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.extensions.common import HTTP_EXTENSION_HEADER
+from a2a.client.transports.utils import update_extension_metadata
 from a2a.grpc import a2a_pb2, a2a_pb2_grpc
 from a2a.types import (
     AgentCard,
@@ -87,9 +85,10 @@ async def send_message(
                 configuration=proto_utils.ToProto.message_send_configuration(
                     request.configuration
                 ),
-                metadata=proto_utils.ToProto.metadata(request.metadata),
+                metadata=update_extension_metadata(
+                    request.metadata, self.extensions
+                ),
             ),
-            metadata=self._update_extension_metadata(),
         )
         if response.HasField('task'):
             return proto_utils.FromProto.task(response.task)
@@ -110,17 +109,17 @@ async def send_message_streaming(
                 configuration=proto_utils.ToProto.message_send_configuration(
                     request.configuration
                 ),
-                metadata=proto_utils.ToProto.metadata(request.metadata),
+                metadata=update_extension_metadata(
+                    request.metadata, self.extensions
+                ),
             ),
-            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[
@@ -129,15 +128,16 @@ async def resubscribe(
         """Reconnects to get task updates."""
         stream = self.stub.TaskSubscription(
             a2a_pb2.TaskSubscriptionRequest(name=f'tasks/{request.id}'),
-            metadata=self._update_extension_metadata(),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         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,
@@ -150,7 +150,9 @@ async def get_task(
                 name=f'tasks/{request.id}',
                 history_length=request.history_length,
             ),
-            metadata=self._update_extension_metadata(),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task(task)
 
@@ -163,7 +165,6 @@ 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._update_extension_metadata(),
         )
         return proto_utils.FromProto.task(task)
 
@@ -182,11 +183,12 @@ async def set_task_callback(
                     request
                 ),
             ),
-            metadata=self._update_extension_metadata(),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
-    # iva todo GetTaskPushNotificationConfigParams has metadata
     async def get_task_callback(
         self,
         request: GetTaskPushNotificationConfigParams,
@@ -198,7 +200,9 @@ async def get_task_callback(
             a2a_pb2.GetTaskPushNotificationConfigRequest(
                 name=f'tasks/{request.id}/pushNotificationConfigs/{request.push_notification_config_id}',
             ),
-            metadata=self._update_extension_metadata(),
+            metadata=update_extension_metadata(
+                request.metadata, self.extensions
+            ),
         )
         return proto_utils.FromProto.task_push_notification_config(config)
 
@@ -222,33 +226,6 @@ 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

@@ -1,7 +1,10 @@
 from typing import Any
 
+from google.protobuf import struct_pb2
+
 from a2a.client.middleware import ClientCallContext
 from a2a.extensions.common import HTTP_EXTENSION_HEADER
+from a2a.utils import proto_utils
 
 
 def get_http_args(context: ClientCallContext | None) -> dict[str, Any] | None:
@@ -14,27 +17,33 @@ def __merge_extensions(
     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
+        ext for ext in new_extensions if ext not in existing_extensions_list
     ]
-
     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 = [
-        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] = __merge_extensions(
-        existing_extensions_str, extensions
-    )
+    if extensions:
+        headers = http_kwargs.setdefault('headers', {})
+        existing_extensions_str = headers.get(HTTP_EXTENSION_HEADER, '')
+
+        headers[HTTP_EXTENSION_HEADER] = __merge_extensions(
+            existing_extensions_str, extensions
+        )
     return http_kwargs
+
+
+def update_extension_metadata(
+    metadata: dict[str, Any] | None, extensions: list[str] | None
+) -> struct_pb2.Struct | None:
+    if metadata is None:
+        metadata = {}
+    if extensions:
+        existing_extensions_str = str(metadata.get(HTTP_EXTENSION_HEADER, ''))
+        metadata[HTTP_EXTENSION_HEADER] = __merge_extensions(
+            existing_extensions_str, extensions
+        )
+    return proto_utils.ToProto.metadata(metadata if metadata else None)