连接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

---

当我试图从客户端向服务器发送数据时

服务器

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

客户端

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

什么都没发生，服务器的控制台没有打印出任何消息，只有调试信息显示客户端已被接受，连接已打开并再次关闭。

---

我不知道我做错了什么，我按照fastAPI文档中关于websocket的这个教程进行操作，其中js websocket的示例完全正常。

- 404kuso

Answer 1

从您的代码片段中可以看出，连接被任一端（客户端或服务器）关闭。您需要在服务器和客户端中都设置一个循环，以便能够await消息并持续发送消息（请参见此处和此处）。

另外，根据FastAPI文档：

当WebSocket连接关闭时，await websocket.receive_text()将引发WebSocketDisconnect异常，您可以像这个例子一样捕获和处理它。

因此，在服务器端，您应该使用try-except块来捕获和处理WebSocketDisconnect异常。以下是一个工作示例，演示了使用Websockets进行客户端（在aiohttp中）-服务器（在FastAPI中）通信。相关示例可以在这里和这里找到，以及这里和这里。

工作示例

服务器

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)

客户端

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

- Chris