WIP

Author: ishymko

src/a2a/server/request_handlers/__init__.py

@@ -6,7 +6,10 @@
     DefaultRequestHandler,
 )
 from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
-from a2a.server.request_handlers.request_handler import RequestHandler
+from a2a.server.request_handlers.request_handler import (
+    RequestHandler,
+    validate_request_params,
+)
 from a2a.server.request_handlers.response_helpers import (
     build_error_response,
     prepare_response_object,
@@ -45,4 +48,5 @@ def __init__(self, *args, **kwargs):
     'RequestHandler',
     'build_error_response',
     'prepare_response_object',
+    'validate_request_params',
 ]

src/a2a/server/request_handlers/default_request_handler.py

@@ -20,6 +20,7 @@
 )
 from a2a.server.request_handlers.request_handler import (
     RequestHandler,
+    validate_request_params,
 )
 from a2a.server.tasks import (
     PushNotificationConfigStore,
@@ -58,7 +59,6 @@
     validate_page_size,
 )
 from a2a.utils.telemetry import SpanKind, trace_class
-from a2a.utils.proto_utils import validate_proto_required_fields
 
 
 logger = logging.getLogger(__name__)
@@ -121,13 +121,13 @@ def __init__(  # noqa: PLR0913
         # asyncio tasks and to surface unexpected exceptions.
         self._background_tasks = set()
 
+    @validate_request_params
     async def on_get_task(
         self,
         params: GetTaskRequest,
         context: ServerCallContext,
     ) -> Task | None:
         """Default handler for 'tasks/get'."""
-        validate_proto_required_fields(params)
         validate_history_length(params)
 
         task_id = params.id
@@ -137,13 +137,13 @@ async def on_get_task(
 
         return apply_history_length(task, params)
 
+    @validate_request_params
     async def on_list_tasks(
         self,
         params: ListTasksRequest,
         context: ServerCallContext,
     ) -> ListTasksResponse:
         """Default handler for 'tasks/list'."""
-        validate_proto_required_fields(params)
         validate_history_length(params)
         if params.HasField('page_size'):
             validate_page_size(params.page_size)
@@ -159,6 +159,7 @@ async def on_list_tasks(
 
         return page
 
+    @validate_request_params
     async def on_cancel_task(
         self,
         params: CancelTaskRequest,
@@ -168,7 +169,6 @@ async def on_cancel_task(
 
         Attempts to cancel the task managed by the `AgentExecutor`.
         """
-        validate_proto_required_fields(params)
         task_id = params.id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
@@ -323,6 +323,7 @@ async def _send_push_notification_if_needed(
         ):
             await self._push_sender.send_notification(task_id, event)
 
+    @validate_request_params
     async def on_message_send(
         self,
         params: SendMessageRequest,
@@ -333,7 +334,6 @@ async def on_message_send(
         Starts the agent execution for the message and waits for the final
         result (Task or Message).
         """
-        validate_proto_required_fields(params)
         validate_history_length(params.configuration)
 
         (
@@ -393,6 +393,7 @@ async def push_notification_callback(event: Event) -> None:
 
         return result
 
+    @validate_request_params
     async def on_message_send_stream(
         self,
         params: SendMessageRequest,
@@ -403,7 +404,6 @@ async def on_message_send_stream(
         Starts the agent execution and yields events as they are produced
         by the agent.
         """
-        validate_proto_required_fields(params)
         (
             _task_manager,
             task_id,
@@ -482,6 +482,7 @@ async def _cleanup_producer(
         async with self._running_agents_lock:
             self._running_agents.pop(task_id, None)
 
+    @validate_request_params
     async def on_create_task_push_notification_config(
         self,
         params: TaskPushNotificationConfig,
@@ -491,7 +492,6 @@ async def on_create_task_push_notification_config(
 
         Requires a `PushNotifier` to be configured.
         """
-        validate_proto_required_fields(params)
         if not self._push_config_store:
             raise UnsupportedOperationError
 
@@ -508,6 +508,7 @@ async def on_create_task_push_notification_config(
 
         return params
 
+    @validate_request_params
     async def on_get_task_push_notification_config(
         self,
         params: GetTaskPushNotificationConfigRequest,
@@ -517,7 +518,6 @@ async def on_get_task_push_notification_config(
 
         Requires a `PushConfigStore` to be configured.
         """
-        validate_proto_required_fields(params)
         if not self._push_config_store:
             raise UnsupportedOperationError
 
@@ -540,6 +540,7 @@ async def on_get_task_push_notification_config(
 
         raise InternalError(message='Push notification config not found')
 
+    @validate_request_params
     async def on_subscribe_to_task(
         self,
         params: SubscribeToTaskRequest,
@@ -550,7 +551,6 @@ async def on_subscribe_to_task(
         Allows a client to re-attach to a running streaming task's event stream.
         Requires the task and its queue to still be active.
         """
-        validate_proto_required_fields(params)
         task_id = params.id
         task: Task | None = await self.task_store.get(task_id, context)
         if not task:
@@ -583,6 +583,7 @@ async def on_subscribe_to_task(
         async for event in result_aggregator.consume_and_emit(consumer):
             yield event
 
+    @validate_request_params
     async def on_list_task_push_notification_configs(
         self,
         params: ListTaskPushNotificationConfigsRequest,
@@ -592,7 +593,6 @@ async def on_list_task_push_notification_configs(
 
         Requires a `PushConfigStore` to be configured.
         """
-        validate_proto_required_fields(params)
         if not self._push_config_store:
             raise UnsupportedOperationError
 
@@ -609,6 +609,7 @@ async def on_list_task_push_notification_configs(
             configs=push_notification_config_list
         )
 
+    @validate_request_params
     async def on_delete_task_push_notification_config(
         self,
         params: DeleteTaskPushNotificationConfigRequest,
@@ -618,7 +619,6 @@ async def on_delete_task_push_notification_config(
 
         Requires a `PushConfigStore` to be configured.
         """
-        validate_proto_required_fields(params)
         if not self._push_config_store:
             raise UnsupportedOperationError
 

src/a2a/server/request_handlers/request_handler.py

@@ -227,8 +227,25 @@ async def on_delete_task_push_notification_config(
         """
 
 
-def _validate_request_params(method: Callable) -> Callable:
+def validate_request_params(method: Callable) -> Callable:
     """Decorator for RequestHandler methods to validate required fields on incoming requests."""
+    if inspect.isasyncgenfunction(method):
+
+        @functools.wraps(method)
+        async def async_gen_wrapper(
+            self: RequestHandler,
+            params: ProtoMessage,
+            context: ServerCallContext,
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
+            if params is not None:
+                validate_proto_required_fields(params)
+            async for item in method(self, params, context, *args, **kwargs):
+                yield item
+
+        return async_gen_wrapper
+
     if inspect.iscoroutinefunction(method):
 
         @functools.wraps(method)