improve socket stability

viser_proxy_manager.py

@@ -27,7 +27,7 @@ class ViserProxyManager:
         app: FastAPI,
         min_local_port: int = 8000,
         max_local_port: int = 9000,
-        max_message_size: int = 100 * 1024 * 1024,  # 100MB default
+        max_message_size: int = 300 * 1024 * 1024,  # 300MB default
     ) -> None:
         self._min_port = min_local_port
         self._max_port = max_local_port
@@ -83,68 +83,79 @@ class ViserProxyManager:
         @app.websocket("/viser/{server_id}")
         async def websocket_proxy(websocket: WebSocket, server_id: str):
             """Proxy WebSocket connections to the appropriate Viser server."""
-            await websocket.accept()
-
-            server = self._server_from_session_hash.get(server_id)
-            if server is None:
-                await websocket.close(code=1008, reason="Not Found")
-                return
-
-            # Determine target WebSocket URL
-            target_ws_url = f"ws://127.0.0.1:{server.get_port()}"
-
-            if not target_ws_url:
-                await websocket.close(code=1008, reason="Not Found")
-                return
-
             try:
-                # Connect to the target WebSocket with increased message size
-                async with websockets.connect(
-                    target_ws_url,
-                    max_size=self._max_message_size,  # Set max message size for receiving
-                ) as ws_target:
-                    # Create tasks for bidirectional communication
-                    async def forward_to_target():
-                        """Forward messages from the client to the target WebSocket."""
-                        try:
-                            while True:
-                                data = await websocket.receive_bytes()
-                                await ws_target.send(data, text=False)
-                        except WebSocketDisconnect:
+                await websocket.accept()
+                
+                server = self._server_from_session_hash.get(server_id)
+                if server is None:
+                    await websocket.close(code=1008, reason="Not Found")
+                    return
+
+                # Determine target WebSocket URL
+                target_ws_url = f"ws://127.0.0.1:{server.get_port()}"
+
+                if not target_ws_url:
+                    await websocket.close(code=1008, reason="Not Found")
+                    return
+
+                try:
+                    # Connect to the target WebSocket with increased message size and timeout
+                    async with websockets.connect(
+                        target_ws_url,
+                        max_size=self._max_message_size,
+                        ping_interval=30,  # Send ping every 30 seconds
+                        ping_timeout=10,   # Wait 10 seconds for pong response
+                        close_timeout=5,   # Wait 5 seconds for close handshake
+                    ) as ws_target:
+                        # Create tasks for bidirectional communication
+                        async def forward_to_target():
+                            """Forward messages from the client to the target WebSocket."""
                             try:
-                                await ws_target.close()
-                            except RuntimeError:
-                                pass
-
-                    async def forward_from_target():
-                        """Forward messages from the target WebSocket to the client."""
-                        try:
-                            while True:
-                                data = await ws_target.recv(decode=False)
-                                await websocket.send_bytes(data)
-                        except websockets.exceptions.ConnectionClosed:
+                                while True:
+                                    data = await websocket.receive_bytes()
+                                    await ws_target.send(data, text=False)
+                            except WebSocketDisconnect:
+                                try:
+                                    await ws_target.close()
+                                except RuntimeError:
+                                    pass
+
+                        async def forward_from_target():
+                            """Forward messages from the target WebSocket to the client."""
                             try:
-                                await websocket.close()
-                            except RuntimeError:
-                                pass
-
-                    # Run both forwarding tasks concurrently
-                    forward_task = asyncio.create_task(forward_to_target())
-                    backward_task = asyncio.create_task(forward_from_target())
-
-                    # Wait for either task to complete (which means a connection was closed)
-                    done, pending = await asyncio.wait(
-                        [forward_task, backward_task],
-                        return_when=asyncio.FIRST_COMPLETED,
-                    )
-
-                    # Cancel the remaining task
-                    for task in pending:
-                        task.cancel()
-
+                                while True:
+                                    data = await ws_target.recv(decode=False)
+                                    await websocket.send_bytes(data)
+                            except websockets.exceptions.ConnectionClosed:
+                                try:
+                                    await websocket.close()
+                                except RuntimeError:
+                                    pass
+
+                        # Run both forwarding tasks concurrently
+                        forward_task = asyncio.create_task(forward_to_target())
+                        backward_task = asyncio.create_task(forward_from_target())
+
+                        # Wait for either task to complete (which means a connection was closed)
+                        done, pending = await asyncio.wait(
+                            [forward_task, backward_task],
+                            return_when=asyncio.FIRST_COMPLETED,
+                        )
+
+                        # Cancel the remaining task
+                        for task in pending:
+                            task.cancel()
+
+                except websockets.exceptions.ConnectionClosedError as e:
+                    print(f"WebSocket connection closed with error: {e}")
+                    await websocket.close(code=1011, reason="Connection to target closed")
+                    
             except Exception as e:
                 print(f"WebSocket proxy error: {e}")
-                await websocket.close(code=1011, reason=str(e))
+                try:
+                    await websocket.close(code=1011, reason=str(e)[:120])  # Limit reason length
+                except:
+                    pass  # Already closed
 
     def start_server(self, server_id: str) -> viser.ViserServer:
         """Start a new Viser server and associate it with the given server ID.