Update server.py

server.py

@@ -17,6 +17,9 @@ app = FastAPI()
 class VirtualGPUServer:
     def __init__(self):
         self.base_path = Path(__file__).parent / "storage"
+        if not self.base_path.exists():
+            self.base_path.mkdir(parents=True)
+            
         self.vram_path = self.base_path / "vram_blocks"
         self.state_path = self.base_path / "gpu_state"
         self.cache_path = self.base_path / "cache"
@@ -66,30 +69,48 @@ class VirtualGPUServer:
         """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:
+                current_time = time.time()
+                session = self.active_sessions.get(session_id)
+                
+                if not session:
                     break
+                    
+                # Check if connection is still alive
+                last_ping = session.get('last_ping', 0)
+                if current_time - last_ping > self.connection_timeout:
+                    if session.get('keep_alive', False):
+                        try:
+                            await websocket.send_json({"type": "ping"})
+                            session['last_ping'] = current_time
+                        except Exception as e:
+                            print(f"Connection lost for session {session_id}: {e}")
+                            break
+                    else:
+                        # Connection timed out and keep-alive is disabled
+                        break
+                        
                 await asyncio.sleep(self.heartbeat_interval)
         finally:
-            await self.handle_disconnect(session_id)
+            # Only disconnect if connection is truly dead
+            if session_id in self.active_connections:
+                try:
+                    if not await self.check_connection_alive(websocket):
+                        await self.handle_disconnect(session_id)
+                except:
+                    await self.handle_disconnect(session_id)
+
+    async def check_connection_alive(self, websocket: WebSocket) -> bool:
+        """Check if a WebSocket connection is still alive"""
+        try:
+            await websocket.send_json({"type": "ping"})
+            return True
+        except Exception:
+            return False
 
     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
+            # Save any pending state before cleanup
             session_data = self.active_sessions[session_id]
             if session_data.get('pending_state'):
                 await self.handle_state_operation({
@@ -361,45 +382,76 @@ async def handle_files():
 # WebSocket endpoint
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
-    session_id = str(uuid.uuid4())
-    server.active_connections[session_id] = websocket
-    server.active_sessions[session_id] = {
-        'start_time': time.time(),
-        'ops_count': 0
-    }
-
+    session_id = None
     try:
+        await websocket.accept()
+        session_id = str(uuid.uuid4())
+        print(f"INFO:     WebSocket connection opened, session: {session_id}")
+        
+        # Initialize session with keep-alive enabled
+        server.active_connections[session_id] = websocket
+        server.active_sessions[session_id] = {
+            'start_time': time.time(),
+            'ops_count': 0,
+            'keep_alive': True,
+            'last_ping': time.time()
+        }
+
         while True:
-            message = await websocket.receive_json()
-            
-            # Route operation to appropriate handler
-            operation_type = message.get('operation')
-            if operation_type == 'vram':
-                response = await server.handle_vram_operation(message)
-            elif operation_type == 'state':
-                response = await server.handle_state_operation(message)
-            elif operation_type == 'cache':
-                response = await server.handle_cache_operation(message)
-            else:
-                response = {
-                    'status': 'error',
-                    'message': 'Unknown operation type'
-                }
-            
-            # Update statistics
-            server.ops_counter += 1
-            server.active_sessions[session_id]['ops_count'] += 1
-            
-            # Send response
-            await websocket.send_json(response)
+            try:
+                # Use a shorter timeout for more responsive connection management
+                message = await websocket.receive_json()
+                # Update last activity timestamp
+                server.active_sessions[session_id]['last_ping'] = time.time()
+                
+                # Handle ping messages
+                if message.get('type') == 'ping':
+                    await websocket.send_json({"type": "pong"})
+                    continue
+
+                # Route operation to appropriate handler
+                operation_type = message.get('operation')
+                if operation_type == 'vram':
+                    response = await server.handle_vram_operation(message)
+                elif operation_type == 'state':
+                    response = await server.handle_state_operation(message)
+                elif operation_type == 'cache':
+                    response = await server.handle_cache_operation(message)
+                else:
+                    response = {
+                        'status': 'error',
+                        'message': 'Unknown operation type'
+                    }
+                
+                # Update statistics
+                server.ops_counter += 1
+                server.active_sessions[session_id]['ops_count'] += 1
                 
+                # Send response
+                await websocket.send_json(response)
+                
+            except asyncio.TimeoutError:
+                # Send ping on timeout
+                try:
+                    await websocket.send_json({"type": "ping"})
+                except:
+                    break  # Connection lost
+                
+    except websockets.exceptions.ConnectionClosed:
+        print(f"INFO:     WebSocket connection closed normally, session: {session_id}")
     except Exception as e:
-        print(f"WebSocket error: {e}")
+        print(f"ERROR:    WebSocket error in session {session_id}: {str(e)}")
     finally:
-        # Cleanup on disconnect
-        del server.active_connections[session_id]
-        del server.active_sessions[session_id]
+        # Cleanup on disconnect, but only if we had a valid session
+        if session_id:
+            try:
+                if session_id in server.active_connections:
+                    del server.active_connections[session_id]
+                if session_id in server.active_sessions:
+                    del server.active_sessions[session_id]
+                print(f"INFO:     Cleaned up session: {session_id}")
+            except Exception as cleanup_error:
+                print(f"WARNING:  Error during session cleanup: {cleanup_error}")
 
 # For running directly (development)
 if __name__ == "__main__":