refactor: cleanup @Validate decorators. (#853)

Use validate_async_generator and add validators to compat grpc handler.

Author: bartek-w

src/a2a/compat/v0_3/grpc_handler.py

@@ -29,7 +29,7 @@
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.types.a2a_pb2 import AgentCard
 from a2a.utils.errors import A2AError, InvalidParamsError
-from a2a.utils.helpers import maybe_await
+from a2a.utils.helpers import maybe_await, validate, validate_async_generator
 
 
 logger = logging.getLogger(__name__)
@@ -170,6 +170,10 @@ async def _handler(
             context, _handler, a2a_v0_3_pb2.SendMessageResponse()
         )
 
+    @validate_async_generator(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def SendStreamingMessage(
         self,
         request: a2a_v0_3_pb2.SendMessageRequest,
@@ -229,6 +233,10 @@ async def _handler(
 
         return await self._handle_unary(context, _handler, a2a_v0_3_pb2.Task())
 
+    @validate_async_generator(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def TaskSubscription(
         self,
         request: a2a_v0_3_pb2.TaskSubscriptionRequest,
@@ -252,6 +260,10 @@ async def _handler(
         async for item in self._handle_stream(context, _handler):
             yield item
 
+    @validate(
+        lambda self: self.agent_card.capabilities.push_notifications,
+        'Push notifications are not supported by the agent',
+    )
     async def CreateTaskPushNotificationConfig(
         self,
         request: a2a_v0_3_pb2.CreateTaskPushNotificationConfigRequest,

src/a2a/compat/v0_3/rest_handler.py

@@ -28,7 +28,7 @@
 from a2a.compat.v0_3 import types as types_v03
 from a2a.compat.v0_3.request_handler import RequestHandler03
 from a2a.server.context import ServerCallContext
-from a2a.utils.helpers import validate
+from a2a.utils.helpers import validate, validate_async_generator
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -78,7 +78,7 @@ async def on_message_send(
         pb2_v03_resp = proto_utils.ToProto.task_or_message(v03_resp)
         return MessageToDict(pb2_v03_resp)
 
-    @validate(
+    @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -134,7 +134,7 @@ async def on_cancel_task(
         pb2_v03_task = proto_utils.ToProto.task(v03_resp)
         return MessageToDict(pb2_v03_task)
 
-    @validate(
+    @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -49,7 +49,7 @@
     UnsupportedOperationError,
     VersionNotSupportedError,
 )
-from a2a.utils.helpers import maybe_await, validate
+from a2a.utils.helpers import maybe_await, validate, validate_async_generator
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -171,7 +171,7 @@ async def on_message_send(
         except A2AError as e:
             return _build_error_response(request_id, e)
 
-    @validate(
+    @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -235,6 +235,10 @@ async def on_cancel_task(
 
         return _build_error_response(request_id, TaskNotFoundError())
 
+    @validate_async_generator(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def on_subscribe_to_task(
         self,
         request: SubscribeToTaskRequest,

src/a2a/server/request_handlers/rest_handler.py

@@ -29,7 +29,7 @@
 )
 from a2a.utils import proto_utils
 from a2a.utils.errors import TaskNotFoundError
-from a2a.utils.helpers import validate
+from a2a.utils.helpers import validate, validate_async_generator
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -87,7 +87,7 @@ async def on_message_send(
             response = a2a_pb2.SendMessageResponse(message=task_or_message)
         return MessageToDict(response)
 
-    @validate(
+    @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -139,7 +139,7 @@ async def on_cancel_task(
             return MessageToDict(task)
         raise TaskNotFoundError
 
-    @validate(
+    @validate_async_generator(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )

tests/compat/v0_3/test_grpc_handler.py

@@ -34,6 +34,10 @@ def sample_agent_card() -> a2a_pb2.AgentCard:
         name='Test Agent',
         description='A test agent',
         version='1.0.0',
+        capabilities=a2a_pb2.AgentCapabilities(
+            streaming=True,
+            push_notifications=True,
+        ),
         supported_interfaces=[
             a2a_pb2.AgentInterface(
                 url='http://jsonrpc.v03.com',
@@ -445,7 +449,10 @@ async def test_get_agent_card_success(
         version='1.0.0',
         protocol_version='0.3',
         preferred_transport='JSONRPC',
-        capabilities=a2a_v0_3_pb2.AgentCapabilities(),
+        capabilities=a2a_v0_3_pb2.AgentCapabilities(
+            streaming=True,
+            push_notifications=True,
+        ),
     )
     assert response == expected_res
 

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -123,6 +123,8 @@ def init_fixtures(self) -> None:
         )
         self.mock_agent_card.capabilities = MagicMock(spec=AgentCapabilities)
         self.mock_agent_card.capabilities.extended_agent_card = True
+        self.mock_agent_card.capabilities.streaming = True
+        self.mock_agent_card.capabilities.push_notifications = True
 
         # Mock supported_interfaces list
         interface = MagicMock(spec=AgentInterface)
@@ -710,6 +712,8 @@ async def test_on_resubscribe_existing_task_success(
             mock_agent_executor, mock_task_store, mock_queue_manager
         )
         self.mock_agent_card = MagicMock(spec=AgentCard)
+        self.mock_agent_card.capabilities = MagicMock(spec=AgentCapabilities)
+        self.mock_agent_card.capabilities.streaming = True
         handler = JSONRPCHandler(self.mock_agent_card, request_handler)
         mock_task = create_task()
         events: list[Any] = [