完成线程启动训练并以websocket的方式进行发送shell

app/api/business/websocket_api.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, WebSocket
+from starlette.websockets import WebSocketState
 
 from app.websocket.web_socket_server import room_manager
 
@@ -19,9 +20,11 @@ async def websocket_room(websocket: WebSocket, room: str):
     try:
         while True:
             data = await websocket.receive_text()
-            await room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
+            room_manager.broadcast_to_room(room, data, exclude_websocket=websocket)
     except Exception as e:
-        print(f"连接关闭: {e}")
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close(code=1000)
     finally:
         await room_manager.remove_from_room(room, websocket)
-        await websocket.close()
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close(code=1001)