df

SYS/worker_manager.py

@@ -19,11 +19,17 @@ logger = logging.getLogger(__name__)
 
 class Worker:
     """Represents a single worker task with state management."""
-    
-    def __init__(self, worker_id: str, worker_type: str, title: str = "", 
-                 description: str = "", manager: Optional['WorkerManager'] = None):
+
+    def __init__(
+        self,
+        worker_id: str,
+        worker_type: str,
+        title: str = "",
+        description: str = "",
+        manager: Optional["WorkerManager"] = None,
+    ):
         """Initialize a worker.
-        
+
         Args:
             worker_id: Unique identifier for this worker
             worker_type: Type of work (e.g., 'download', 'search', 'import')
@@ -43,10 +49,10 @@ class Worker:
         self.result = "pending"
         self._stdout_buffer = []
         self._steps_buffer = []
-    
+
     def log_step(self, step_text: str) -> None:
         """Log a step for this worker.
-        
+
         Args:
             step_text: Text describing the step
         """
@@ -57,10 +63,10 @@ class Worker:
                 logger.info(f"[{self.id}] {step_text}")
         except Exception as e:
             logger.error(f"Error logging step for worker {self.id}: {e}")
-    
+
     def append_stdout(self, text: str) -> None:
         """Append text to stdout log.
-        
+
         Args:
             text: Text to append
         """
@@ -71,10 +77,10 @@ class Worker:
                 self._stdout_buffer.append(text)
         except Exception as e:
             logger.error(f"Error appending stdout for worker {self.id}: {e}")
-    
+
     def get_stdout(self) -> str:
         """Get all stdout for this worker.
-        
+
         Returns:
             Complete stdout text
         """
@@ -86,10 +92,10 @@ class Worker:
         except Exception as e:
             logger.error(f"Error getting stdout for worker {self.id}: {e}")
             return ""
-    
+
     def get_steps(self) -> str:
         """Get all steps for this worker.
-        
+
         Returns:
             Complete steps text
         """
@@ -101,10 +107,10 @@ class Worker:
         except Exception as e:
             logger.error(f"Error getting steps for worker {self.id}: {e}")
             return ""
-    
+
     def update_progress(self, progress: str = "", details: str = "") -> None:
         """Update worker progress.
-        
+
         Args:
             progress: Progress string (e.g., "50%")
             details: Additional details
@@ -116,10 +122,10 @@ class Worker:
                 self.manager.update_worker(self.id, progress, details)
         except Exception as e:
             logger.error(f"Error updating worker {self.id}: {e}")
-    
+
     def finish(self, result: str = "completed", message: str = "") -> None:
         """Mark worker as finished.
-        
+
         Args:
             result: Result status ('completed', 'error', 'cancelled')
             message: Result message/error details
@@ -139,12 +145,16 @@ class Worker:
 
 class WorkerLoggingHandler(logging.StreamHandler):
     """Custom logging handler that captures logs for a worker."""
-    
-    def __init__(self, worker_id: str, db: API_folder_store,
-                 manager: Optional['WorkerManager'] = None,
-                 buffer_size: int = 50):
+
+    def __init__(
+        self,
+        worker_id: str,
+        db: API_folder_store,
+        manager: Optional["WorkerManager"] = None,
+        buffer_size: int = 50,
+    ):
         """Initialize the handler.
-        
+
         Args:
             worker_id: ID of the worker to capture logs for
             db: Reference to LocalLibraryDB for storing logs
@@ -157,14 +167,13 @@ class WorkerLoggingHandler(logging.StreamHandler):
         self.buffer_size = buffer_size
         self.buffer = []
         self._lock = Lock()
-        
+
         # Set a format that includes timestamp and level
         formatter = logging.Formatter(
-            '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-            datefmt='%Y-%m-%d %H:%M:%S'
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
         )
         self.setFormatter(formatter)
-    
+
     def emit(self, record):
         """Emit a log record."""
         try:
@@ -183,41 +192,42 @@ class WorkerLoggingHandler(logging.StreamHandler):
                 except (TypeError, ValueError):
                     # If that fails too, just use the raw message string
                     msg = str(record.msg)
-                
+
                 # Add timestamp and level if not already in message
                 import time
-                timestamp = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(record.created))
+
+                timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(record.created))
                 msg = f"{timestamp} - {record.name} - {record.levelname} - {msg}"
-            
+
             with self._lock:
                 self.buffer.append(msg)
-                
+
                 # Flush to DB when buffer reaches size
                 if len(self.buffer) >= self.buffer_size:
                     self._flush()
         except Exception:
             self.handleError(record)
-    
+
     def _flush(self):
         """Flush buffered logs to database."""
         if self.buffer:
-            log_text = '\n'.join(self.buffer)
+            log_text = "\n".join(self.buffer)
             try:
                 if self.manager:
-                    self.manager.append_worker_stdout(self.worker_id, log_text, channel='log')
+                    self.manager.append_worker_stdout(self.worker_id, log_text, channel="log")
                 else:
-                    self.db.append_worker_stdout(self.worker_id, log_text, channel='log')
+                    self.db.append_worker_stdout(self.worker_id, log_text, channel="log")
             except Exception as e:
                 # If we can't write to DB, at least log it
                 log(f"Error flushing worker logs: {e}")
             self.buffer = []
-    
+
     def flush(self):
         """Flush any buffered records."""
         with self._lock:
             self._flush()
         super().flush()
-    
+
     def close(self):
         """Close the handler."""
         self.flush()
@@ -226,10 +236,10 @@ class WorkerLoggingHandler(logging.StreamHandler):
 
 class WorkerManager:
     """Manages persistent worker tasks with auto-refresh capability."""
-    
+
     def __init__(self, library_root: Path, auto_refresh_interval: float = 2.0):
         """Initialize the worker manager.
-        
+
         Args:
             library_root: Root directory for the local library database
             auto_refresh_interval: Seconds between auto-refresh checks (0 = disabled)
@@ -243,7 +253,7 @@ class WorkerManager:
         self._lock = Lock()
         self.worker_handlers: Dict[str, WorkerLoggingHandler] = {}  # Track active handlers
         self._worker_last_step: Dict[str, str] = {}
-    
+
     def close(self) -> None:
         """Close the database connection."""
         if self.db:
@@ -251,18 +261,18 @@ class WorkerManager:
                 self.db.close()
             except Exception:
                 pass
-    
+
     def __enter__(self):
         """Context manager entry."""
         return self
-    
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Context manager exit - close database."""
         self.close()
-    
+
     def add_refresh_callback(self, callback: Callable[[List[Dict[str, Any]]], None]) -> None:
         """Register a callback to be called on worker updates.
-        
+
         Args:
             callback: Function that receives list of active workers
         """
@@ -297,25 +307,25 @@ class WorkerManager:
         except Exception as exc:
             logger.error(f"Failed to expire stale workers: {exc}", exc_info=True)
             return 0
-    
+
     def remove_refresh_callback(self, callback: Callable) -> None:
         """Remove a refresh callback.
-        
+
         Args:
             callback: The callback function to remove
         """
         with self._lock:
             if callback in self.refresh_callbacks:
                 self.refresh_callbacks.remove(callback)
-    
+
     def enable_logging_for_worker(self, worker_id: str) -> Optional[WorkerLoggingHandler]:
         """Enable logging capture for a worker.
-        
+
         Creates a logging handler that captures all logs for this worker.
-        
+
         Args:
             worker_id: ID of the worker to capture logs for
-            
+
         Returns:
             The logging handler that was created, or None if there was an error
         """
@@ -323,46 +333,57 @@ class WorkerManager:
             handler = WorkerLoggingHandler(worker_id, self.db, manager=self)
             with self._lock:
                 self.worker_handlers[worker_id] = handler
-            
+
             # Add the handler to the root logger so it captures all logs
             root_logger = logging.getLogger()
             root_logger.addHandler(handler)
             root_logger.setLevel(logging.DEBUG)  # Capture all levels
-            
+
             logger.debug(f"[WorkerManager] Enabled logging for worker: {worker_id}")
             return handler
         except Exception as e:
-            logger.error(f"[WorkerManager] Error enabling logging for worker {worker_id}: {e}", exc_info=True)
+            logger.error(
+                f"[WorkerManager] Error enabling logging for worker {worker_id}: {e}", exc_info=True
+            )
             return None
-    
+
     def disable_logging_for_worker(self, worker_id: str) -> None:
         """Disable logging capture for a worker and flush any pending logs.
-        
+
         Args:
             worker_id: ID of the worker to stop capturing logs for
         """
         try:
             with self._lock:
                 handler = self.worker_handlers.pop(worker_id, None)
-            
+
             if handler:
                 # Flush and close the handler
                 handler.flush()
                 handler.close()
-                
+
                 # Remove from root logger
                 root_logger = logging.getLogger()
                 root_logger.removeHandler(handler)
-                
+
                 logger.debug(f"[WorkerManager] Disabled logging for worker: {worker_id}")
         except Exception as e:
-            logger.error(f"[WorkerManager] Error disabling logging for worker {worker_id}: {e}", exc_info=True)
-    
-    def track_worker(self, worker_id: str, worker_type: str, title: str = "", 
-                    description: str = "", total_steps: int = 0,
-                    pipe: Optional[str] = None) -> bool:
+            logger.error(
+                f"[WorkerManager] Error disabling logging for worker {worker_id}: {e}",
+                exc_info=True,
+            )
+
+    def track_worker(
+        self,
+        worker_id: str,
+        worker_type: str,
+        title: str = "",
+        description: str = "",
+        total_steps: int = 0,
+        pipe: Optional[str] = None,
+    ) -> bool:
         """Start tracking a new worker.
-        
+
         Args:
             worker_id: Unique identifier for the worker
             worker_type: Type of worker (e.g., 'download', 'search', 'import')
@@ -370,12 +391,14 @@ class WorkerManager:
             description: Worker description
             total_steps: Total number of steps for progress tracking
             pipe: Text of the originating pipe/prompt, if any
-            
+
         Returns:
             True if worker was inserted successfully
         """
         try:
-            result = self.db.insert_worker(worker_id, worker_type, title, description, total_steps, pipe=pipe)
+            result = self.db.insert_worker(
+                worker_id, worker_type, title, description, total_steps, pipe=pipe
+            )
             if result > 0:
                 logger.debug(f"[WorkerManager] Tracking worker: {worker_id} ({worker_type})")
                 self._start_refresh_if_needed()
@@ -384,65 +407,69 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error tracking worker: {e}", exc_info=True)
             return False
-    
-    def update_worker(self, worker_id: str, progress: float = 0.0, current_step: str = "",
-                     details: str = "", error: str = "") -> bool:
+
+    def update_worker(
+        self,
+        worker_id: str,
+        progress: float = 0.0,
+        current_step: str = "",
+        details: str = "",
+        error: str = "",
+    ) -> bool:
         """Update worker progress and status.
-        
+
         Args:
             worker_id: Unique identifier for the worker
             progress: Progress percentage (0-100)
             current_step: Current step description
             details: Additional details
             error: Error message if any
-            
+
         Returns:
             True if update was successful
         """
         try:
             kwargs = {}
             if progress > 0:
-                kwargs['progress'] = progress
+                kwargs["progress"] = progress
             if current_step:
-                kwargs['current_step'] = current_step
+                kwargs["current_step"] = current_step
             if details:
-                kwargs['description'] = details
+                kwargs["description"] = details
             if error:
-                kwargs['error_message'] = error
-            
+                kwargs["error_message"] = error
+
             if kwargs:
-                kwargs['last_updated'] = datetime.now().isoformat()
-                if 'current_step' in kwargs and kwargs['current_step']:
-                    self._worker_last_step[worker_id] = str(kwargs['current_step'])
+                kwargs["last_updated"] = datetime.now().isoformat()
+                if "current_step" in kwargs and kwargs["current_step"]:
+                    self._worker_last_step[worker_id] = str(kwargs["current_step"])
                 return self.db.update_worker(worker_id, **kwargs)
             return True
         except Exception as e:
             logger.error(f"[WorkerManager] Error updating worker {worker_id}: {e}", exc_info=True)
             return False
-    
-    def finish_worker(self, worker_id: str, result: str = "completed", 
-                     error_msg: str = "", result_data: str = "") -> bool:
+
+    def finish_worker(
+        self, worker_id: str, result: str = "completed", error_msg: str = "", result_data: str = ""
+    ) -> bool:
         """Mark a worker as finished.
-        
+
         Args:
             worker_id: Unique identifier for the worker
             result: Result status ('completed', 'error', 'cancelled')
             error_msg: Error message if any
             result_data: Result data as JSON string
-            
+
         Returns:
             True if update was successful
         """
         try:
-            kwargs = {
-                'status': result,
-                'completed_at': datetime.now().isoformat()
-            }
+            kwargs = {"status": result, "completed_at": datetime.now().isoformat()}
             if error_msg:
-                kwargs['error_message'] = error_msg
+                kwargs["error_message"] = error_msg
             if result_data:
-                kwargs['result_data'] = result_data
-            
+                kwargs["result_data"] = result_data
+
             success = self.db.update_worker(worker_id, **kwargs)
             logger.info(f"[WorkerManager] Worker finished: {worker_id} ({result})")
             self._worker_last_step.pop(worker_id, None)
@@ -450,10 +477,10 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error finishing worker {worker_id}: {e}", exc_info=True)
             return False
-    
+
     def get_active_workers(self) -> List[Dict[str, Any]]:
         """Get all active (running) workers.
-        
+
         Returns:
             List of active worker dictionaries
         """
@@ -462,31 +489,33 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error getting active workers: {e}", exc_info=True)
             return []
-    
+
     def get_finished_workers(self, limit: int = 100) -> List[Dict[str, Any]]:
         """Get all finished workers (completed, errored, or cancelled).
-        
+
         Args:
             limit: Maximum number of workers to retrieve
-        
+
         Returns:
             List of finished worker dictionaries
         """
         try:
             all_workers = self.db.get_all_workers(limit=limit)
             # Filter to only finished workers
-            finished = [w for w in all_workers if w.get('status') in ['completed', 'error', 'cancelled']]
+            finished = [
+                w for w in all_workers if w.get("status") in ["completed", "error", "cancelled"]
+            ]
             return finished
         except Exception as e:
             logger.error(f"[WorkerManager] Error getting finished workers: {e}", exc_info=True)
             return []
-    
+
     def get_worker(self, worker_id: str) -> Optional[Dict[str, Any]]:
         """Get a specific worker's data.
-        
+
         Args:
             worker_id: Unique identifier for the worker
-            
+
         Returns:
             Worker data or None if not found
         """
@@ -495,18 +524,18 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error getting worker {worker_id}: {e}", exc_info=True)
             return None
-    
+
     def get_worker_events(self, worker_id: str, limit: int = 500) -> List[Dict[str, Any]]:
         """Fetch recorded worker timeline events."""
         return self.db.get_worker_events(worker_id, limit)
-    
+
     def log_step(self, worker_id: str, step_text: str) -> bool:
         """Log a step to a worker's step history.
-        
+
         Args:
             worker_id: Unique identifier for the worker
             step_text: Step description to log
-            
+
         Returns:
             True if successful
         """
@@ -516,43 +545,49 @@ class WorkerManager:
                 self._worker_last_step[worker_id] = step_text
             return success
         except Exception as e:
-            logger.error(f"[WorkerManager] Error logging step for worker {worker_id}: {e}", exc_info=True)
+            logger.error(
+                f"[WorkerManager] Error logging step for worker {worker_id}: {e}", exc_info=True
+            )
             return False
-    
+
     def _get_last_step(self, worker_id: str) -> Optional[str]:
         """Return the most recent step description for a worker."""
         return self._worker_last_step.get(worker_id)
-    
+
     def get_steps(self, worker_id: str) -> str:
         """Get step logs for a worker.
-        
+
         Args:
             worker_id: Unique identifier for the worker
-            
+
         Returns:
             Steps text or empty string if not found
         """
         try:
             return self.db.get_worker_steps(worker_id)
         except Exception as e:
-            logger.error(f"[WorkerManager] Error getting steps for worker {worker_id}: {e}", exc_info=True)
-            return ''
-    
+            logger.error(
+                f"[WorkerManager] Error getting steps for worker {worker_id}: {e}", exc_info=True
+            )
+            return ""
+
     def start_auto_refresh(self) -> None:
         """Start the auto-refresh thread for periodic worker updates."""
         if self.auto_refresh_interval <= 0:
             logger.debug("[WorkerManager] Auto-refresh disabled (interval <= 0)")
             return
-        
+
         if self.refresh_thread and self.refresh_thread.is_alive():
             logger.debug("[WorkerManager] Auto-refresh already running")
             return
-        
-        logger.info(f"[WorkerManager] Starting auto-refresh with {self.auto_refresh_interval}s interval")
+
+        logger.info(
+            f"[WorkerManager] Starting auto-refresh with {self.auto_refresh_interval}s interval"
+        )
         self._stop_refresh = False
         self.refresh_thread = Thread(target=self._auto_refresh_loop, daemon=True)
         self.refresh_thread.start()
-    
+
     def stop_auto_refresh(self) -> None:
         """Stop the auto-refresh thread."""
         logger.info("[WorkerManager] Stopping auto-refresh")
@@ -560,46 +595,48 @@ class WorkerManager:
         if self.refresh_thread:
             self.refresh_thread.join(timeout=5)
             self.refresh_thread = None
-    
+
     def _start_refresh_if_needed(self) -> None:
         """Start auto-refresh if we have active workers and callbacks."""
         active = self.get_active_workers()
         if active and self.refresh_callbacks and not self._stop_refresh:
             self.start_auto_refresh()
-    
+
     def _auto_refresh_loop(self) -> None:
         """Main auto-refresh loop that periodically queries and notifies."""
         try:
             while not self._stop_refresh:
                 time.sleep(self.auto_refresh_interval)
-                
+
                 # Check if there are active workers
                 active = self.get_active_workers()
-                
+
                 if not active:
                     # No more active workers, stop refreshing
                     logger.debug("[WorkerManager] No active workers, stopping auto-refresh")
                     break
-                
+
                 # Call all registered callbacks with the active workers
                 with self._lock:
                     for callback in self.refresh_callbacks:
                         try:
                             callback(active)
                         except Exception as e:
-                            logger.error(f"[WorkerManager] Error in refresh callback: {e}", exc_info=True)
-        
+                            logger.error(
+                                f"[WorkerManager] Error in refresh callback: {e}", exc_info=True
+                            )
+
         except Exception as e:
             logger.error(f"[WorkerManager] Error in auto-refresh loop: {e}", exc_info=True)
         finally:
             logger.debug("[WorkerManager] Auto-refresh loop ended")
-    
+
     def cleanup_old_workers(self, days: int = 7) -> int:
         """Clean up completed/errored workers older than specified days.
-        
+
         Args:
             days: Delete workers completed more than this many days ago
-            
+
         Returns:
             Number of workers deleted
         """
@@ -611,15 +648,15 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error cleaning up old workers: {e}", exc_info=True)
             return 0
-    
+
     def append_stdout(self, worker_id: str, text: str, channel: str = "stdout") -> bool:
         """Append text to a worker's stdout log.
-        
+
         Args:
             worker_id: Unique identifier for the worker
             text: Text to append
             channel: Logical channel (stdout, stderr, log, etc.)
-            
+
         Returns:
             True if append was successful
         """
@@ -629,13 +666,13 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error appending stdout: {e}", exc_info=True)
             return False
-    
+
     def get_stdout(self, worker_id: str) -> str:
         """Get stdout logs for a worker.
-        
+
         Args:
             worker_id: Unique identifier for the worker
-            
+
         Returns:
             Worker's stdout or empty string
         """
@@ -644,17 +681,17 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error getting stdout: {e}", exc_info=True)
             return ""
-    
+
     def append_worker_stdout(self, worker_id: str, text: str, channel: str = "stdout") -> bool:
         """Compatibility wrapper for append_stdout."""
         return self.append_stdout(worker_id, text, channel=channel)
-    
+
     def clear_stdout(self, worker_id: str) -> bool:
         """Clear stdout logs for a worker.
-        
+
         Args:
             worker_id: Unique identifier for the worker
-            
+
         Returns:
             True if clear was successful
         """
@@ -663,7 +700,7 @@ class WorkerManager:
         except Exception as e:
             logger.error(f"[WorkerManager] Error clearing stdout: {e}", exc_info=True)
             return False
-    
+
     def close(self) -> None:
         """Close the worker manager and database connection."""
         self.stop_auto_refresh()