feat: Add auto-organize progress tracking and WebSocket broadcasting in BaseModelRoutes and WebSocketManager

py/routes/base_model_routes.py

@@ -54,6 +54,7 @@ class BaseModelRoutes(ABC):
         app.router.add_post(f'/api/{prefix}/move_model', self.move_model)
         app.router.add_post(f'/api/{prefix}/move_models_bulk', self.move_models_bulk)
         app.router.add_get(f'/api/{prefix}/auto-organize', self.auto_organize_models)
+        app.router.add_get(f'/api/{prefix}/auto-organize-progress', self.get_auto_organize_progress)
         
         # Common query routes
         app.router.add_get(f'/api/{prefix}/top-tags', self.get_top_tags)
@@ -749,6 +750,43 @@ class BaseModelRoutes(ABC):
     
     async def auto_organize_models(self, request: web.Request) -> web.Response:
         """Auto-organize all models based on current settings"""
+        try:
+            # Check if auto-organize is already running
+            if ws_manager.is_auto_organize_running():
+                return web.json_response({
+                    'success': False,
+                    'error': 'Auto-organize is already running. Please wait for it to complete.'
+                }, status=409)
+            
+            # Acquire lock to prevent concurrent auto-organize operations
+            auto_organize_lock = await ws_manager.get_auto_organize_lock()
+            
+            if auto_organize_lock.locked():
+                return web.json_response({
+                    'success': False,
+                    'error': 'Auto-organize is already running. Please wait for it to complete.'
+                }, status=409)
+            
+            async with auto_organize_lock:
+                return await self._perform_auto_organize()
+                
+        except Exception as e:
+            logger.error(f"Error in auto_organize_models: {e}", exc_info=True)
+            
+            # Send error message via WebSocket and cleanup
+            await ws_manager.broadcast_auto_organize_progress({
+                'type': 'auto_organize_progress',
+                'status': 'error',
+                'error': str(e)
+            })
+            
+            return web.json_response({
+                'success': False,
+                'error': str(e)
+            }, status=500)
+    
+    async def _perform_auto_organize(self) -> web.Response:
+        """Perform the actual auto-organize operation"""
         try:
             # Get all models from cache
             cache = await self.service.scanner.get_cached_data()
@@ -757,6 +795,11 @@ class BaseModelRoutes(ABC):
             # Get model roots for this scanner
             model_roots = self.service.get_model_roots()
             if not model_roots:
+                await ws_manager.broadcast_auto_organize_progress({
+                    'type': 'auto_organize_progress',
+                    'status': 'error',
+                    'error': 'No model roots configured'
+                })
                 return web.json_response({
                     'success': False,
                     'error': 'No model roots configured'
@@ -775,7 +818,7 @@ class BaseModelRoutes(ABC):
             skipped_count = 0
             
             # Send initial progress via WebSocket
-            await ws_manager.broadcast({
+            await ws_manager.broadcast_auto_organize_progress({
                 'type': 'auto_organize_progress',
                 'status': 'started',
                 'total': total_models,
@@ -906,7 +949,7 @@ class BaseModelRoutes(ABC):
                         processed += 1
                 
                 # Send progress update after each batch
-                await ws_manager.broadcast({
+                await ws_manager.broadcast_auto_organize_progress({
                     'type': 'auto_organize_progress',
                     'status': 'processing',
                     'total': total_models,
@@ -920,7 +963,7 @@ class BaseModelRoutes(ABC):
                 await asyncio.sleep(0.1)
             
             # Send completion message
-            await ws_manager.broadcast({
+            await ws_manager.broadcast_auto_organize_progress({
                 'type': 'auto_organize_progress',
                 'status': 'cleaning',
                 'total': total_models,
@@ -939,7 +982,7 @@ class BaseModelRoutes(ABC):
                 cleanup_counts[root] = removed
                 
             # Send cleanup completed message
-            await ws_manager.broadcast({
+            await ws_manager.broadcast_auto_organize_progress({
                 'type': 'auto_organize_progress',
                 'status': 'completed',
                 'total': total_models,
@@ -974,15 +1017,34 @@ class BaseModelRoutes(ABC):
             return web.json_response(response_data)
             
         except Exception as e:
-            logger.error(f"Error in auto_organize_models: {e}", exc_info=True)
+            logger.error(f"Error in _perform_auto_organize: {e}", exc_info=True)
             
             # Send error message via WebSocket
-            await ws_manager.broadcast({
+            await ws_manager.broadcast_auto_organize_progress({
                 'type': 'auto_organize_progress',
                 'status': 'error',
                 'error': str(e)
             })
             
+            raise e
+    
+    async def get_auto_organize_progress(self, request: web.Request) -> web.Response:
+        """Get current auto-organize progress for polling"""
+        try:
+            progress_data = ws_manager.get_auto_organize_progress()
+            
+            if progress_data is None:
+                return web.json_response({
+                    'success': False,
+                    'error': 'No auto-organize operation in progress'
+                }, status=404)
+            
+            return web.json_response({
+                'success': True,
+                'progress': progress_data
+            })
+        except Exception as e:
+            logger.error(f"Error getting auto-organize progress: {e}", exc_info=True)
             return web.json_response({
                 'success': False,
                 'error': str(e)

py/services/websocket_manager.py

@@ -16,6 +16,9 @@ class WebSocketManager:
         self._download_websockets: Dict[str, web.WebSocketResponse] = {}  # New dict for download-specific clients
         # Add progress tracking dictionary
         self._download_progress: Dict[str, Dict] = {}
+        # Add auto-organize progress tracking
+        self._auto_organize_progress: Optional[Dict] = None
+        self._auto_organize_lock = asyncio.Lock()
         
     async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
         """Handle new WebSocket connection"""
@@ -134,6 +137,33 @@ class WebSocketManager:
         except Exception as e:
             logger.error(f"Error sending download progress: {e}")
             
+    async def broadcast_auto_organize_progress(self, data: Dict):
+        """Broadcast auto-organize progress to connected clients"""
+        # Store progress data in memory
+        self._auto_organize_progress = data
+        
+        # Broadcast via WebSocket
+        await self.broadcast(data)
+    
+    def get_auto_organize_progress(self) -> Optional[Dict]:
+        """Get current auto-organize progress"""
+        return self._auto_organize_progress
+    
+    def cleanup_auto_organize_progress(self):
+        """Clear auto-organize progress data"""
+        self._auto_organize_progress = None
+    
+    def is_auto_organize_running(self) -> bool:
+        """Check if auto-organize is currently running"""
+        if not self._auto_organize_progress:
+            return False
+        status = self._auto_organize_progress.get('status')
+        return status in ['started', 'processing', 'cleaning']
+    
+    async def get_auto_organize_lock(self):
+        """Get the auto-organize lock"""
+        return self._auto_organize_lock
+    
     def get_download_progress(self, download_id: str) -> Optional[Dict]:
         """Get progress information for a specific download"""
         return self._download_progress.get(download_id)