Update fun.py

fun.py

@@ -34,7 +34,6 @@ async def proxy(request: Request, path: str):
         headers=dict(resp.headers)
     )
 
-@app.websocket("/{path:path}")
 async def websocket_proxy(websocket: WebSocket, path: str):
     await websocket.accept()
     target_url = f"{TARGET_WS_BASE}/{path}"
@@ -42,64 +41,40 @@ async def websocket_proxy(websocket: WebSocket, path: str):
 
     try:
         async with websockets.connect(target_url) as target_ws:
-            disconnect_event = asyncio.Event()
+            # 1. 先发送初始化命令（必需）
+            await target_ws.send(json.dumps({
+                "id": 1,
+                "method": "Runtime.enable"
+            }))
+            logger.info("Sent Runtime.enable")
 
+            # 2. 启动转发任务
             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()
+                    while True:
+                        data = await websocket.receive_text()
+                        logger.info(f"CLIENT->CHROME: {data[:200]}")
+                        await target_ws.send(data)
+                except WebSocketDisconnect:
+                    logger.info("Client disconnected")
 
             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())
+                    while True:
+                        response = await target_ws.recv()
+                        logger.info(f"CHROME->CLIENT: {response[:200]}")
+                        await websocket.send_text(response)
+                except websockets.exceptions.ConnectionClosed:
+                    logger.info("Browser disconnected")
 
-            # 等待任意任务完成
-            done, pending = await asyncio.wait(
-                [client_task, server_task],
-                return_when=asyncio.FIRST_COMPLETED
+            await asyncio.gather(
+                forward_client_to_server(),
+                forward_server_to_client()
             )
-
-            # 清理未完成的任务
-            disconnect_event.set()
-            for task in pending:
-                task.cancel()
-            await asyncio.gather(*pending, return_exceptions=True)
-
     except Exception as e:
-        logger.error(f"Connection error: {e}")
+        logger.error(f"Error: {e}")
     finally:
-        try:
-            await websocket.close()
-        except Exception:
-            pass
+        await websocket.close()
 
  
 if __name__ == "__main__":