Update fun.py

fun.py

@@ -42,21 +42,65 @@ async def websocket_proxy(websocket: WebSocket, path: str):
 
     try:
         async with websockets.connect(target_url) as target_ws:
-            async def forward(from_ws, to_ws, direction):
-                while True:
-                    data = await from_ws.receive()
-                    logger.info(f"{direction} RAW DATA: {data}")  # 打印原始数据
-                    if isinstance(data, str):
-                        await to_ws.send(data)
-                    else:
-                        await to_ws.send(data.bytes if hasattr(data, 'bytes') else data)
-
-            await asyncio.gather(
-                forward(websocket, target_ws, "CLIENT->CHROME"),
-                forward(target_ws, websocket, "CHROME->CLIENT")
+            disconnect_event = asyncio.Event()
+
+            async def forward_client_to_server():
+                try:
+                    while not disconnect_event.is_set():
+                        # 正确处理 FastAPI WebSocket 接收
+                        message = await websocket.receive()
+                        logger.info(f"CLIENT->CHROME RAW DATA: {message}")
+
+                        if message["type"] == "websocket.disconnect":
+                            logger.info("Client disconnected")
+                            disconnect_event.set()
+                            break
+
+                        if message["type"] == "websocket.receive.text":
+                            await target_ws.send(message["text"])
+                        elif message["type"] == "websocket.receive.bytes":
+                            await target_ws.send(message["bytes"])
+                except Exception as e:
+                    logger.error(f"Client->Server error: {e}")
+                    disconnect_event.set()
+
+            async def forward_server_to_client():
+                try:
+                    while not disconnect_event.is_set():
+                        # 正确处理 websockets 库接收
+                        message = await target_ws.recv()
+                        if isinstance(message, str):
+                            await websocket.send_text(message)
+                        elif isinstance(message, bytes):
+                            await websocket.send_bytes(message)
+                except Exception as e:
+                    logger.error(f"Server->Client error: {e}")
+                    disconnect_event.set()
+
+            # 并行运行两个转发任务
+            client_task = asyncio.create_task(forward_client_to_server())
+            server_task = asyncio.create_task(forward_server_to_client())
+
+            # 等待任意任务完成
+            done, pending = await asyncio.wait(
+                [client_task, server_task],
+                return_when=asyncio.FIRST_COMPLETED
             )
+
+            # 清理未完成的任务
+            disconnect_event.set()
+            for task in pending:
+                task.cancel()
+            await asyncio.gather(*pending, return_exceptions=True)
+
     except Exception as e:
-        logger.error(f"Error: {e}")
+        logger.error(f"Connection error: {e}")
+    finally:
+        try:
+            await websocket.close()
+        except Exception:
+            pass
+
  
 if __name__ == "__main__":
     uvicorn.run("fun:app", host="0.0.0.0", port=8000, reload=True)