fix

Author: guglielmo-san

src/a2a/server/routes/rest_dispatcher.py

@@ -1,7 +1,8 @@
 import json
 import logging
+
 from collections.abc import AsyncIterator, Awaitable, Callable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from google.protobuf.json_format import MessageToDict, Parse
 
@@ -52,6 +53,9 @@
 
 logger = logging.getLogger(__name__)
 
+TResponse = TypeVar('TResponse')
+
+
 @trace_class(kind=SpanKind.SERVER)
 class RestDispatcher:
     """Dispatches incoming REST requests to the appropriate handler methods.
@@ -108,29 +112,24 @@ def _build_call_context(self, request: Request) -> ServerCallContext:
             call_context.tenant = request.path_params['tenant']
         return call_context
 
-    @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
-    async def on_message_send(self, request: Request) -> Response:
-        """Handles the 'message/send' REST method."""
+    async def _handle_non_streaming(
+        self,
+        request: Request,
+        handler_func: Callable[[ServerCallContext], Awaitable[TResponse]],
+    ) -> TResponse:
+        """Centralized error handling and context management for unary calls."""
         context = self._build_call_context(request)
-        body = await request.body()
-        params = a2a_pb2.SendMessageRequest()
-        Parse(body, params)
-        task_or_message = await self.request_handler.on_message_send(params, context)
-        if isinstance(task_or_message, a2a_pb2.Task):
-            response = a2a_pb2.SendMessageResponse(task=task_or_message)
-        else:
-            response = a2a_pb2.SendMessageResponse(message=task_or_message)
-        return JSONResponse(content=MessageToDict(response))
+        return await handler_func(context)
 
-    @rest_stream_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate(
-        lambda self: self.agent_card.capabilities.streaming,
-        'Streaming is not supported by the agent',
-    )
-    async def on_message_send_stream(self, request: Request) -> EventSourceResponse:
-        """Handles the 'message/stream' REST method."""
+    async def _handle_streaming(
+        self,
+        request: Request,
+        handler_func: Callable[[ServerCallContext], AsyncIterator[Any]],
+    ) -> EventSourceResponse:
+        """Centralized error handling and context management for streaming calls."""
+        # Pre-consume and cache the request body to prevent deadlock in streaming context
+        # This is required because Starlette's request.body() can only be consumed once,
+        # and attempting to consume it after EventSourceResponse starts causes deadlock
         try:
             await request.body()
         except (ValueError, RuntimeError, OSError) as e:
@@ -139,139 +138,234 @@ async def on_message_send_stream(self, request: Request) -> EventSourceResponse:
             ) from e
 
         context = self._build_call_context(request)
-        body = await request.body()
-        params = a2a_pb2.SendMessageRequest()
-        Parse(body, params)
 
-        stream = aiter(self.request_handler.on_message_send_stream(params, context))
+        # Eagerly fetch the first item from the stream so that errors raised
+        # before any event is yielded (e.g. validation, parsing, or handler
+        # failures) propagate here and are caught by
+        # @rest_stream_error_handler, which returns a JSONResponse with
+        # the correct HTTP status code instead of starting an SSE stream.
+        # Without this, the error would be raised after SSE headers are
+        # already sent, and the client would see a broken stream instead
+        stream = aiter(handler_func(context))
         try:
-            first_event = await anext(stream)
+            first_item = await anext(stream)
         except StopAsyncIteration:
             return EventSourceResponse(iter([]))
 
         async def event_generator() -> AsyncIterator[str]:
-            yield json.dumps(MessageToDict(proto_utils.to_stream_response(first_event)))
-            async for event in stream:
-                yield json.dumps(MessageToDict(proto_utils.to_stream_response(event)))
+            yield json.dumps(first_item)
+            async for item in stream:
+                yield json.dumps(item)
 
         return EventSourceResponse(event_generator())
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
+    async def on_message_send(self, request: Request) -> Response:
+        """Handles the 'message/send' REST method."""
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(
+            context: ServerCallContext,
+        ) -> a2a_pb2.SendMessageResponse:
+            body = await request.body()
+            params = a2a_pb2.SendMessageRequest()
+            Parse(body, params)
+            task_or_message = await self.request_handler.on_message_send(
+                params, context
+            )
+            if isinstance(task_or_message, a2a_pb2.Task):
+                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))
+
+    @rest_stream_error_handler
+    async def on_message_send_stream(
+        self, request: Request
+    ) -> EventSourceResponse:
+        """Handles the 'message/stream' REST method."""
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        @validate(
+            lambda _: self.agent_card.capabilities.streaming,
+            'Streaming is not supported by the agent',
+        )
+        async def _handler(
+            context: ServerCallContext,
+        ) -> AsyncIterator[dict[str, Any]]:
+            body = await request.body()
+            params = a2a_pb2.SendMessageRequest()
+            Parse(body, params)
+            async for event in self.request_handler.on_message_send_stream(
+                params, context
+            ):
+                response = proto_utils.to_stream_response(event)
+                yield MessageToDict(response)
+
+        return await self._handle_streaming(request, _handler)
+
+    @rest_error_handler
     async def on_cancel_task(self, request: Request) -> Response:
         """Handles the 'tasks/cancel' REST method."""
-        context = self._build_call_context(request)
-        task_id = request.path_params['id']
-        task = await self.request_handler.on_cancel_task(CancelTaskRequest(id=task_id), context)
-        if task:
-            return JSONResponse(content=MessageToDict(task))
-        raise TaskNotFoundError
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(context: ServerCallContext) -> a2a_pb2.Task:
+            task_id = request.path_params['id']
+            task = await self.request_handler.on_cancel_task(
+                CancelTaskRequest(id=task_id), context
+            )
+            if task:
+                return task
+            raise TaskNotFoundError
+
+        response = await self._handle_non_streaming(request, _handler)
+        return JSONResponse(content=MessageToDict(response))
 
     @rest_stream_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate(
-        lambda self: self.agent_card.capabilities.streaming,
-        'Streaming is not supported by the agent',
-    )
-    async def on_subscribe_to_task(self, request: Request) -> EventSourceResponse:
+    async def on_subscribe_to_task(
+        self, request: Request
+    ) -> EventSourceResponse:
         """Handles the 'SubscribeToTask' REST method."""
-        try:
-            await request.body()
-        except (ValueError, RuntimeError, OSError) as e:
-            raise InvalidRequestError(
-                message=f'Failed to pre-consume request body: {e}'
-            ) from e
-
-        context = self._build_call_context(request)
         task_id = request.path_params['id']
-        
-        stream = aiter(self.request_handler.on_subscribe_to_task(SubscribeToTaskRequest(id=task_id), context))
-        try:
-            first_event = await anext(stream)
-        except StopAsyncIteration:
-            return EventSourceResponse(iter([]))
 
-        async def event_generator() -> AsyncIterator[str]:
-            yield json.dumps(MessageToDict(proto_utils.to_stream_response(first_event)))
-            async for event in stream:
-                yield json.dumps(MessageToDict(proto_utils.to_stream_response(event)))
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        @validate(
+            lambda _: self.agent_card.capabilities.streaming,
+            'Streaming is not supported by the agent',
+        )
+        async def _handler(
+            context: ServerCallContext,
+        ) -> AsyncIterator[dict[str, Any]]:
+            async for event in self.request_handler.on_subscribe_to_task(
+                SubscribeToTaskRequest(id=task_id), context
+            ):
+                response = proto_utils.to_stream_response(event)
+                yield MessageToDict(response)
 
-        return EventSourceResponse(event_generator())
+        return await self._handle_streaming(request, _handler)
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def on_get_task(self, request: Request) -> Response:
         """Handles the 'tasks/{id}' REST method."""
-        context = self._build_call_context(request)
-        params = a2a_pb2.GetTaskRequest()
-        proto_utils.parse_params(request.query_params, params)
-        params.id = request.path_params['id']
-        task = await self.request_handler.on_get_task(params, context)
-        if task:
-            return JSONResponse(content=MessageToDict(task))
-        raise TaskNotFoundError
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(context: ServerCallContext) -> a2a_pb2.Task:
+            params = a2a_pb2.GetTaskRequest()
+            proto_utils.parse_params(request.query_params, params)
+            params.id = request.path_params['id']
+            task = await self.request_handler.on_get_task(params, context)
+            if task:
+                return task
+            raise TaskNotFoundError
+
+        response = await self._handle_non_streaming(request, _handler)
+        return JSONResponse(content=MessageToDict(response))
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def get_push_notification(self, request: Request) -> Response:
         """Handles the 'tasks/pushNotificationConfig/get' REST method."""
-        context = self._build_call_context(request)
-        task_id = request.path_params['id']
-        push_id = request.path_params['push_id']
-        params = GetTaskPushNotificationConfigRequest(task_id=task_id, id=push_id)
-        config = await self.request_handler.on_get_task_push_notification_config(params, context)
-        return JSONResponse(content=MessageToDict(config))
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(
+            context: ServerCallContext,
+        ) -> a2a_pb2.TaskPushNotificationConfig:
+            task_id = request.path_params['id']
+            push_id = request.path_params['push_id']
+            params = GetTaskPushNotificationConfigRequest(
+                task_id=task_id, id=push_id
+            )
+            return (
+                await self.request_handler.on_get_task_push_notification_config(
+                    params, context
+                )
+            )
+
+        response = await self._handle_non_streaming(request, _handler)
+        return JSONResponse(content=MessageToDict(response))
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def delete_push_notification(self, request: Request) -> Response:
         """Handles the 'tasks/pushNotificationConfig/delete' REST method."""
-        context = self._build_call_context(request)
-        task_id = request.path_params['id']
-        push_id = request.path_params['push_id']
-        params = a2a_pb2.DeleteTaskPushNotificationConfigRequest(task_id=task_id, id=push_id)
-        await self.request_handler.on_delete_task_push_notification_config(params, context)
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(context: ServerCallContext) -> None:
+            task_id = request.path_params['id']
+            push_id = request.path_params['push_id']
+            params = a2a_pb2.DeleteTaskPushNotificationConfigRequest(
+                task_id=task_id, id=push_id
+            )
+            await self.request_handler.on_delete_task_push_notification_config(
+                params, context
+            )
+
+        await self._handle_non_streaming(request, _handler)
         return JSONResponse(content={})
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate(
-        lambda self: self.agent_card.capabilities.push_notifications,
-        'Push notifications are not supported by the agent',
-    )
     async def set_push_notification(self, request: Request) -> Response:
         """Handles the 'tasks/pushNotificationConfig/set' REST method."""
-        context = self._build_call_context(request)
-        body = await request.body()
-        params = a2a_pb2.TaskPushNotificationConfig()
-        Parse(body, params)
-        params.task_id = request.path_params['id']
-        config = await self.request_handler.on_create_task_push_notification_config(params, context)
-        return JSONResponse(content=MessageToDict(config))
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        @validate(
+            lambda _: self.agent_card.capabilities.push_notifications,
+            'Push notifications are not supported by the agent',
+        )
+        async def _handler(
+            context: ServerCallContext,
+        ) -> a2a_pb2.TaskPushNotificationConfig:
+            body = await request.body()
+            params = a2a_pb2.TaskPushNotificationConfig()
+            Parse(body, params)
+            params.task_id = request.path_params['id']
+            return await self.request_handler.on_create_task_push_notification_config(
+                params, context
+            )
+
+        response = await self._handle_non_streaming(request, _handler)
+        return JSONResponse(content=MessageToDict(response))
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def list_push_notifications(self, request: Request) -> Response:
         """Handles the 'tasks/pushNotificationConfig/list' REST method."""
-        context = self._build_call_context(request)
-        params = a2a_pb2.ListTaskPushNotificationConfigsRequest()
-        proto_utils.parse_params(request.query_params, params)
-        params.task_id = request.path_params['id']
-        result = await self.request_handler.on_list_task_push_notification_configs(params, context)
-        return JSONResponse(content=MessageToDict(result))
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(
+            context: ServerCallContext,
+        ) -> a2a_pb2.ListTaskPushNotificationConfigsResponse:
+            params = a2a_pb2.ListTaskPushNotificationConfigsRequest()
+            proto_utils.parse_params(request.query_params, params)
+            params.task_id = request.path_params['id']
+            return await self.request_handler.on_list_task_push_notification_configs(
+                params, context
+            )
+
+        response = await self._handle_non_streaming(request, _handler)
+        return JSONResponse(content=MessageToDict(response))
 
     @rest_error_handler
-    @validate_version(constants.PROTOCOL_VERSION_1_0)
     async def list_tasks(self, request: Request) -> Response:
         """Handles the 'tasks/list' REST method."""
-        context = self._build_call_context(request)
-        params = a2a_pb2.ListTasksRequest()
-        proto_utils.parse_params(request.query_params, params)
-        result = await self.request_handler.on_list_tasks(params, context)
-        return JSONResponse(content=MessageToDict(result, always_print_fields_with_no_presence=True))
+
+        @validate_version(constants.PROTOCOL_VERSION_1_0)
+        async def _handler(
+            context: ServerCallContext,
+        ) -> a2a_pb2.ListTasksResponse:
+            params = a2a_pb2.ListTasksRequest()
+            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
+            )
+        )
 
     @rest_error_handler
-    async def handle_authenticated_agent_card(self, request: Request) -> Response:
+    async def handle_authenticated_agent_card(
+        self, request: Request
+    ) -> Response:
         """Handles the 'extendedAgentCard' REST method."""
         if not self.agent_card.capabilities.extended_agent_card:
             raise ExtendedAgentCardNotConfiguredError(
@@ -288,5 +382,7 @@ async def handle_authenticated_agent_card(self, request: Request) -> Response:
             card_to_serve = await maybe_await(self.card_modifier(card_to_serve))
 
         return JSONResponse(
-            content=MessageToDict(card_to_serve, preserving_proto_field_name=True)
+            content=MessageToDict(
+                card_to_serve, preserving_proto_field_name=True
+            )
         )