Update fun.py

fun.py

@@ -41,55 +41,30 @@ async def websocket_proxy(websocket: WebSocket, path: str):
     logger.info(f"Forwarding WebSocket to: {target_url}")
 
     try:
-        # 创建自定义连接
-        async def create_connection():
-            return await websockets.connect(
-                target_url,
-                # extra_headers={"Host": "127.0.0.1"},
-                ping_interval=20,
-                ping_timeout=30
-            )
-
-        target_ws = await create_connection()
-
-        async def forward_client_to_server():
-            try:
+        async with websockets.connect(target_url) as target_ws:
+            async def forward_client_to_server():
                 while True:
-                    # 明确处理文本和二进制消息
                     message = await websocket.receive()
                     if message["type"] == "websocket.receive.text":
                         await target_ws.send(message["text"])
-                        logger.info(f"Sent text: {message['text'][:100]}")
-                    elif message["type"] == "websocket.receive.bytes":
-                        await target_ws.send(message["bytes"])
-                        logger.info(f"Sent bytes: {message['bytes'][:100]}")
-            except Exception as e:
-                logger.error(f"Client->Server error: {e}")
-                raise
+                        logger.info(f"Sent to Chrome: {message['text'][:100]}")
 
-        async def forward_server_to_client():
-            try:
+            async def forward_server_to_client():
                 while True:
                     message = await target_ws.recv()
                     if isinstance(message, str):
                         await websocket.send_text(message)
-                        logger.info(f"Received text: {message[:100]}")
+                        logger.info(f"Received from Chrome: {message[:100]}")
                     elif isinstance(message, bytes):
                         await websocket.send_bytes(message)
-                        logger.info(f"Received bytes: {message[:100]}")
-            except Exception as e:
-                logger.error(f"Server->Client error: {e}")
-                raise
 
-        await asyncio.gather(
-            forward_client_to_server(),
-            forward_server_to_client()
-        )
-
-    except WebSocketDisconnect:
-        logger.info("Client disconnected")
+            await asyncio.gather(
+                forward_client_to_server(),
+                forward_server_to_client()
+            )
     except Exception as e:
-        logger.error(f"WebSocket error: {str(e)}")
+        logger.error(f"Error: {e}")
         await websocket.close()
+ 
 if __name__ == "__main__":
     uvicorn.run("fun:app", host="0.0.0.0", port=8000, reload=True)