fix: propagate activated extensions for REST

Author: ishymko

src/a2a/server/routes/rest_dispatcher.py

@@ -6,6 +6,7 @@
 
 from google.protobuf.json_format import MessageToDict, Parse
 
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.routes.common import (
@@ -99,14 +100,29 @@ def _build_call_context(self, request: Request) -> ServerCallContext:
             call_context.tenant = request.path_params['tenant']
         return call_context
 
+    def _extension_headers(self, context: ServerCallContext) -> dict[str, str]:
+        """Builds response headers carrying the activated extensions, if any."""
+        if exts := context.activated_extensions:
+            return {HTTP_EXTENSION_HEADER: ', '.join(sorted(exts))}
+        return {}
+
     async def _handle_non_streaming(
         self,
         request: Request,
         handler_func: Callable[[ServerCallContext], Awaitable[TResponse]],
-    ) -> TResponse:
-        """Centralized error handling and context management for unary calls."""
+        serializer: Callable[[TResponse], Any] = MessageToDict,
+    ) -> JSONResponse:
+        """Centralized error handling and context management for unary calls.
+
+        Builds the call context, invokes the handler, and wraps the result in
+        a `JSONResponse` carrying any activated-extension headers.
+        """
         context = self._build_call_context(request)
-        return await handler_func(context)
+        response = await handler_func(context)
+        return JSONResponse(
+            content=serializer(response),
+            headers=self._extension_headers(context),
+        )
 
     async def _handle_streaming(
         self,
@@ -137,7 +153,9 @@ async def _handle_streaming(
         try:
             first_item = await anext(stream)
         except StopAsyncIteration:
-            return EventSourceResponse(iter([]))
+            return EventSourceResponse(
+                iter([]), headers=self._extension_headers(context)
+            )
 
         async def event_generator() -> AsyncIterator[ServerSentEvent]:
             yield ServerSentEvent(data=json.dumps(first_item))
@@ -151,7 +169,9 @@ async def event_generator() -> AsyncIterator[ServerSentEvent]:
                     event='error',
                 )
 
-        return EventSourceResponse(event_generator())
+        return EventSourceResponse(
+            event_generator(), headers=self._extension_headers(context)
+        )
 
     @rest_error_handler
     async def on_message_send(self, request: Request) -> Response:
@@ -171,8 +191,7 @@ async def _handler(
                 return a2a_pb2.SendMessageResponse(task=task_or_message)
             return a2a_pb2.SendMessageResponse(message=task_or_message)
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_stream_error_handler
     async def on_message_send_stream(
@@ -209,8 +228,7 @@ async def _handler(context: ServerCallContext) -> a2a_pb2.Task:
                 return task
             raise TaskNotFoundError
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_stream_error_handler
     async def on_subscribe_to_task(
@@ -245,8 +263,7 @@ async def _handler(context: ServerCallContext) -> a2a_pb2.Task:
                 return task
             raise TaskNotFoundError
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_error_handler
     async def get_push_notification(self, request: Request) -> Response:
@@ -267,8 +284,7 @@ async def _handler(
                 )
             )
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_error_handler
     async def delete_push_notification(self, request: Request) -> Response:
@@ -285,8 +301,9 @@ async def _handler(context: ServerCallContext) -> None:
                 params, context
             )
 
-        await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content={})
+        return await self._handle_non_streaming(
+            request, _handler, serializer=lambda _: {}
+        )
 
     @rest_error_handler
     async def set_push_notification(self, request: Request) -> Response:
@@ -304,8 +321,7 @@ async def _handler(
                 params, context
             )
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_error_handler
     async def list_push_notifications(self, request: Request) -> Response:
@@ -322,8 +338,7 @@ async def _handler(
                 params, context
             )
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)
 
     @rest_error_handler
     async def list_tasks(self, request: Request) -> Response:
@@ -337,11 +352,12 @@ async def _handler(
             proto_utils.parse_params(request.query_params, params)
             return await self.request_handler.on_list_tasks(params, context)
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(
-            content=MessageToDict(
-                response, always_print_fields_with_no_presence=True
-            )
+        return await self._handle_non_streaming(
+            request,
+            _handler,
+            serializer=lambda r: MessageToDict(
+                r, always_print_fields_with_no_presence=True
+            ),
         )
 
     @rest_error_handler
@@ -359,5 +375,4 @@ async def _handler(
                 params, context
             )
 
-        response = await self._handle_non_streaming(request, _handler)
-        return JSONResponse(content=MessageToDict(response))
+        return await self._handle_non_streaming(request, _handler)

tests/integration/test_end_to_end.py

@@ -15,6 +15,7 @@
     ServiceParametersFactory,
     with_a2a_extensions,
 )
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import EventQueue
 from a2a.server.events.in_memory_queue_manager import InMemoryQueueManager
@@ -835,3 +836,56 @@ async def test_end_to_end_extensions_propagation(transport_setups, streaming):
         response.message, Role.ROLE_AGENT, 'extensions echoed'
     )
     assert set(response.message.extensions) == set(SUPPORTED_EXTENSION_URIS)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'transport_fixture',
+    [
+        pytest.param('rest_setup', id='REST'),
+        pytest.param('jsonrpc_setup', id='JSON-RPC'),
+    ],
+)
+async def test_end_to_end_extensions_response_header(
+    request, transport_fixture
+):
+    """Test that activated extensions are returned in the X-A2A-Extensions
+    response header for HTTP-based transports."""
+    setup = request.getfixturevalue(transport_fixture)
+    client = setup.client
+    client._config.streaming = False
+
+    captured_headers: list[httpx.Headers] = []
+
+    async def capture_response(response: httpx.Response) -> None:
+        captured_headers.append(response.headers)
+
+    client._transport.httpx_client.event_hooks['response'].append(
+        capture_response
+    )
+
+    service_params = ServiceParametersFactory.create(
+        [with_a2a_extensions(SUPPORTED_EXTENSION_URIS)]
+    )
+    context = ClientCallContext(service_parameters=service_params)
+
+    message_to_send = Message(
+        role=Role.ROLE_USER,
+        message_id='msg-ext-header',
+        parts=[Part(text='Extensions: echo')],
+    )
+
+    async for _ in client.send_message(
+        request=SendMessageRequest(message=message_to_send),
+        context=context,
+    ):
+        pass
+
+    assert captured_headers, 'No HTTP response was captured'
+    response_headers = captured_headers[-1]
+    assert HTTP_EXTENSION_HEADER in response_headers
+    returned = {
+        ext.strip()
+        for ext in response_headers[HTTP_EXTENSION_HEADER].split(',')
+    }
+    assert returned == set(SUPPORTED_EXTENSION_URIS)