Add WebSocketsSansIOProtocol (#2540)

Author: Kludex

docs/settings.md

@@ -91,7 +91,7 @@ Using Uvicorn with watchfiles will enable the following options (which are other
 
 * `--loop <str>` - Set the event loop implementation. The uvloop implementation provides greater performance, but is not compatible with Windows or PyPy. **Options:** *'auto', 'asyncio', 'uvloop'.* **Default:** *'auto'*.
 * `--http <str>` - Set the HTTP protocol implementation. The httptools implementation provides greater performance, but it not compatible with PyPy. **Options:** *'auto', 'h11', 'httptools'.* **Default:** *'auto'*.
-* `--ws <str>` - Set the WebSockets protocol implementation. Either of the `websockets` and `wsproto` packages are supported. Use `'none'` to ignore all websocket requests. **Options:** *'auto', 'none', 'websockets', 'wsproto'.* **Default:** *'auto'*.
+* `--ws <str>` - Set the WebSockets protocol implementation. Either of the `websockets` and `wsproto` packages are supported. There are two versions of `websockets` supported: `websockets` and `websockets-sansio`. Use `'none'` to ignore all websocket requests. **Options:** *'auto', 'none', 'websockets', 'websockets-sansio', 'wsproto'.* **Default:** *'auto'*.
 * `--ws-max-size <int>` - Set the WebSockets max message size, in bytes. Only available with the `websockets` protocol. **Default:** *16777216* (16 MB).
 * `--ws-max-queue <int>` - Set the maximum length of the WebSocket incoming message queue. Only available with the `websockets` protocol. **Default:** *32*.
 * `--ws-ping-interval <float>` - Set the WebSockets ping interval, in seconds. Only available with the `websockets` protocol. **Default:** *20.0*.

pyproject.toml

@@ -92,6 +92,9 @@ filterwarnings = [
     "ignore:Uvicorn's native WSGI implementation is deprecated.*:DeprecationWarning",
     "ignore: 'cgi' is deprecated and slated for removal in Python 3.13:DeprecationWarning",
     "ignore: remove second argument of ws_handler:DeprecationWarning:websockets",
+    "ignore: websockets.legacy is deprecated.*:DeprecationWarning",
+    "ignore: websockets.server.WebSocketServerProtocol is deprecated.*:DeprecationWarning",
+    "ignore: websockets.client.connect is deprecated.*:DeprecationWarning",
 ]
 
 [tool.coverage.run]

tests/conftest.py

@@ -233,9 +233,9 @@ def unused_tcp_port() -> int:
             marks=pytest.mark.skipif(not importlib.util.find_spec("wsproto"), reason="wsproto not installed."),
             id="wsproto",
         ),
+        pytest.param("uvicorn.protocols.websockets.websockets_impl:WebSocketProtocol", id="websockets"),
         pytest.param(
-            "uvicorn.protocols.websockets.websockets_impl:WebSocketProtocol",
-            id="websockets",
+            "uvicorn.protocols.websockets.websockets_sansio_impl:WebSocketsSansIOProtocol", id="websockets-sansio"
         ),
     ]
 )

tests/middleware/test_logging.py

@@ -5,12 +5,12 @@
 import socket
 import sys
 from collections.abc import Iterator
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import httpx
 import pytest
-import websockets
 import websockets.client
+from websockets.protocol import State
 
 from tests.utils import run_server
 from uvicorn import Config
@@ -50,7 +50,7 @@ async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
     await send({"type": "http.response.body", "body": b"", "more_body": False})
 
 
-async def test_trace_logging(caplog: pytest.LogCaptureFixture, logging_config, unused_tcp_port: int):
+async def test_trace_logging(caplog: pytest.LogCaptureFixture, logging_config: dict[str, Any], unused_tcp_port: int):
     config = Config(
         app=app,
         log_level="trace",
@@ -92,8 +92,8 @@ async def test_trace_logging_on_http_protocol(http_protocol_cls, caplog, logging
 
 async def test_trace_logging_on_ws_protocol(
     ws_protocol_cls: WSProtocol,
-    caplog,
-    logging_config,
+    caplog: pytest.LogCaptureFixture,
+    logging_config: dict[str, Any],
     unused_tcp_port: int,
 ):
     async def websocket_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
@@ -105,9 +105,9 @@ async def websocket_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISe
             elif message["type"] == "websocket.disconnect":
                 break
 
-    async def open_connection(url):
+    async def open_connection(url: str):
         async with websockets.client.connect(url) as websocket:
-            return websocket.open
+            return websocket.state is State.OPEN
 
     config = Config(
         app=websocket_app,
@@ -127,7 +127,9 @@ async def open_connection(url):
 
 
 @pytest.mark.parametrize("use_colors", [(True), (False), (None)])
-async def test_access_logging(use_colors: bool, caplog: pytest.LogCaptureFixture, logging_config, unused_tcp_port: int):
+async def test_access_logging(
+    use_colors: bool, caplog: pytest.LogCaptureFixture, logging_config: dict[str, Any], unused_tcp_port: int
+):
     config = Config(app=app, use_colors=use_colors, log_config=logging_config, port=unused_tcp_port)
     with caplog_for_logger(caplog, "uvicorn.access"):
         async with run_server(config):
@@ -141,7 +143,7 @@ async def test_access_logging(use_colors: bool, caplog: pytest.LogCaptureFixture
 
 @pytest.mark.parametrize("use_colors", [(True), (False)])
 async def test_default_logging(
-    use_colors: bool, caplog: pytest.LogCaptureFixture, logging_config, unused_tcp_port: int
+    use_colors: bool, caplog: pytest.LogCaptureFixture, logging_config: dict[str, Any], unused_tcp_port: int
 ):
     config = Config(app=app, use_colors=use_colors, log_config=logging_config, port=unused_tcp_port)
     with caplog_for_logger(caplog, "uvicorn.access"):

tests/middleware/test_proxy_headers.py

@@ -465,6 +465,7 @@ async def websocket_app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISe
         host, port = scope["client"]
         await send({"type": "websocket.accept"})
         await send({"type": "websocket.send", "text": f"{scheme}://{host}:{port}"})
+        await send({"type": "websocket.close"})
 
     app_with_middleware = ProxyHeadersMiddleware(websocket_app, trusted_hosts="*")
     config = Config(

tests/protocols/test_websocket.py

@@ -597,12 +597,9 @@ async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable
             await send_accept_task.wait()
         disconnect_message = await receive()  # type: ignore
 
-    response: httpx.Response | None = None
-
     async def websocket_session(uri: str):
-        nonlocal response
         async with httpx.AsyncClient() as client:
-            response = await client.get(
+            await client.get(
                 f"http://127.0.0.1:{unused_tcp_port}",
                 headers={
                     "upgrade": "websocket",
@@ -619,9 +616,6 @@ async def websocket_session(uri: str):
         send_accept_task.set()
         await asyncio.sleep(0.1)
 
-    assert response is not None
-    assert response.status_code == 500, response.text
-    assert response.text == "Internal Server Error"
     assert disconnect_message == {"type": "websocket.disconnect", "code": 1006}
     await task
 
@@ -916,6 +910,9 @@ async def websocket_session(url: str):
 async def test_server_reject_connection_with_invalid_msg(
     ws_protocol_cls: WSProtocol, http_protocol_cls: HTTPProtocol, unused_tcp_port: int
 ):
+    if ws_protocol_cls.__name__ == "WebSocketsSansIOProtocol":
+        pytest.skip("WebSocketsSansIOProtocol sends both start and body messages in one message.")
+
     async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
         assert scope["type"] == "websocket"
         assert "extensions" in scope and "websocket.http.response" in scope["extensions"]
@@ -947,6 +944,9 @@ async def websocket_session(url: str):
 async def test_server_reject_connection_with_missing_body(
     ws_protocol_cls: WSProtocol, http_protocol_cls: HTTPProtocol, unused_tcp_port: int
 ):
+    if ws_protocol_cls.__name__ == "WebSocketsSansIOProtocol":
+        pytest.skip("WebSocketsSansIOProtocol sends both start and body messages in one message.")
+
     async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):
         assert scope["type"] == "websocket"
         assert "extensions" in scope and "websocket.http.response" in scope["extensions"]
@@ -982,6 +982,8 @@ async def test_server_multiple_websocket_http_response_start_events(
     The server should raise an exception if it sends multiple
     websocket.http.response.start events.
     """
+    if ws_protocol_cls.__name__ == "WebSocketsSansIOProtocol":
+        pytest.skip("WebSocketsSansIOProtocol sends both start and body messages in one message.")
     exception_message: str | None = None
 
     async def app(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallable):

uvicorn/config.py

@@ -25,7 +25,7 @@
 from uvicorn.middleware.wsgi import WSGIMiddleware
 
 HTTPProtocolType = Literal["auto", "h11", "httptools"]
-WSProtocolType = Literal["auto", "none", "websockets", "wsproto"]
+WSProtocolType = Literal["auto", "none", "websockets", "websockets-sansio", "wsproto"]
 LifespanType = Literal["auto", "on", "off"]
 LoopSetupType = Literal["none", "auto", "asyncio", "uvloop"]
 InterfaceType = Literal["auto", "asgi3", "asgi2", "wsgi"]
@@ -47,6 +47,7 @@
     "auto": "uvicorn.protocols.websockets.auto:AutoWebSocketsProtocol",
     "none": None,
     "websockets": "uvicorn.protocols.websockets.websockets_impl:WebSocketProtocol",
+    "websockets-sansio": "uvicorn.protocols.websockets.websockets_sansio_impl:WebSocketsSansIOProtocol",
     "wsproto": "uvicorn.protocols.websockets.wsproto_impl:WSProtocol",
 }
 LIFESPAN: dict[LifespanType, str] = {