TMP

Author: ishymko

src/a2a/server/request_handlers/default_request_handler.py

@@ -20,7 +20,6 @@
 )
 from a2a.server.request_handlers.request_handler import (
     RequestHandler,
-    validate_request_params,
 )
 from a2a.server.tasks import (
     PushNotificationConfigStore,
@@ -59,6 +58,7 @@
     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,7 +159,6 @@ async def on_list_tasks(
 
         return page
 
-    @validate_request_params
     async def on_cancel_task(
         self,
         params: CancelTaskRequest,
@@ -169,6 +168,7 @@ 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,7 +323,6 @@ 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,
@@ -334,6 +333,7 @@ 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,7 +393,6 @@ async def push_notification_callback(event: Event) -> None:
 
         return result
 
-    @validate_request_params
     async def on_message_send_stream(
         self,
         params: SendMessageRequest,
@@ -404,6 +403,7 @@ 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,7 +482,6 @@ 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,
@@ -492,6 +491,7 @@ 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,7 +508,6 @@ async def on_create_task_push_notification_config(
 
         return params
 
-    @validate_request_params
     async def on_get_task_push_notification_config(
         self,
         params: GetTaskPushNotificationConfigRequest,
@@ -518,6 +517,7 @@ 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,7 +540,6 @@ 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,
@@ -551,6 +550,7 @@ 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,7 +583,6 @@ 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,
@@ -593,6 +592,7 @@ 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,7 +609,6 @@ 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,
@@ -619,6 +618,7 @@ 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,7 +227,7 @@ 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.iscoroutinefunction(method):