Merge "Validate agent capabilities before invoking api" into main

rajeshvelicheti · Gerrit Code Review · commit 045f8a47c97f · 2025-05-13T19:24:40.000-07:00
diff --git a/src/a2a/server/request_handlers/jsonrpc_handler.py b/src/a2a/server/request_handlers/jsonrpc_handler.py
@@ -3,9 +3,7 @@
 from collections.abc import AsyncIterable
 
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.server.request_handlers.response_helpers import (
-    prepare_response_object,
-)
+from a2a.server.request_handlers.response_helpers import prepare_response_object
 from a2a.types import (
     AgentCard,
     CancelTaskRequest,
@@ -37,6 +35,7 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import ServerError
+from a2a.utils.helpers import validate
 
 
 logger = logging.getLogger(__name__)
@@ -83,6 +82,10 @@ async def on_message_send(
             )
 
     # message/stream
+    @validate(
+        lambda self: self.agent_card.capabilities.streaming,
+        'Streaming is not supported by the agent',
+    )
     async def on_message_send_stream(
         self, request: SendStreamingMessageRequest
     ) -> AsyncIterable[SendStreamingMessageResponse]:
@@ -183,6 +186,10 @@ async def get_push_notification(
             )
 
     # tasks/pushNotification/set
+    @validate(
+        lambda self: self.agent_card.capabilities.pushNotifications,
+        'Push notifications are not supported by the agent',
+    )
     async def set_push_notification(
         self, request: SetTaskPushNotificationConfigRequest
     ) -> SetTaskPushNotificationConfigResponse:
diff --git a/src/a2a/utils/helpers.py b/src/a2a/utils/helpers.py
@@ -12,6 +12,7 @@
     TaskStatus,
     TextPart,
 )
+from a2a.utils.errors import ServerError, UnsupportedOperationError
 
 
 logger = logging.getLogger(__name__)
@@ -82,3 +83,22 @@ def build_text_artifact(text: str, artifact_id: str) -> Artifact:
     text_part = TextPart(text=text)
     part = Part(root=text_part)
     return Artifact(parts=[part], artifactId=artifact_id)
+
+
+def validate(expression, error_message=None):
+    """Decorator that validates if the given expression evaluates to True."""
+
+    def decorator(function):
+        def wrapper(self, *args, **kwargs):
+            if not expression(self):
+                if not error_message:
+                    message = str(expression)
+                logger.error(f'Unsuppported Operation: {error_message}')
+                raise ServerError(
+                    UnsupportedOperationError(message=error_message)
+                )
+            return function(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
