test: add extension propagation test in test_end_to_end.py

Author: ishymko

tests/integration/test_end_to_end.py

@@ -5,15 +5,21 @@
 import httpx
 import pytest
 import pytest_asyncio
+
+from google.protobuf.struct_pb2 import Struct, Value
 from starlette.applications import Starlette
 
 from a2a.client.base_client import BaseClient
-from a2a.client.client import ClientConfig
+from a2a.client.client import ClientCallContext, ClientConfig
 from a2a.client.client_factory import ClientFactory
+from a2a.client.service_parameters import (
+    ServiceParametersFactory,
+    with_a2a_extensions,
+)
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import EventQueue
 from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager
-from a2a.server.request_handlers import GrpcHandler, DefaultRequestHandler
+from a2a.server.request_handlers import DefaultRequestHandler, GrpcHandler
 from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
 from a2a.server.routes.rest_routes import create_rest_routes
 from a2a.server.tasks import TaskUpdater
@@ -87,6 +93,29 @@ class MockAgentExecutor(AgentExecutor):
     async def execute(self, context: RequestContext, event_queue: EventQueue):
         user_input = context.get_user_input()
 
+        # Extensions echo: report requested extensions and activate them.
+        if user_input.startswith('Extensions:'):
+            requested = sorted(context.requested_extensions)
+            for ext_uri in requested:
+                context.add_activated_extension(ext_uri)
+            activated = sorted(context.call_context.activated_extensions)
+
+            payload = Struct()
+            payload.update(
+                {
+                    'requested_extensions': requested,
+                    'activated_extensions': activated,
+                }
+            )
+            await event_queue.enqueue_event(
+                Message(
+                    role=Role.ROLE_AGENT,
+                    message_id='ext-reply-1',
+                    parts=[Part(data=Value(struct_value=payload))],
+                )
+            )
+            return
+
         # Direct message response (no task created).
         if user_input.startswith('Message:'):
             await event_queue.enqueue_event(
@@ -757,3 +786,53 @@ async def test_end_to_end_direct_message_return_immediately(transport_setups):
         Role.ROLE_AGENT,
         'Direct reply to: Message: Quick question',
     )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'streaming',
+    [
+        pytest.param(False, id='blocking'),
+        pytest.param(True, id='streaming'),
+    ],
+)
+async def test_end_to_end_extensions_propagation(transport_setups, streaming):
+    """Test that extensions sent by the client reach the agent executor."""
+    client = transport_setups.client
+    client._config.streaming = streaming
+
+    extensions = [
+        'https://example.com/ext/v1',
+        'https://example.com/ext/v2',
+    ]
+    service_params = ServiceParametersFactory.create(
+        [with_a2a_extensions(extensions)]
+    )
+    context = ClientCallContext(service_parameters=service_params)
+
+    message_to_send = Message(
+        role=Role.ROLE_USER,
+        message_id='msg-ext-propagation',
+        parts=[Part(text='Extensions: echo')],
+    )
+
+    events = [
+        event
+        async for event in client.send_message(
+            request=SendMessageRequest(message=message_to_send),
+            context=context,
+        )
+    ]
+
+    assert len(events) == 1
+    response = events[0]
+    assert response.HasField('message')
+    part = response.message.parts[0]
+    assert part.HasField('data')
+
+    payload = part.data.struct_value
+    requested = set(payload['requested_extensions'])
+    activated = set(payload['activated_extensions'])
+
+    assert requested == set(extensions)
+    assert activated == set(extensions)