Problem with `await self.task_manager.on_send_task_subscribe` call

[joyoungzhang]

Problem Description

In the current implementation, there is an issue with the await self.task_manager.on_send_task_subscribe(json_rpc_request) call in server.py. The problem occurs because:

1. The method on_send_task_subscribe() is declared to return Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]

2. When it returns an AsyncIterable (asynchronous generator), we cannot directly use await on it

3. This causes the error: TypeError: object async_generator can't be used in 'await' expression

Expected Behavior

For streaming requests (SendTaskStreamingRequest), the system should properly handle the asynchronous streaming response without raising exceptions.

Current Behavior

The code attempts to await an asynchronous generator, which is not a valid operation in Python, resulting in the mentioned error.

Solution Options

Option 1: Change the calling code (Recommended)

elif isinstance(json_rpc_request, SendTaskStreamingRequest):
    # Don't await the async generator, pass it directly
    result = self.task_manager.on_send_task_subscribe(json_rpc_request)

Option 2: Modify the method implementation (If keeping await is necessary)
Change on_send_task_subscribe() to return a coroutine that collects all streaming responses:

async def on_send_task_subscribe(
    self, request: SendTaskStreamingRequest
) -> Union[List[SendTaskStreamingResponse], JSONRPCResponse]:
    responses = []
    async for item in original_async_gen:
        responses.append(item)
    return responses

Recommendation

Option 1 is preferred as it maintains the streaming nature of the response. Option 2 would convert the streaming response into a batch response, which may not be desirable for the intended use case.

[holtskinner]

Hello! Thank you for filing this detailed bug report and for providing such a clear analysis and recommended solution.

Based on my review of the current codebase, it appears this issue has been addressed and the implementation aligns with your recommended "Option 1".

In the a2a-python repository, the file src/a2a/server/apps/jsonrpc/jsonrpc_app.py is responsible for handling incoming RPC requests. Specifically, in the _process_streaming_request method, it calls the handler for streaming messages (on_message_send_stream) and correctly handles the returned asynchronous generator without using await on it directly.

# From src/a2a/server/apps/jsonrpc/jsonrpc_app.py

# ...
        if isinstance(
            request_obj,
            SendStreamingMessageRequest,
        ):
            handler_result = self.handler.on_message_send_stream(
                request_obj, context
            )
# ...

The handler_result, which is an AsyncGenerator, is then passed to the _create_response method. This method correctly wraps the generator in a starlette.responses.EventSourceResponse to stream the events to the client.

It's possible that the method and request names have been updated since the version of the code you were referencing. For instance, SendTaskStreamingRequest is now SendStreamingMessageRequest, and the handler method is on_message_send_stream.

Your analysis was spot-on, and the current implementation avoids the TypeError you described by not awaiting the async generator. Please feel free to check against the latest version on the main branch. I'll be closing this issue, but please let us know if you encounter any further problems.