Skip to content

Commit b7f323f

Browse files
committed
Remove validate_async_generator
1 parent 8ea34a4 commit b7f323f

7 files changed

Lines changed: 14 additions & 101 deletions

File tree

src/a2a/compat/v0_3/grpc_handler.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from a2a.server.request_handlers.request_handler import RequestHandler
3030
from a2a.types.a2a_pb2 import AgentCard
3131
from a2a.utils.errors import A2AError, InvalidParamsError
32-
from a2a.utils.helpers import maybe_await, validate, validate_async_generator
32+
from a2a.utils.helpers import maybe_await, validate
3333

3434

3535
logger = logging.getLogger(__name__)
@@ -177,7 +177,7 @@ async def SendStreamingMessage(
177177
) -> AsyncIterable[a2a_v0_3_pb2.StreamResponse]:
178178
"""Handles the 'SendStreamingMessage' gRPC method (v0.3)."""
179179

180-
@validate_async_generator(
180+
@validate(
181181
lambda _: self.agent_card.capabilities.streaming,
182182
'Streaming is not supported by the agent',
183183
)
@@ -240,7 +240,7 @@ async def TaskSubscription(
240240
) -> AsyncIterable[a2a_v0_3_pb2.StreamResponse]:
241241
"""Handles the 'TaskSubscription' gRPC method (v0.3)."""
242242

243-
@validate_async_generator(
243+
@validate(
244244
lambda _: self.agent_card.capabilities.streaming,
245245
'Streaming is not supported by the agent',
246246
)

src/a2a/compat/v0_3/rest_handler.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
from a2a.utils import constants
3232
from a2a.utils.helpers import (
3333
validate,
34-
validate_async_generator,
3534
validate_version,
3635
)
3736
from a2a.utils.telemetry import SpanKind, trace_class
@@ -85,7 +84,7 @@ async def on_message_send(
8584
return MessageToDict(pb2_v03_resp)
8685

8786
@validate_version(constants.PROTOCOL_VERSION_0_3)
88-
@validate_async_generator(
87+
@validate(
8988
lambda self: self.agent_card.capabilities.streaming,
9089
'Streaming is not supported by the agent',
9190
)
@@ -143,7 +142,7 @@ async def on_cancel_task(
143142
return MessageToDict(pb2_v03_task)
144143

145144
@validate_version(constants.PROTOCOL_VERSION_0_3)
146-
@validate_async_generator(
145+
@validate(
147146
lambda self: self.agent_card.capabilities.streaming,
148147
'Streaming is not supported by the agent',
149148
)

src/a2a/server/request_handlers/grpc_handler.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
A2AError,
4141
TaskNotFoundError,
4242
)
43-
from a2a.utils.helpers import maybe_await, validate, validate_async_generator
43+
from a2a.utils.helpers import maybe_await, validate
4444

4545

4646
logger = logging.getLogger(__name__)
@@ -197,7 +197,7 @@ async def SendStreamingMessage(
197197
) -> AsyncIterable[a2a_pb2.StreamResponse]:
198198
"""Handles the 'StreamMessage' gRPC method."""
199199

200-
@validate_async_generator(
200+
@validate(
201201
lambda _: self.agent_card.capabilities.streaming,
202202
'Streaming is not supported by the agent',
203203
)
@@ -238,7 +238,7 @@ async def SubscribeToTask(
238238
) -> AsyncIterable[a2a_pb2.StreamResponse]:
239239
"""Handles the 'SubscribeToTask' gRPC method."""
240240

241-
@validate_async_generator(
241+
@validate(
242242
lambda _: self.agent_card.capabilities.streaming,
243243
'Streaming is not supported by the agent',
244244
)

src/a2a/server/request_handlers/jsonrpc_handler.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@
5252
from a2a.utils.helpers import (
5353
maybe_await,
5454
validate,
55-
validate_async_generator,
5655
validate_version,
5756
)
5857
from a2a.utils.telemetry import SpanKind, trace_class
@@ -178,7 +177,7 @@ async def on_message_send(
178177
return _build_error_response(request_id, e)
179178

180179
@validate_version(constants.PROTOCOL_VERSION_1_0)
181-
@validate_async_generator(
180+
@validate(
182181
lambda self: self.agent_card.capabilities.streaming,
183182
'Streaming is not supported by the agent',
184183
)
@@ -244,7 +243,7 @@ async def on_cancel_task(
244243
return _build_error_response(request_id, TaskNotFoundError())
245244

246245
@validate_version(constants.PROTOCOL_VERSION_1_0)
247-
@validate_async_generator(
246+
@validate(
248247
lambda self: self.agent_card.capabilities.streaming,
249248
'Streaming is not supported by the agent',
250249
)

src/a2a/server/request_handlers/rest_handler.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
from a2a.utils.errors import TaskNotFoundError
3232
from a2a.utils.helpers import (
3333
validate,
34-
validate_async_generator,
3534
validate_version,
3635
)
3736
from a2a.utils.telemetry import SpanKind, trace_class
@@ -93,7 +92,7 @@ async def on_message_send(
9392
return MessageToDict(response)
9493

9594
@validate_version(constants.PROTOCOL_VERSION_1_0)
96-
@validate_async_generator(
95+
@validate(
9796
lambda self: self.agent_card.capabilities.streaming,
9897
'Streaming is not supported by the agent',
9998
)
@@ -147,7 +146,7 @@ async def on_cancel_task(
147146
raise TaskNotFoundError
148147

149148
@validate_version(constants.PROTOCOL_VERSION_1_0)
150-
@validate_async_generator(
149+
@validate(
151150
lambda self: self.agent_card.capabilities.streaming,
152151
'Streaming is not supported by the agent',
153152
)

src/a2a/utils/helpers.py

Lines changed: 0 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -232,90 +232,6 @@ def sync_wrapper(self: Any, *args, **kwargs) -> Any:
232232
return decorator
233233

234234

235-
def validate_async_generator(
236-
expression: Callable[[Any], bool], error_message: str | None = None
237-
):
238-
"""Decorator that validates if a given expression evaluates to True for async generators.
239-
240-
Typically used on class methods to check capabilities or configuration
241-
before executing the method's logic. If the expression is False,
242-
an `UnsupportedOperationError` is raised.
243-
244-
Args:
245-
expression: A callable that takes the instance (`self`) as its argument
246-
and returns a boolean.
247-
error_message: An optional custom error message for the `UnsupportedOperationError`.
248-
If None, the string representation of the expression will be used.
249-
250-
Examples:
251-
Streaming capability validation with success case:
252-
>>> import asyncio
253-
>>> from a2a.utils.errors import UnsupportedOperationError
254-
>>>
255-
>>> class StreamingAgent:
256-
... def __init__(self, streaming_enabled: bool):
257-
... self.streaming_enabled = streaming_enabled
258-
...
259-
... @validate_async_generator(
260-
... lambda self: self.streaming_enabled,
261-
... 'Streaming is not supported by this agent',
262-
... )
263-
... async def stream_messages(self, count: int):
264-
... for i in range(count):
265-
... yield f'Message {i}'
266-
>>>
267-
>>> async def run_streaming_test():
268-
... # Successful streaming
269-
... agent = StreamingAgent(streaming_enabled=True)
270-
... async for msg in agent.stream_messages(2):
271-
... print(msg)
272-
>>>
273-
>>> asyncio.run(run_streaming_test())
274-
Message 0
275-
Message 1
276-
277-
Error case - validation fails:
278-
>>> class FeatureAgent:
279-
... def __init__(self):
280-
... self.features = {'real_time': False}
281-
...
282-
... @validate_async_generator(
283-
... lambda self: self.features.get('real_time', False),
284-
... 'Real-time feature must be enabled to stream updates',
285-
... )
286-
... async def real_time_updates(self):
287-
... yield 'This should not be yielded'
288-
>>>
289-
>>> async def run_error_test():
290-
... agent = FeatureAgent()
291-
... try:
292-
... async for _ in agent.real_time_updates():
293-
... pass
294-
... except UnsupportedOperationError as e:
295-
... print(e.message)
296-
>>>
297-
>>> asyncio.run(run_error_test())
298-
Real-time feature must be enabled to stream updates
299-
300-
Note:
301-
This decorator is specifically for async generator methods (async def with yield).
302-
The validation happens before the generator starts yielding values.
303-
"""
304-
305-
def decorator(function):
306-
@functools.wraps(function)
307-
def wrapper(self, *args, **kwargs):
308-
if not expression(self):
309-
final_message = error_message or str(expression)
310-
logger.error('Unsupported Operation: %s', final_message)
311-
raise UnsupportedOperationError(message=final_message)
312-
return function(self, *args, **kwargs)
313-
314-
return wrapper
315-
316-
return decorator
317-
318-
319235
def are_modalities_compatible(
320236
server_output_modes: list[str] | None, client_output_modes: list[str] | None
321237
) -> bool:

tests/integration/test_client_server_integration.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,10 +1032,10 @@ async def test_validate_decorator_push_notifications_disabled(
10321032

10331033

10341034
@pytest.mark.asyncio
1035-
async def test_validate_async_generator_streaming_disabled(
1035+
async def test_validate_streaming_disabled(
10361036
error_handling_setups, agent_card: AgentCard
10371037
) -> None:
1038-
"""Integration test for @validate_async_generator decorator when streaming is disabled."""
1038+
"""Integration test for @validate decorator when streaming is disabled."""
10391039
client = error_handling_setups.client
10401040
transport = client._transport
10411041

0 commit comments

Comments
 (0)