Remove validate_async_generator

Author: ishymko

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, validate, validate_async_generator
+from a2a.utils.helpers import maybe_await, validate
 
 
 logger = logging.getLogger(__name__)
@@ -177,7 +177,7 @@ async def SendStreamingMessage(
     ) -> AsyncIterable[a2a_v0_3_pb2.StreamResponse]:
         """Handles the 'SendStreamingMessage' gRPC method (v0.3)."""
 
-        @validate_async_generator(
+        @validate(
             lambda _: self.agent_card.capabilities.streaming,
             'Streaming is not supported by the agent',
         )
@@ -240,7 +240,7 @@ async def TaskSubscription(
     ) -> AsyncIterable[a2a_v0_3_pb2.StreamResponse]:
         """Handles the 'TaskSubscription' gRPC method (v0.3)."""
 
-        @validate_async_generator(
+        @validate(
             lambda _: self.agent_card.capabilities.streaming,
             'Streaming is not supported by the agent',
         )

src/a2a/compat/v0_3/rest_handler.py

@@ -31,7 +31,6 @@
 from a2a.utils import constants
 from a2a.utils.helpers import (
     validate,
-    validate_async_generator,
     validate_version,
 )
 from a2a.utils.telemetry import SpanKind, trace_class
@@ -85,7 +84,7 @@ async def on_message_send(
         return MessageToDict(pb2_v03_resp)
 
     @validate_version(constants.PROTOCOL_VERSION_0_3)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -143,7 +142,7 @@ async def on_cancel_task(
         return MessageToDict(pb2_v03_task)
 
     @validate_version(constants.PROTOCOL_VERSION_0_3)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )

src/a2a/server/request_handlers/grpc_handler.py

@@ -40,7 +40,7 @@
     A2AError,
     TaskNotFoundError,
 )
-from a2a.utils.helpers import maybe_await, validate, validate_async_generator
+from a2a.utils.helpers import maybe_await, validate
 
 
 logger = logging.getLogger(__name__)
@@ -197,7 +197,7 @@ async def SendStreamingMessage(
     ) -> AsyncIterable[a2a_pb2.StreamResponse]:
         """Handles the 'StreamMessage' gRPC method."""
 
-        @validate_async_generator(
+        @validate(
             lambda _: self.agent_card.capabilities.streaming,
             'Streaming is not supported by the agent',
         )
@@ -238,7 +238,7 @@ async def SubscribeToTask(
     ) -> AsyncIterable[a2a_pb2.StreamResponse]:
         """Handles the 'SubscribeToTask' gRPC method."""
 
-        @validate_async_generator(
+        @validate(
             lambda _: self.agent_card.capabilities.streaming,
             'Streaming is not supported by the agent',
         )

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -52,7 +52,6 @@
 from a2a.utils.helpers import (
     maybe_await,
     validate,
-    validate_async_generator,
     validate_version,
 )
 from a2a.utils.telemetry import SpanKind, trace_class
@@ -178,7 +177,7 @@ async def on_message_send(
             return _build_error_response(request_id, e)
 
     @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -244,7 +243,7 @@ async def on_cancel_task(
         return _build_error_response(request_id, TaskNotFoundError())
 
     @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )

src/a2a/server/request_handlers/rest_handler.py

@@ -31,7 +31,6 @@
 from a2a.utils.errors import TaskNotFoundError
 from a2a.utils.helpers import (
     validate,
-    validate_async_generator,
     validate_version,
 )
 from a2a.utils.telemetry import SpanKind, trace_class
@@ -93,7 +92,7 @@ async def on_message_send(
         return MessageToDict(response)
 
     @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )
@@ -147,7 +146,7 @@ async def on_cancel_task(
         raise TaskNotFoundError
 
     @validate_version(constants.PROTOCOL_VERSION_1_0)
-    @validate_async_generator(
+    @validate(
         lambda self: self.agent_card.capabilities.streaming,
         'Streaming is not supported by the agent',
     )

src/a2a/utils/helpers.py

@@ -232,90 +232,6 @@ def sync_wrapper(self: Any, *args, **kwargs) -> Any:
     return decorator
 
 
-def validate_async_generator(
-    expression: Callable[[Any], bool], error_message: str | None = None
-):
-    """Decorator that validates if a given expression evaluates to True for async generators.
-
-    Typically used on class methods to check capabilities or configuration
-    before executing the method's logic. If the expression is False,
-    an `UnsupportedOperationError` is raised.
-
-    Args:
-        expression: A callable that takes the instance (`self`) as its argument
-                    and returns a boolean.
-        error_message: An optional custom error message for the `UnsupportedOperationError`.
-                       If None, the string representation of the expression will be used.
-
-    Examples:
-        Streaming capability validation with success case:
-        >>> import asyncio
-        >>> from a2a.utils.errors import UnsupportedOperationError
-        >>>
-        >>> class StreamingAgent:
-        ...     def __init__(self, streaming_enabled: bool):
-        ...         self.streaming_enabled = streaming_enabled
-        ...
-        ...     @validate_async_generator(
-        ...         lambda self: self.streaming_enabled,
-        ...         'Streaming is not supported by this agent',
-        ...     )
-        ...     async def stream_messages(self, count: int):
-        ...         for i in range(count):
-        ...             yield f'Message {i}'
-        >>>
-        >>> async def run_streaming_test():
-        ...     # Successful streaming
-        ...     agent = StreamingAgent(streaming_enabled=True)
-        ...     async for msg in agent.stream_messages(2):
-        ...         print(msg)
-        >>>
-        >>> asyncio.run(run_streaming_test())
-        Message 0
-        Message 1
-
-        Error case - validation fails:
-        >>> class FeatureAgent:
-        ...     def __init__(self):
-        ...         self.features = {'real_time': False}
-        ...
-        ...     @validate_async_generator(
-        ...         lambda self: self.features.get('real_time', False),
-        ...         'Real-time feature must be enabled to stream updates',
-        ...     )
-        ...     async def real_time_updates(self):
-        ...         yield 'This should not be yielded'
-        >>>
-        >>> async def run_error_test():
-        ...     agent = FeatureAgent()
-        ...     try:
-        ...         async for _ in agent.real_time_updates():
-        ...             pass
-        ...     except UnsupportedOperationError as e:
-        ...         print(e.message)
-        >>>
-        >>> asyncio.run(run_error_test())
-        Real-time feature must be enabled to stream updates
-
-    Note:
-        This decorator is specifically for async generator methods (async def with yield).
-        The validation happens before the generator starts yielding values.
-    """
-
-    def decorator(function):
-        @functools.wraps(function)
-        def wrapper(self, *args, **kwargs):
-            if not expression(self):
-                final_message = error_message or str(expression)
-                logger.error('Unsupported Operation: %s', final_message)
-                raise UnsupportedOperationError(message=final_message)
-            return function(self, *args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
 def are_modalities_compatible(
     server_output_modes: list[str] | None, client_output_modes: list[str] | None
 ) -> bool:

tests/integration/test_client_server_integration.py

@@ -1032,10 +1032,10 @@ async def test_validate_decorator_push_notifications_disabled(
 
 
 @pytest.mark.asyncio
-async def test_validate_async_generator_streaming_disabled(
+async def test_validate_streaming_disabled(
     error_handling_setups, agent_card: AgentCard
 ) -> None:
-    """Integration test for @validate_async_generator decorator when streaming is disabled."""
+    """Integration test for @validate decorator when streaming is disabled."""
     client = error_handling_setups.client
     transport = client._transport