eager eval on jsonrpc

Author: guglielmo-san

src/a2a/server/routes/jsonrpc_dispatcher.py

@@ -376,10 +376,34 @@ async def _process_streaming_request(
         if stream is None:
             raise UnsupportedOperationError(message='Stream not supported')
 
+        # Eagerly fetch the first item from the stream so that errors raised
+        # before any event is yielded (e.g. task terminal state checks)
+        # propagate here and result in a standard JSON-RPC error response
+        # instead of establishing a broken SSE stream.
+        stream = aiter(stream)
+        try:
+            first_event = await anext(stream)
+        except StopAsyncIteration:
+
+            async def _empty_gen() -> AsyncGenerator[dict[str, Any], None]:
+                if False:
+                    yield {}
+
+            return _empty_gen()
+
         async def _wrap_stream(
+            first_evt: Any,
             st: AsyncGenerator,
         ) -> AsyncGenerator[dict[str, Any], None]:
             try:
+                # Yield the first event
+                stream_response = proto_utils.to_stream_response(first_evt)
+                result = MessageToDict(
+                    stream_response, preserving_proto_field_name=False
+                )
+                yield JSONRPC20Response(result=result, _id=request_id).data
+
+                # Yield the rest of the events
                 async for event in st:
                     stream_response = proto_utils.to_stream_response(event)
                     result = MessageToDict(
@@ -389,7 +413,7 @@ async def _wrap_stream(
             except A2AError as e:
                 yield build_error_response(request_id, e)
 
-        return _wrap_stream(stream)
+        return _wrap_stream(first_event, stream)
 
     async def _handle_send_message(
         self, request_obj: SendMessageRequest, context: ServerCallContext

tests/integration/test_client_server_integration.py

@@ -1019,6 +1019,65 @@ async def mock_generator(*args, **kwargs):
     await client.close()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'error_cls,handler_attr,client_method,request_params',
+    [
+        pytest.param(
+            UnsupportedOperationError,
+            'on_subscribe_to_task',
+            'subscribe',
+            SubscribeToTaskRequest(id='some-id'),
+            id='subscribe',
+        ),
+    ],
+)
+async def test_server_rejects_stream_on_validation_error(
+    transport_setups, error_cls, handler_attr, client_method, request_params
+) -> None:
+    """Verify that the server returns an error directly and doesn't open a stream on validation error."""
+    client = transport_setups.client
+    handler = transport_setups.handler
+
+    async def mock_generator(*args, **kwargs):
+        raise error_cls('Validation failed')
+        yield
+
+    getattr(handler, handler_attr).side_effect = mock_generator
+
+    transport = client._transport
+    
+    if isinstance(transport, (RestTransport, JsonRpcTransport)):
+        # Spy on httpx client to check response headers
+        original_send = transport.httpx_client.send
+        response_headers = {}
+        
+        async def mock_send(*args, **kwargs):
+            resp = await original_send(*args, **kwargs)
+            response_headers['Content-Type'] = resp.headers.get('Content-Type')
+            return resp
+            
+        transport.httpx_client.send = mock_send
+        
+        try:
+            with pytest.raises(error_cls):
+                async for _ in getattr(client, client_method)(request=request_params):
+                    pass
+        finally:
+            transport.httpx_client.send = original_send
+            
+        # Verify that the response content type was NOT text/event-stream
+        assert not response_headers.get('Content-Type', '').startswith('text/event-stream')
+    else:
+        # For gRPC, we just verify it raises the error
+        with pytest.raises(error_cls):
+            async for _ in getattr(client, client_method)(request=request_params):
+                pass
+                
+    getattr(handler, handler_attr).side_effect = None
+    await client.close()
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     'request_kwargs, expected_error_code',