From 0e36dc263bb3cb5d4c23d89c9cf923f28c9cd9f0 Mon Sep 17 00:00:00 2001
From: Ivan Shymko <ishymko@google.com>
Date: Wed, 18 Feb 2026 13:18:57 +0000
Subject: [PATCH] refactor: unify on_message_send in server

JSON-RPC diverged a bit, update it to match other transports. Domain request handler returns `Task | Message` from `on_message_send` already.
---
 src/a2a/server/request_handlers/grpc_handler.py    |  1 -
 src/a2a/server/request_handlers/jsonrpc_handler.py | 10 ++--------
 src/a2a/server/request_handlers/rest_handler.py    |  1 -
 3 files changed, 2 insertions(+), 10 deletions(-)

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: