Receiving on a disconnected websocket emits an incorrect error message

[bryzhang]

Hi, I'm messaging on behalf of Baseten. We've noticed that when we call "receive" on a disconnected websocket, we get an incorrect error message, 'WebSocket is not connected. Need to call "accept" first.'. Instead, it would be helpful to return a clearer message for a disconnected websocket, like, WebSocket is already disconnected. Cannot receive.

I've attached a patch that essentially makes this change (since I had some trouble pushing to create a PR):
0001-emit-clearer-error-message-when-receiving-on-a-disco.patch

[Kludex]

Can you provide a minimal reproducible example?

[bryzhang]

Yeah, I've attached a zip with a simple server and client code. Just run python3 server.py and then python3 client.py. You will see that the output's first error on the server is the message I alluded to.

starlette_mre.zip

[Kludex]

I'll not open a zip file. Please share the snippet here.

[bryzhang]

Server code:

from starlette.applications import Starlette
from starlette.routing import WebSocketRoute
from starlette.websockets import WebSocket
import uvicorn


async def websocket_handler(websocket: WebSocket):
    await websocket.accept()
    print("Client connected")

    try:
        while True:
            data = await websocket.receive_text()
            try:
                await websocket.send_text(data)
            except Exception as e:
                continue
    except Exception as e:
        print(f"Connection closed: {e}")
    finally:
        await websocket.close()
        print("Connection closed")


app = Starlette(
    debug=True,
    routes=[
        WebSocketRoute("/ws", websocket_handler)
    ]
)

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8052)

Client code:

import asyncio
import websockets


async def connect():
    uri = "ws://127.0.0.1:8052/ws"
    async with websockets.connect(uri) as websocket:
        await websocket.send("hi")
        print(await websocket.recv())
        await websocket.send("hi")
        print(await websocket.recv())
        await websocket.send("hi")
        # exit before receiving, so that the server's send() fails
        exit()
        
asyncio.run(connect())

[Kludex]

I can't see the issue you are describing. I see this:

    raise RuntimeError(msg % message_type)
RuntimeError: Unexpected ASGI message 'websocket.close', after sending 'websocket.close' or response already completed.

Which seems correct...

[bryzhang]

Yeah, it's not a perfect example of what we saw. We had a different RuntimeError that was less clear (it didn't indicate that the websocket was closed already). But we still had the message at the top, preceding the exception, which is the less-clear "need to call 'accept' first". Even in this case, I think it would be preferable to simply say that the websocket is disconnected?