diff --git a/src/a2a/server/request_handlers/grpc_handler.py b/src/a2a/server/request_handlers/grpc_handler.py index aab011357..4735ebc53 100644 --- a/src/a2a/server/request_handlers/grpc_handler.py +++ b/src/a2a/server/request_handlers/grpc_handler.py @@ -132,7 +132,6 @@ async def SendMessage( request, server_context ) self._set_extension_metadata(context, server_context) - # Wrap in SendMessageResponse based on type if isinstance(task_or_message, a2a_pb2.Task): return a2a_pb2.SendMessageResponse(task=task_or_message) return a2a_pb2.SendMessageResponse(message=task_or_message) diff --git a/src/a2a/server/request_handlers/jsonrpc_handler.py b/src/a2a/server/request_handlers/jsonrpc_handler.py index de168719f..06ca917f2 100644 --- a/src/a2a/server/request_handlers/jsonrpc_handler.py +++ b/src/a2a/server/request_handlers/jsonrpc_handler.py @@ -26,7 +26,6 @@ GetTaskRequest, ListTaskPushNotificationConfigRequest, ListTasksRequest, - Message, SendMessageRequest, SendMessageResponse, SubscribeToTaskRequest, @@ -171,15 +170,10 @@ async def on_message_send( task_or_message = await self.request_handler.on_message_send( request, context ) - # Build result based on return type - response = SendMessageResponse() if isinstance(task_or_message, Task): - response.task.CopyFrom(task_or_message) - elif isinstance(task_or_message, Message): - response.message.CopyFrom(task_or_message) + response = SendMessageResponse(task=task_or_message) else: - # Should we handle this fallthrough? - pass + response = SendMessageResponse(message=task_or_message) result = MessageToDict(response) return _build_success_response(request_id, result) diff --git a/src/a2a/server/request_handlers/rest_handler.py b/src/a2a/server/request_handlers/rest_handler.py index afa362147..61e063570 100644 --- a/src/a2a/server/request_handlers/rest_handler.py +++ b/src/a2a/server/request_handlers/rest_handler.py @@ -84,7 +84,6 @@ async def on_message_send( task_or_message = await self.request_handler.on_message_send( params, context ) - # Wrap the result in a SendMessageResponse if isinstance(task_or_message, a2a_pb2.Task): response = a2a_pb2.SendMessageResponse(task=task_or_message) else: