连接到 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())

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

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!')

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

---

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

Answer 1

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

此外，根据 FastAPI 文档：

当 WebSocket 连接关闭时，await websocket.receive_text() 将引发 WebSocketDisconnect
异常，然后您可以像本示例中那样捕获并处理该异常。

因此，在服务器端，您应该使用 try-except 块来捕获和处理 WebSocketDisconnect 异常，以及 websockets.exceptions.ConnectionClosed例外情况，如此答案中所述。下面是一个工作示例，演示了使用 websockets 进行客户端（在 aiohttp 中）- 服务器（在 FastAPI 中）通信。相关示例可以在此处和这里，以及此处和此处。

工作示例

服务器

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)

客户端

使用 websockets 库而不是 aiohttp 的示例可以在此处，以及此处和此处。

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