Update server.py

server.py

@@ -34,16 +34,89 @@ class VirtualGPUServer:
         # Active connections and sessions
         self.active_connections: Dict[str, WebSocket] = {}
         self.active_sessions: Dict[str, Dict[str, Any]] = {}
+        self.heartbeat_interval = 5  # seconds
+        self.connection_timeout = 30  # seconds
         
         # Performance monitoring
         self.ops_counter = 0
         self.start_time = time.time()
 
+    def _make_json_serializable(self, obj):
+        """Convert non-JSON-serializable objects to serializable format"""
+        if isinstance(obj, dict):
+            return {k: self._make_json_serializable(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._make_json_serializable(i) for i in obj]
+        elif isinstance(obj, tuple):
+            return list(obj)
+        elif isinstance(obj, (np.ndarray, np.generic)):
+            return obj.tolist()
+        elif isinstance(obj, (Path, uuid.UUID)):
+            return str(obj)
+        elif hasattr(obj, '__dict__'):
+            # Handle custom objects by converting their __dict__ to serializable format
+            return self._make_json_serializable(obj.__dict__)
+        elif isinstance(obj, (int, float, str, bool, type(None))):
+            return obj
+        else:
+            # Convert any other types to string representation
+            return str(obj)
+
+    async def monitor_connection(self, websocket: WebSocket, session_id: str):
+        """Monitor connection health and handle reconnection"""
+        try:
+            while session_id in self.active_connections:
+                try:
+                    await asyncio.wait_for(websocket.receive_text(), timeout=self.heartbeat_interval)
+                except asyncio.TimeoutError:
+                    try:
+                        await websocket.send_json({"type": "ping"})
+                    except:
+                        print(f"Connection lost for session {session_id}")
+                        break
+                except Exception:
+                    break
+                await asyncio.sleep(self.heartbeat_interval)
+        finally:
+            await self.handle_disconnect(session_id)
+
+    async def handle_disconnect(self, session_id: str):
+        """Clean up resources when a client disconnects"""
+        if session_id in self.active_connections:
+            try:
+                await self.active_connections[session_id].close()
+            except:
+                pass
+            del self.active_connections[session_id]
+        if session_id in self.active_sessions:
+            # Save any pending state before removing session
+            session_data = self.active_sessions[session_id]
+            if session_data.get('pending_state'):
+                await self.handle_state_operation({
+                    'type': 'save',
+                    'component': 'session',
+                    'state_id': session_id,
+                    'data': session_data['pending_state']
+                })
+            del self.active_sessions[session_id]
+
     async def handle_vram_operation(self, operation: dict) -> dict:
         """Handle VRAM read/write operations"""
-        op_type = operation.get('type')
-        block_id = operation.get('block_id')
-        data = operation.get('data')
+        try:
+            op_type = operation.get('type')
+            block_id = operation.get('block_id')
+            data = operation.get('data')
+            
+            # Handle non-serializable objects
+            if data and isinstance(data, (dict, list)):
+                data = self._make_json_serializable(data)
+
+            if op_type == 'write':
+                # Additional error handling and logging
+                if not block_id:
+                    raise ValueError("Missing block_id for VRAM write operation")
+                if data is None:
+                    raise ValueError("Missing data for VRAM write operation")
 
         if op_type == 'write':
             # Generate unique file path for this block