连接到 FastAPI 端点后 Websocket 立即关闭

Question

我正在尝试将 websocket aiohttp 客户端连接到 fastapi websocket 端点，但我无法发送或接收任何数据，因为似乎 websocket 在连接到端点后立即关闭。

---

服务器

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)

客户

import aiohttp
import asyncio

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

asyncio.run(main())

---

当我尝试在建立连接时将数据从服务器发送到客户端时

服务器

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

客户

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

我总是在客户的控制台中打印

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

在服务器的调试控制台中它说

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

---

当我尝试将数据从客户端发送到服务器时

服务器

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

客户

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

什么也没发生，我在服务器的控制台中没有打印任何消息，只有调试消息说客户端已被接受，连接再次打开和关闭。

---

我不知道我做错了什么，我按照fastAPI 文档中的 websocket教程进行操作，并且 js websocket 的示例工作得很好。

Answer 1

连接由任一端（客户端或服务器）关闭，如代码片段所示。您需要在服务器和客户端中都有一个循环，以便能够await连续地获取消息以及发送消息（请查看此处和此处）。

此外，根据FastAPI 的文档：

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, as well as websockets.exceptions.ConnectionClosed exceptions, as explained in this answer. 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
from websockets.exceptions import ConnectionClosed
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, ConnectionClosed):
        print("Client disconnected")

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

Client

Examples using the websockets library instead of aiohttp can be found here, as well as here and here.

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())