fix: resolve streaming endpoint deadlock by pre-consuming request body (#426)

- Fix infinite hang issue in `/v1/message:stream` endpoint
- Pre-consume `request.body()` in `_handle_streaming_request` to prevent
deadlock
- `EventSourceResponse` context was causing `request.body()` consumption
to block
- Add comprehensive error handling for body consumption failures
- Add regression tests for streaming endpoint request body handling

Fixes deadlock where `request.body()` consumption inside
`EventSourceResponse` context caused the event loop to hang
indefinitely.

Breaking changes: None
Backward compatibility: Maintained

Fixes [#1001](#431) 🦕

Release-As: 0.3.4

Author: youngchannelforyou

src/a2a/server/apps/rest/rest_adapter.py

@@ -39,7 +39,7 @@
     rest_error_handler,
     rest_stream_error_handler,
 )
-from a2a.utils.errors import ServerError
+from a2a.utils.errors import InvalidRequestError, ServerError
 
 
 logger = logging.getLogger(__name__)
@@ -120,6 +120,18 @@ async def _handle_streaming_request(
         method: Callable[[Request, ServerCallContext], AsyncIterable[Any]],
         request: Request,
     ) -> EventSourceResponse:
+        # Pre-consume and cache the request body to prevent deadlock in streaming context
+        # This is required because Starlette's request.body() can only be consumed once,
+        # and attempting to consume it after EventSourceResponse starts causes deadlock
+        try:
+            await request.body()
+        except (ValueError, RuntimeError, OSError) as e:
+            raise ServerError(
+                error=InvalidRequestError(
+                    message=f'Failed to pre-consume request body: {e}'
+                )
+            ) from e
+
         call_context = self._context_builder.build(request)
 
         async def event_generator(

tests/server/apps/rest/test_rest_fastapi_app.py

@@ -34,6 +34,27 @@ async def agent_card() -> AgentCard:
     mock_agent_card = MagicMock(spec=AgentCard)
     mock_agent_card.url = 'http://mockurl.com'
     mock_agent_card.supports_authenticated_extended_card = False
+
+    # Mock the capabilities object with streaming disabled
+    mock_capabilities = MagicMock()
+    mock_capabilities.streaming = False
+    mock_agent_card.capabilities = mock_capabilities
+
+    return mock_agent_card
+
+
+@pytest.fixture
+async def streaming_agent_card() -> AgentCard:
+    """Agent card that supports streaming for testing streaming endpoints."""
+    mock_agent_card = MagicMock(spec=AgentCard)
+    mock_agent_card.url = 'http://mockurl.com'
+    mock_agent_card.supports_authenticated_extended_card = False
+
+    # Mock the capabilities object with streaming enabled
+    mock_capabilities = MagicMock()
+    mock_capabilities.streaming = True
+    mock_agent_card.capabilities = mock_capabilities
+
     return mock_agent_card
 
 
@@ -42,6 +63,25 @@ async def request_handler() -> RequestHandler:
     return MagicMock(spec=RequestHandler)
 
 
+@pytest.fixture
+async def streaming_app(
+    streaming_agent_card: AgentCard, request_handler: RequestHandler
+) -> FastAPI:
+    """Builds the FastAPI application for testing streaming endpoints."""
+
+    return A2ARESTFastAPIApplication(
+        streaming_agent_card, request_handler
+    ).build(agent_card_url='/well-known/agent-card.json', rpc_url='')
+
+
+@pytest.fixture
+async def streaming_client(streaming_app: FastAPI) -> AsyncClient:
+    """HTTP client for the streaming FastAPI application."""
+    return AsyncClient(
+        transport=ASGITransport(app=streaming_app), base_url='http://test'
+    )
+
+
 @pytest.fixture
 async def app(
     agent_card: AgentCard, request_handler: RequestHandler
@@ -222,5 +262,95 @@ async def test_send_message_success_task(
     assert expected_response == actual_response
 
 
+@pytest.mark.anyio
+async def test_streaming_message_request_body_consumption(
+    streaming_client: AsyncClient, request_handler: MagicMock
+) -> None:
+    """Test that streaming endpoint properly handles request body consumption.
+
+    This test verifies the fix for the deadlock issue where request.body()
+    was being consumed inside the EventSourceResponse context, causing
+    the application to hang indefinitely.
+    """
+
+    # Mock the async generator response from the request handler
+    async def mock_stream_response():
+        """Mock streaming response generator."""
+        yield Message(
+            message_id='stream_msg_1',
+            role=Role.agent,
+            parts=[Part(TextPart(text='First streaming response'))],
+        )
+        yield Message(
+            message_id='stream_msg_2',
+            role=Role.agent,
+            parts=[Part(TextPart(text='Second streaming response'))],
+        )
+
+    request_handler.on_message_send_stream.return_value = mock_stream_response()
+
+    # Create a valid streaming request
+    request = a2a_pb2.SendMessageRequest(
+        request=a2a_pb2.Message(
+            message_id='test_stream_msg',
+            role=a2a_pb2.ROLE_USER,
+            content=[a2a_pb2.Part(text='Test streaming message')],
+        ),
+        configuration=a2a_pb2.SendMessageConfiguration(),
+    )
+
+    # This should not hang indefinitely (previously it would due to the deadlock)
+    response = await streaming_client.post(
+        '/v1/message:stream',
+        json=json_format.MessageToDict(request),
+        headers={'Accept': 'text/event-stream'},
+        timeout=10.0,  # Reasonable timeout to prevent hanging in tests
+    )
+
+    # The response should be successful
+    response.raise_for_status()
+    assert response.status_code == 200
+    assert 'text/event-stream' in response.headers.get('content-type', '')
+
+    # Verify that the request handler was called
+    request_handler.on_message_send_stream.assert_called_once()
+
+
+@pytest.mark.anyio
+async def test_streaming_endpoint_with_invalid_content_type(
+    streaming_client: AsyncClient, request_handler: MagicMock
+) -> None:
+    """Test streaming endpoint behavior with invalid content type."""
+
+    async def mock_stream_response():
+        yield Message(
+            message_id='stream_msg_1',
+            role=Role.agent,
+            parts=[Part(TextPart(text='Response'))],
+        )
+
+    request_handler.on_message_send_stream.return_value = mock_stream_response()
+
+    request = a2a_pb2.SendMessageRequest(
+        request=a2a_pb2.Message(
+            message_id='test_stream_msg',
+            role=a2a_pb2.ROLE_USER,
+            content=[a2a_pb2.Part(text='Test message')],
+        ),
+        configuration=a2a_pb2.SendMessageConfiguration(),
+    )
+
+    # Send request without proper event-stream headers
+    response = await streaming_client.post(
+        '/v1/message:stream',
+        json=json_format.MessageToDict(request),
+        timeout=10.0,
+    )
+
+    # Should still succeed (the adapter handles content-type internally)
+    response.raise_for_status()
+    assert response.status_code == 200
+
+
 if __name__ == '__main__':
     pytest.main([__file__])