WebSocketDisconnect not caught while waiting on async tasks

First Check

• I added a very descriptive title to this issue.
• I used the GitHub search to find a similar issue and didn’t find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google “How to X in FastAPI” and didn’t find any information.
• I already read and followed all the tutorial in the docs and didn’t find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.

Commit to Help

• I commit to help with one of those options 👆

Example Code

@app.websocket("/ws")
async def my_websocket(websocket: WebSocket):
    await websocket.accept()

    try:
        await websocket.send_text("Hello!")
        await asyncio.sleep(5)
        await websocket.send_text("Finished")
    except WebSocketDisconnect:
        print("Client disconnected")

Description

Example first, background below.

Steps to reproduce using provided example

• Connect to the /ws WebSocket endpoint.
• Confirm client receives “Hello!” message from the server.
• Disconnect the client within 5 seconds.

Expected result

The WebSocketDisconnect error is caught as soon as the client closes the connection and “Client disconnected” is printed in the console.

Actual result

No WebSocketDisconnect exception is raised and instead, websockets.exceptions.ConnectionClosedError is raised, only when the server attempts to send a message back to the client.

With debug logging enabled, I can see the following at the exact moment the client disconnects:

DEBUG:websockets.protocol:server - event = data_received(<6 bytes>)
DEBUG:websockets.protocol:server < Frame(fin=True, opcode=<Opcode.CLOSE: 8>, data=b'', rsv1=False, rsv2=False, rsv3=False)
DEBUG:websockets.protocol:server - state = CLOSING
DEBUG:websockets.protocol:server > Frame(fin=True, opcode=<Opcode.CLOSE: 8>, data=b'', rsv1=False, rsv2=False, rsv3=False)
DEBUG:websockets.protocol:server x half-closing TCP connection
DEBUG:websockets.protocol:server - event = eof_received()
DEBUG:websockets.protocol:server - event = connection_lost(None)
DEBUG:websockets.protocol:server - state = CLOSED
DEBUG:websockets.protocol:server x code = 1005, reason = [no reason]

So, the websockets library does receive an event as soon as the connection is closed. But this is not propagated to a WebSocketDisconnect error and I’ve found no way to catch this as it happens.

Background

In my real code, at some point, the server must await a long-running task and send a message to the client when it has finished. This task is potentially expensive - as such, it is strongly preferable to cancel it as soon as a client disconnects, instead of waiting for it to finish only to find out at that point that the client is no longer there.

In my testing, I’ve seen that the WebSocketDisconnect error is raised in some cases, of course, for example:

@app.websocket("/ws")
async def my_websocket(websocket: WebSocket):
    await websocket.accept()

    try:
        await websocket.receive_text()
    except WebSocketDisconnect:
        print("Client disconnected")

If the client connects and then disconnects before sending text, then the disconnection is caught successfully. But, not if the client disconnects during a task that the server is waiting on.

Operating System

macOS

Operating System Details

No response

FastAPI Version

0.68.0

Python Version

3.9.5

Additional Context

No response

[goDeni]

There is no correct solution if your coroutine only send data and not receive

Only trying to receive data to sync websocket state?..

But this is not good solution because why more simpler raise cancell task exception for this coroutine And debug logs know about connection lost

Is there a plan to fix this issue in the future? @tiangolo

Thanks for answer in advance.

[trondhindenes]

I just had the same issue. What you can do, is to create two tasks: one for your long-running thing and one for awaiting client text. The latter will raise an exception the moment the client disconnects, so you can use asyncio.FIRST_COMPLETED to return out of the task awaiter. Here’s how it could look (my long-running task is a blocking redis call that waits for a list):

async def get_receive_text(websocket) -> Tuple[Union[str, None], bool]:
    """The second return value bolean will return True if the client disconnected"""
    try:
        return await websocket.receive_text(), False
    except WebSocketDisconnect as e:
        return None, True


async def get_redis_data(list_name) -> Tuple[Union[str, None], bool]:
    """The second return value bolean will never return true, just there for task consistency"""
    redis = aioredis.from_url("redis://localhost:6379")
    async with redis.client() as conn:
        result = await conn.blpop(keys=[list_name], timeout=30)
        if result:
            return result[1].decode(), False
        return None, False


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    conn_id = random.randint(0, 999)
    print(f'{conn_id=}')
    client_is_connected = True
    while client_is_connected:
        # Wait for either text from the websocket channel, OR the redis list (which will ultimately time out)
        task1 = get_receive_text(websocket)
        task2 = get_redis_data(list_name=conn_id)
        if websocket.client_state.name != 'CONNECTED':
            client_is_connected = False
            break
        # asyncio wait will return when either text is received from the client, 
        # or the redis blocking call returned data
        # or the redis blocking call timed out
        # or the client disconnected
        done, pending = await asyncio.wait([task1, task2], return_when=asyncio.FIRST_COMPLETED)
        done_task = done.pop()
        while len(pending) > 0:
            # cancel the other tasks in the waiter
            pending_task = pending.pop()
            pending_task.cancel()
        result, connection_closed = done_task.result()
        if connection_closed:
            client_is_connected = False
            break
        if result:
            print(result)
            await websocket.send_text(result)