refactor: unify on_message_send in server

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)

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)

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: