Websocket getting closed immediately after connecting to FastAPI Endpoint

Question:

I’m trying to connect a websocket aiohttp client to a fastapi websocket endpoint, but I can’t send or recieve any data because it seems that the websocket gets closed immediately after connecting to the endpoint.


server

import uvicorn
from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket('/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    ...


if __name__ == '__main__':
    uvicorn.run('test:app', debug=True, reload=True)

client

import aiohttp
import asyncio

async def main():
    s = aiohttp.ClientSession()
    ws = await s.ws_connect('ws://localhost:8000/ws')
    while True:
        ...

asyncio.run(main())

When I try to send data from the server to the client when a connection is made

server

@app.websocket('/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_text('yo')

client

while True:
   print(await ws.receive())

I always get printed in my client’s console

WSMessage(type=<WSMsgType.CLOSED: 257>, data=None, extra=None)

While in the server’s debug console it says

INFO:     ('127.0.0.1', 59792) - "WebSocket /ws" [accepted]
INFO:     connection open
INFO:     connection closed

When I try to send data from the client to the server

server

@app.websocket('/ws')
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        await websocket.receive_text()

client

ws = await s.ws_connect('ws://localhost:8000/ws')
await ws.send_str('client!')

Nothing happens, I get no message printed out in the server’s console, just the debug message saying the client got accepted, connection opened and closed again.


I have no idea what I’m doing wrong, I followed this tutorial in the fastAPI docs for a websocket and the example there with the js websocket works completely fine.

Asked By: 404kuso

||

Answers:

The connection is closed by either end (client or server), as shown from your code snippets. You would need to have a loop in both the server and the client for being able to await for messages, as well as send messages, continuously (have a look here and here).

Additionally, as per FastAPI’s documentation:

When a WebSocket connection is closed, the await websocket.receive_text() will raise a WebSocketDisconnect
exception, which you can then catch and handle like in this example.

Thus, on server side, you should use a try-except block to catch and handle WebSocketDisconnect exceptions. Below is a working example demonstrating a client (in aiohttp) – server (in FastAPI) communication using websockets. Related examples can be found here and here, as well as here and here.

Working Example

Server

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import uvicorn

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # await for connections
    await websocket.accept()
    
    try:
        # send "Connection established" message to client
        await websocket.send_text("Connection established!")
        
        # await for messages and send messages
        while True:
            msg = await websocket.receive_text()
            if msg.lower() == "close":
                await websocket.close()
                break
            else:
                print(f'CLIENT says - {msg}')
                await websocket.send_text(f"Your message was: {msg}")
                
    except WebSocketDisconnect:
        print("Client disconnected")

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

Client

import aiohttp
import asyncio

async def main():
    async with aiohttp.ClientSession() as session:
        async with session.ws_connect('ws://127.0.0.1:8000/ws') as ws:
            # await for messages and send messages
            async for msg in ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    print(f'SERVER says - {msg.data}')
                    text = input('Enter a message: ')
                    await ws.send_str(text)
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    break

asyncio.run(main())
Answered By: Chris