fastapi: WebSocket disconnected state is not propagated to the application code (proper closures, ping timeouts)

Code first, explanation below.


import asyncio
import uvicorn
from fastapi import FastAPI, WebSocket
from starlette.websockets import WebSocketState, WebSocketDisconnect

app = FastAPI(debug=True)


@app.websocket('/ws')
async def registration(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            if websocket.client_state != WebSocketState.CONNECTED:
                print('Connection broken')  # This line is never reached even if the connection is gone
                break

            some_condition = False  # will be set to True when I need to send something to the client
            if some_condition:
                await websocket.send_text('some text')

            await asyncio.sleep(1)
    except WebSocketDisconnect as e:
        print(f'Connection closed {e.code}') # This line is never reached unless I send anything to the client

if __name__ == '__main__':
    uvicorn.run("main:app", port=5000)

I used the code from the WS web chat example from here: https://fastapi.tiangolo.com/advanced/websockets/ The issue is that unless you call websocket.receive_text() there is no way to get if the connection is already dead or not, even if this fact is already established. An idle WS connection can get closed via three ways that I can think of 1) proper connection procedure (client sends close frame, server confirms) 2) TCP connection dies and ping fails to be sent 3) Ping-pong response timed out. All of them are handled by the websockets library that is used inside FastAPI/Starlette/uvicorn. Since the ping-pong is enabled by default, the actual state of connection is known.

But I don’t see any way to check the underlying connection state, websocket.client_state is alwaysWebSocketState.CONNECTED even after it is closed.

The only way to “probe” the state of the connection and trigger WebSocketDisconnect if it is closed is to periodically run this code:

            try:
                await asyncio.wait_for(
                    websocket.receive_text(), 0.0001
                )
            except asyncio.TimeoutError:
                pass

It works in a sense that it triggers the check of the actual state and raises an exception if it is gone, but it is a wasteful call is the connection is actually still okay (needless receive call). Meanwhile if using websockets library raw there is websocket.closed which is actually updated once the connection is gone with a reason as well.

Is there a way to get the actual connection state without this polling receive_text?

About this issue

Original URL
State: closed
Created 3 years ago
Reactions: 9
Comments: 17

Commits related to this issue

Reimplement infinite loop * WebSocketDisconnect doesn't propagate if you never call websocket.receive or websocket.receive_text as seen here tiangolo/fastapi#3008 * As such, since i don't expect to r... — committed to Espolvoritas/BackEnd by marcost2 3 years ago
Reimplement infinite loop * WebSocketDisconnect doesn't propagate if you never call websocket.receive or websocket.receive_text as seen here tiangolo/fastapi#3008 * As such, since i don't expect to r... — committed to Espolvoritas/BackEnd by marcost2 3 years ago

Most upvoted comments

@asyschikov I’m in essentially the same situation - can see the ping/pong messages, but can’t detect if the client disconnects without closing the session via the starlette websocket. Did you find a workaround? Is there a way to attach the lower level python websockets class to the method to detect the connection state?

dm-intropic on Sep 24, 2021

My solution is



async def _alive_task(websocket: WebSocket):
    try:
        await websocket.receive_text()
    except (WebSocketDisconnect, ConnectionClosedError):
        pass
        
async def _send_data(websocket: WebSocket):
    try:
        while True:
            data = await wait_data()
            await websocket.send_json(data)
    except (WebSocketDisconnect, ConnectionClosedError):
        pass

@router.websocket_route("/api/ws")
async def handle_something(websocket: WebSocket):
    await websocket.accept()
    
    loop = asyncio.get_running_loop()
    alive_task = loop.create_task(
        _alive_task(websocket),
        name=f"WS alive check: {websocket.client}",
    )
    send_task: asyncio.Task = loop.create_task(
        _send_data(websocket),
        name=f"WS data sending: {websocket.client}",
    )
    
    alive_task.add_done_callback(send_task.cancel)
    send_task.add_done_callback(alive_task.cancel)
    
    await asyncio.wait({alive_task, send_task})

goDeni on Feb 7, 2022

@dm-intropic the only workaround I found (and tested in production and it works) is to occasionally “poll” the connection this way:

        try:
            await asyncio.wait_for(
                websocket.receive_text(), 0.0001
            )
        except asyncio.TimeoutError:
            pass

Somewhere under the hood websocket.receive_text actually checks if the connection is closed or not and it it is closed it will raise WebSocketDisconnect that you can catch and act on it.

asyschikov on Nov 11, 2021

I’m having the same issue, any updates on this?

jnordberg on Nov 10, 2022