diff --git a/__init__.py b/__init__.py index ab717be..54a5f34 100644 --- a/__init__.py +++ b/__init__.py @@ -3,10 +3,162 @@ ComfyUI Queue Manager Custom Node A comprehensive queue management system for ComfyUI workflows. """ -from .queue_manager_node import QueueManagerNode +import mimetypes +from pathlib import Path +import shutil +import sys -NODE_CLASS_MAPPINGS = {"QueueManagerNode": QueueManagerNode} +# Add the current directory to Python path for imports +current_dir = Path(__file__).parent +if str(current_dir) not in sys.path: + sys.path.insert(0, str(current_dir)) -NODE_DISPLAY_NAME_MAPPINGS = {"QueueManagerNode": "Queue Manager"} +# Import ComfyUI modules with error handling +try: + from aiohttp import web as aiohttp_web + import folder_paths + import server +except ImportError: + # ComfyUI modules not available, will handle gracefully + folder_paths = None + server = None + aiohttp_web = None -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +from queue_manager_node import QueueManagerNode + +# ComfyUI custom node registration +NODE_CLASS_MAPPINGS = { + "QueueManagerNode": QueueManagerNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "QueueManagerNode": "Queue Manager", +} + +# Web directory for serving static files +WEB_DIRECTORY = "./web" + +# ComfyUI Web Extensions +WEB_EXTENSIONS = [ + {"name": "queue_manager_extension", "path": "queue_manager_extension.js"} +] + +# Extension metadata +__version__ = "1.0.0" +__author__ = "AC ComfyUI Queue Manager" +__description__ = "A comprehensive queue management system for ComfyUI workflows" + + +def register_menu_extension(): + """Register the Queue Manager menu extension with ComfyUI.""" + if not server or not aiohttp_web: + print("[Queue Manager] ComfyUI server modules not available") + return False + + try: + # Register with ComfyUI's web extension system + if not hasattr(server, "PromptServer"): + return False + + prompt_server = server.PromptServer.instance + if not prompt_server or not hasattr(prompt_server, "app"): + return False + + async def serve_extension_file(request): + """Serve extension files.""" + filename = request.match_info.get("filename", "") + file_path = current_dir / "web" / filename + + if not file_path.exists(): + return aiohttp_web.Response(status=404, text="File not found") + + # Determine content type + content_type, _ = mimetypes.guess_type(str(file_path)) + if not content_type: + if filename.endswith(".js"): + content_type = "application/javascript" + elif filename.endswith(".html"): + content_type = "text/html" + elif filename.endswith(".css"): + content_type = "text/css" + else: + content_type = "text/plain" + + try: + with file_path.open(encoding="utf-8") as f: + content = f.read() + return aiohttp_web.Response(text=content, content_type=content_type) + except OSError as e: + return aiohttp_web.Response( + status=500, text=f"Error reading file: {e}" + ) + + # Register routes + prompt_server.app.router.add_get( + "/extensions/comfyui-queue-manager/{filename}", serve_extension_file + ) + + # Copy extension to ComfyUI's web extensions directory + if folder_paths: + try: + web_extensions_dir = folder_paths.get_folder_paths("web_extensions") + if web_extensions_dir: + extension_src = current_dir / "web" / "queue_manager_extension.js" + for ext_dir in web_extensions_dir: + ext_dest = Path(ext_dir) / "queue_manager_extension.js" + if extension_src.exists(): + shutil.copy2(extension_src, ext_dest) + print(f"[Queue Manager] Extension copied to {ext_dest}") + except OSError as e: + print(f"[Queue Manager] Could not copy extension: {e}") + + print("[Queue Manager] Menu extension routes registered") + return True + + except Exception as e: + print(f"[Queue Manager] Failed to register menu extension: {e}") + return False + + +def initialize_queue_manager(): + """Initialize the queue manager system.""" + try: + # Import and initialize core services + from api_routes import setup_api_routes + from database import SQLiteDatabase + from queue_service import QueueService + + # Initialize database + db_path = current_dir / "queue_manager.db" + database = SQLiteDatabase(str(db_path)) + database.initialize() + + # Initialize queue service + queue_service = QueueService(database) + + # Set up API routes (will be used by the web interface) + setup_api_routes(queue_service) + + print(f"[Queue Manager] Initialized successfully (v{__version__})") + return True + + except Exception as e: + print(f"[Queue Manager] Failed to initialize: {e}") + return False + + +# Initialize on import +_initialized = initialize_queue_manager() +_menu_registered = register_menu_extension() + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY", + "WEB_EXTENSIONS", + "__author__", + "__description__", + "__version__", + "initialize_queue_manager", + "register_menu_extension", +] diff --git a/execution_monitor.py b/execution_monitor.py new file mode 100644 index 0000000..c31d44e --- /dev/null +++ b/execution_monitor.py @@ -0,0 +1,402 @@ +""" +Execution monitoring system for ComfyUI Queue Manager. +Monitors workflow execution status and provides real-time updates. +""" + +from __future__ import annotations + +import logging +import threading +import time +from datetime import datetime, timezone, timedelta +from typing import Any, Callable, Dict, List, Optional + +from models import QueueStatus + +logger = logging.getLogger(__name__) + + +class ExecutionMonitor: + """ + Monitors workflow execution status and provides real-time updates. + + This class tracks running workflows, monitors their progress, and provides + callbacks for status changes and completion events. + """ + + def __init__(self, queue_service=None, workflow_executor=None): + """Initialize the execution monitor. + + Args: + queue_service: Optional queue service for status updates + workflow_executor: Optional workflow executor for execution monitoring + """ + self.queue_service = queue_service + self.workflow_executor = workflow_executor + + self._monitoring_thread: Optional[threading.Thread] = None + self._stop_monitoring = threading.Event() + self._lock = threading.RLock() + + # Monitoring state + self._monitored_workflows: Dict[str, Dict[str, Any]] = {} + self._status_callbacks: Dict[str, Callable] = {} + self._monitoring_interval = 1.0 # seconds + self._is_monitoring = False + + logger.info("Execution monitor initialized") + + def start_monitoring(self) -> bool: + """Start the execution monitoring thread. + + Returns: + True if monitoring was started successfully + """ + with self._lock: + if self._is_monitoring: + logger.warning("Execution monitoring is already running") + return True + + try: + self._stop_monitoring.clear() + self._monitoring_thread = threading.Thread( + target=self._monitoring_loop, + name="ExecutionMonitor", + daemon=True + ) + self._monitoring_thread.start() + self._is_monitoring = True + + logger.info("Execution monitoring started") + return True + + except Exception as e: + logger.error(f"Failed to start execution monitoring: {e}") + return False + + def stop_monitoring(self) -> bool: + """Stop the execution monitoring thread. + + Returns: + True if monitoring was stopped successfully + """ + with self._lock: + if not self._is_monitoring: + return True + + try: + self._stop_monitoring.set() + + if self._monitoring_thread and self._monitoring_thread.is_alive(): + self._monitoring_thread.join(timeout=5.0) + + if self._monitoring_thread.is_alive(): + logger.warning("Monitoring thread did not stop gracefully") + + self._is_monitoring = False + self._monitoring_thread = None + + logger.info("Execution monitoring stopped") + return True + + except Exception as e: + logger.error(f"Failed to stop execution monitoring: {e}") + return False + + def is_monitoring(self) -> bool: + """Check if execution monitoring is active.""" + return self._is_monitoring + + def set_monitoring_interval(self, interval: float) -> None: + """Set the monitoring interval in seconds. + + Args: + interval: Monitoring interval in seconds + """ + if interval <= 0: + raise ValueError("Monitoring interval must be positive") + + self._monitoring_interval = interval + logger.info(f"Monitoring interval set to {interval} seconds") + + def add_workflow_to_monitor(self, workflow_id: str, workflow_info: Dict[str, Any]) -> None: + """Add a workflow to the monitoring system. + + Args: + workflow_id: Unique identifier for the workflow + workflow_info: Information about the workflow + """ + with self._lock: + # Preserve existing added_at if provided, otherwise use current time + added_at = workflow_info.get("added_at", datetime.now(timezone.utc)) + + self._monitored_workflows[workflow_id] = { + **workflow_info, + "added_at": added_at, + "last_checked": None, + "check_count": 0, + } + + logger.debug(f"Added workflow to monitor: {workflow_id}") + + def remove_workflow_from_monitor(self, workflow_id: str) -> bool: + """Remove a workflow from the monitoring system. + + Args: + workflow_id: Identifier of the workflow to remove + + Returns: + True if the workflow was removed + """ + with self._lock: + removed = self._monitored_workflows.pop(workflow_id, None) is not None + + if removed: + logger.debug(f"Removed workflow from monitor: {workflow_id}") + + return removed + + def register_status_callback(self, callback_id: str, callback: Callable) -> None: + """Register a callback for status change events. + + Args: + callback_id: Unique identifier for the callback + callback: Function to call on status changes + """ + self._status_callbacks[callback_id] = callback + logger.debug(f"Registered status callback: {callback_id}") + + def unregister_status_callback(self, callback_id: str) -> None: + """Unregister a status callback. + + Args: + callback_id: Identifier of the callback to remove + """ + self._status_callbacks.pop(callback_id, None) + logger.debug(f"Unregistered status callback: {callback_id}") + + def get_monitored_workflows(self) -> Dict[str, Dict[str, Any]]: + """Get information about all monitored workflows. + + Returns: + Dictionary mapping workflow IDs to their information + """ + with self._lock: + return {wf_id: info.copy() for wf_id, info in self._monitored_workflows.items()} + + def get_workflow_status(self, workflow_id: str) -> QueueStatus | None: + """Get the current status of a monitored workflow. + + Args: + workflow_id: The workflow ID + + Returns: + Current status or None if not monitored + """ + with self._lock: + if workflow_id not in self._monitored_workflows: + return None + + workflow_info = self._monitored_workflows[workflow_id] + return workflow_info.get("status") + + def _monitoring_loop(self) -> None: + """Main monitoring loop that runs in a separate thread.""" + logger.info("Execution monitoring loop started") + + try: + while not self._stop_monitoring.is_set(): + try: + self._check_workflow_statuses() + + # Wait for the monitoring interval or stop signal + if self._stop_monitoring.wait(timeout=self._monitoring_interval): + break + + except Exception as e: + logger.error(f"Error in monitoring loop: {e}") + # Continue monitoring even if there's an error + time.sleep(1.0) + + except Exception as e: + logger.error(f"Fatal error in monitoring loop: {e}") + + finally: + logger.info("Execution monitoring loop stopped") + + def _check_workflow_statuses(self) -> None: + """Check the status of all monitored workflows.""" + with self._lock: + workflows_to_check = list(self._monitored_workflows.items()) + + for workflow_id, workflow_info in workflows_to_check: + try: + self._check_single_workflow_status(workflow_id, workflow_info) + except Exception as e: + logger.error(f"Error checking workflow {workflow_id}: {e}") + + def _check_single_workflow_status(self, workflow_id: str, workflow_info: Dict[str, Any]) -> None: + """Check the status of a single workflow. + + Args: + workflow_id: The workflow ID + workflow_info: Current workflow information + """ + old_status = workflow_info.get("status") + new_status = None + + # Update check information + now = datetime.now(timezone.utc) + workflow_info["last_checked"] = now + workflow_info["check_count"] = workflow_info.get("check_count", 0) + 1 + + # Check status from workflow executor if available + if self.workflow_executor: + try: + new_status = self.workflow_executor.get_workflow_status(workflow_id) + except Exception as e: + logger.error(f"Failed to get workflow status from executor: {e}") + + # Check status from queue service if available and executor didn't provide status + if new_status is None and self.queue_service: + try: + queue_item_id = workflow_info.get("queue_item_id") + if queue_item_id: + queue_item = self.queue_service.get_queue_item(queue_item_id) + if queue_item: + new_status = queue_item.status + except Exception as e: + logger.error(f"Failed to get workflow status from queue service: {e}") + + # Update status if it changed + if new_status and new_status != old_status: + workflow_info["status"] = new_status + workflow_info["status_changed_at"] = now + + # Notify callbacks + self._notify_status_callbacks(workflow_id, old_status, new_status, workflow_info) + + # Remove completed/failed workflows from monitoring after a delay + if new_status in (QueueStatus.COMPLETED, QueueStatus.FAILED): + self._schedule_workflow_removal(workflow_id) + + def _notify_status_callbacks( + self, + workflow_id: str, + old_status: QueueStatus | None, + new_status: QueueStatus, + workflow_info: Dict[str, Any] + ) -> None: + """Notify all registered callbacks about a status change. + + Args: + workflow_id: The workflow ID + old_status: Previous status + new_status: New status + workflow_info: Current workflow information + """ + callback_data = { + "workflow_id": workflow_id, + "old_status": old_status.value if old_status else None, + "new_status": new_status.value, + "workflow_info": workflow_info.copy(), + "timestamp": datetime.now(timezone.utc).isoformat(), + } + + for callback_id, callback in list(self._status_callbacks.items()): + try: + callback(callback_data) + except Exception as e: + logger.error(f"Error in status callback {callback_id}: {e}") + + def _schedule_workflow_removal(self, workflow_id: str, delay_seconds: float = 60.0) -> None: + """Schedule a workflow for removal from monitoring after a delay. + + Args: + workflow_id: The workflow ID + delay_seconds: Delay before removal in seconds + """ + def remove_after_delay(): + time.sleep(delay_seconds) + self.remove_workflow_from_monitor(workflow_id) + + removal_thread = threading.Thread( + target=remove_after_delay, + name=f"RemoveWorkflow-{workflow_id}", + daemon=True + ) + removal_thread.start() + + def get_monitoring_statistics(self) -> Dict[str, Any]: + """Get statistics about the monitoring system. + + Returns: + Dictionary containing monitoring statistics + """ + with self._lock: + stats = { + "is_monitoring": self._is_monitoring, + "monitoring_interval": self._monitoring_interval, + "monitored_workflows": len(self._monitored_workflows), + "registered_callbacks": len(self._status_callbacks), + "status_counts": {}, + } + + # Count workflows by status + for workflow_info in self._monitored_workflows.values(): + status = workflow_info.get("status") + if status: + status_value = status.value if hasattr(status, 'value') else str(status) + stats["status_counts"][status_value] = stats["status_counts"].get(status_value, 0) + 1 + + return stats + + def cleanup_old_workflows(self, max_age_hours: int = 24) -> int: + """Clean up old workflow monitoring records. + + Args: + max_age_hours: Maximum age in hours for workflow records + + Returns: + Number of workflows cleaned up + """ + cutoff_time = datetime.now(timezone.utc) - timedelta(hours=max_age_hours) + cleaned_count = 0 + + with self._lock: + workflows_to_remove = [] + + for workflow_id, workflow_info in self._monitored_workflows.items(): + added_at = workflow_info.get("added_at") + if added_at and added_at < cutoff_time: + # Only remove completed or failed workflows + status = workflow_info.get("status") + if status in (QueueStatus.COMPLETED, QueueStatus.FAILED): + workflows_to_remove.append(workflow_id) + + for workflow_id in workflows_to_remove: + del self._monitored_workflows[workflow_id] + cleaned_count += 1 + + if cleaned_count > 0: + logger.info(f"Cleaned up {cleaned_count} old workflow monitoring records") + + return cleaned_count + + def close(self) -> None: + """Close the execution monitor and clean up resources.""" + self.stop_monitoring() + + with self._lock: + self._monitored_workflows.clear() + self._status_callbacks.clear() + + logger.info("Execution monitor closed") + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.close() \ No newline at end of file diff --git a/queue_manager_node.py b/queue_manager_node.py index 2c76881..ee4cd12 100644 --- a/queue_manager_node.py +++ b/queue_manager_node.py @@ -4,7 +4,12 @@ Main ComfyUI custom node for the Queue Manager. from __future__ import annotations -from typing import Any +import json +import uuid +from datetime import datetime +from typing import Any, Optional + +from models import QueueStatus class QueueManagerNode: @@ -19,47 +24,281 @@ class QueueManagerNode: return { "required": {}, "optional": { - "workflow_name": ("STRING", {"default": ""}), + "workflow_name": ("STRING", {"default": "Workflow"}), "auto_queue": ("BOOLEAN", {"default": True}), + "priority": ("INT", {"default": 0, "min": -10, "max": 10}), + "tags": ("STRING", {"default": ""}), }, } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("queue_item_id",) + RETURN_TYPES = ("STRING", "BOOLEAN") + RETURN_NAMES = ("queue_item_id", "queued_successfully") FUNCTION = "process" CATEGORY = "Queue Management" DESCRIPTION = ( - "Manages workflow execution queues with persistence and control features" + "Manages workflow execution queues with persistence and control features. " + "Automatically captures and queues workflows for later execution." ) - + + # ComfyUI node metadata + OUTPUT_NODE = False + def __init__(self): """Initialize the queue manager node.""" - self.queue_service = None # Will be initialized when needed + self._queue_service = None + self._workflow_executor = None + self._execution_monitor = None - def process(self, workflow_name: str = "", auto_queue: bool = True) -> tuple[str]: + @property + def queue_service(self): + """Lazy initialization of queue service.""" + if self._queue_service is None: + try: + from database import SQLiteDatabase + from queue_service import QueueService + from pathlib import Path + + # Initialize database and queue service + db_path = Path(__file__).parent / "queue_manager.db" + database = SQLiteDatabase(str(db_path)) + database.initialize() + self._queue_service = QueueService(database) + + # Initialize workflow execution components + self._initialize_workflow_execution() + + except Exception as e: + print(f"[Queue Manager] Failed to initialize queue service: {e}") + # Don't set to None, leave it as None to indicate failure + return None + return self._queue_service + + def _initialize_workflow_execution(self): + """Initialize workflow execution integration components.""" + try: + from workflow_executor import WorkflowExecutor + from execution_monitor import ExecutionMonitor + from workflow_interceptor import get_workflow_interceptor + + # Initialize workflow executor + self._workflow_executor = WorkflowExecutor(self._queue_service) + + # Initialize execution monitor + self._execution_monitor = ExecutionMonitor( + queue_service=self._queue_service, + workflow_executor=self._workflow_executor + ) + + # Set up workflow interceptor + interceptor = get_workflow_interceptor() + interceptor.set_workflow_executor(self._workflow_executor) + + # Start execution monitoring + self._execution_monitor.start_monitoring() + + # Enable workflow interception + interceptor.enable() + + print("[Queue Manager] Workflow execution integration initialized") + + except Exception as e: + print(f"[Queue Manager] Failed to initialize workflow execution: {e}") + # Continue without workflow execution integration + + def process( + self, + workflow_name: str = "Workflow", + auto_queue: bool = True, + priority: int = 0, + tags: str = "" + ) -> tuple[str, bool]: """ - Process the node execution. + Process the node execution and optionally add to queue. Args: - workflow_name: Optional name for the workflow + workflow_name: Name for the workflow auto_queue: Whether to automatically add to queue + priority: Priority level for queue processing + tags: Comma-separated tags for the workflow Returns: - Tuple containing the queue item ID + Tuple containing (queue_item_id, queued_successfully) """ - # For now, return a placeholder ID - # This will be implemented in later tasks when the queue service is available - queue_item_id = "placeholder_id" - - return (queue_item_id,) + queue_item_id = str(uuid.uuid4()) + queued_successfully = False + + if auto_queue and self.queue_service: + try: + # Create workflow data structure + workflow_data = { + "name": workflow_name, + "priority": priority, + "tags": [tag.strip() for tag in tags.split(",") if tag.strip()], + "created_at": datetime.now().isoformat(), + "node_id": queue_item_id, + } + + # Add to queue + actual_id = self.queue_service.add_workflow(workflow_data) + if actual_id: + queue_item_id = actual_id + queued_successfully = True + print(f"[Queue Manager] Added workflow '{workflow_name}' to queue: {queue_item_id}") + else: + print(f"[Queue Manager] Failed to add workflow '{workflow_name}' to queue") + + except Exception as e: + print(f"[Queue Manager] Error adding workflow to queue: {e}") + + elif not auto_queue: + print(f"[Queue Manager] Auto-queue disabled for workflow '{workflow_name}'") + + return (queue_item_id, queued_successfully) @classmethod - def IS_CHANGED(cls, **kwargs: Any) -> float: # noqa: N802, ARG003 - """Determine if the node needs to be re-executed.""" - # Always execute for now + def IS_CHANGED(cls, **kwargs: Any) -> float: # noqa: N802 + """ + Determine if the node needs to be re-executed. + + For queue management, we want to execute every time to capture workflows. + """ + # Return NaN to always execute return float("nan") @classmethod - def VALIDATE_INPUTS(cls, **kwargs: Any) -> bool: # noqa: N802, ARG003 + def VALIDATE_INPUTS(cls, **kwargs: Any) -> bool: # noqa: N802 """Validate the inputs to the node.""" + # Basic validation + workflow_name = kwargs.get("workflow_name", "") + priority = kwargs.get("priority", 0) + + # Validate workflow name is not empty if provided + if workflow_name and not isinstance(workflow_name, str): + return False + + # Validate priority is within range + if not isinstance(priority, int) or priority < -10 or priority > 10: + return False + return True + + def get_queue_status(self) -> dict[str, Any]: + """ + Get current queue status information. + + Returns: + Dictionary containing queue status information + """ + if not self.queue_service: + return {"error": "Queue service not available"} + + try: + items = self.queue_service.get_queue_items() + status_counts = {} + + for item in items: + status = item.status.value + status_counts[status] = status_counts.get(status, 0) + 1 + + # Check if pause functionality is available + is_paused = False + if hasattr(self.queue_service, 'is_paused'): + is_paused = self.queue_service.is_paused() + + return { + "total_items": len(items), + "status_counts": status_counts, + "is_paused": is_paused, + } + except Exception as e: + return {"error": f"Failed to get queue status: {e}"} + + def get_workflow_executor(self): + """Get the workflow executor instance.""" + # Ensure queue service is initialized + _ = self.queue_service + return self._workflow_executor + + def get_execution_monitor(self): + """Get the execution monitor instance.""" + # Ensure queue service is initialized + _ = self.queue_service + return self._execution_monitor + + def execute_workflow_directly(self, workflow_data: dict[str, Any]) -> dict[str, Any]: + """Execute a workflow directly through the workflow executor. + + Args: + workflow_data: The workflow data to execute + + Returns: + Dictionary containing execution results + """ + if not self._workflow_executor: + raise RuntimeError("Workflow executor not initialized") + + try: + return self._workflow_executor.execute_workflow(workflow_data) + except Exception as e: + print(f"[Queue Manager] Failed to execute workflow: {e}") + raise + + def get_execution_statistics(self) -> dict[str, Any]: + """Get statistics about workflow executions. + + Returns: + Dictionary containing execution statistics + """ + stats = {} + + if self._workflow_executor: + stats["executor"] = self._workflow_executor.get_execution_statistics() + + if self._execution_monitor: + stats["monitor"] = self._execution_monitor.get_monitoring_statistics() + + return stats + + def cleanup_execution_resources(self) -> dict[str, int]: + """Clean up old execution resources. + + Returns: + Dictionary containing cleanup counts + """ + cleanup_counts = {} + + if self._workflow_executor: + try: + count = self._workflow_executor.cleanup_completed_workflows() + cleanup_counts["executor_workflows"] = count + except Exception as e: + print(f"[Queue Manager] Failed to cleanup executor workflows: {e}") + + if self._execution_monitor: + try: + count = self._execution_monitor.cleanup_old_workflows() + cleanup_counts["monitor_workflows"] = count + except Exception as e: + print(f"[Queue Manager] Failed to cleanup monitor workflows: {e}") + + return cleanup_counts + + @classmethod + def get_node_info(cls) -> dict[str, Any]: + """ + Get information about this node for ComfyUI. + + Returns: + Dictionary containing node information + """ + return { + "name": "QueueManagerNode", + "display_name": "Queue Manager", + "category": "Queue Management", + "description": cls.DESCRIPTION, + "version": "1.0.0", + "author": "ComfyUI Queue Manager", + "input_types": cls.INPUT_TYPES(), + "return_types": cls.RETURN_TYPES, + "return_names": cls.RETURN_NAMES, + } diff --git a/test_menu_functionality.py b/test_menu_functionality.py new file mode 100644 index 0000000..f868225 --- /dev/null +++ b/test_menu_functionality.py @@ -0,0 +1,203 @@ +""" +Functional tests for ComfyUI Queue Manager menu integration. +Tests the actual menu functionality and integration points. +""" + +from pathlib import Path +import sys +import unittest + +# Add the current directory to Python path for imports +current_dir = Path(__file__).parent +sys.path.insert(0, str(current_dir)) + + +class TestMenuFunctionality(unittest.TestCase): + """Test cases for menu functionality.""" + + def setUp(self): + """Set up test fixtures.""" + self.web_dir = current_dir / "web" + + def test_menu_extension_loads_without_errors(self): + """Test that the menu extension can be loaded without syntax errors.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Basic syntax validation + assert "SyntaxError" not in content + assert "app.registerExtension" in content + + # Check for proper JavaScript structure + brace_count = content.count("{") - content.count("}") + assert brace_count == 0, "Braces should be balanced" + + paren_count = content.count("(") - content.count(")") + assert paren_count == 0, "Parentheses should be balanced" + + def test_menu_button_configuration(self): + """Test menu button configuration and styling.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for menu button configuration + assert "Queue Manager" in content + assert "textContent" in content + assert "addEventListener" in content + assert "click" in content + + def test_window_management_configuration(self): + """Test window management configuration.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for window management features + assert "window.open" in content + assert "popup" in content + assert "tab" in content + assert "focus" in content + + def test_url_generation(self): + """Test URL generation for queue manager interface.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for proper URL generation + assert "/extensions/comfyui-queue-manager/" in content + assert "index.html" in content + + def test_error_handling_implementation(self): + """Test error handling implementation.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for error handling + assert "try" in content + assert "catch" in content + assert "console.error" in content + + def test_keyboard_shortcut_implementation(self): + """Test keyboard shortcut implementation.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for keyboard shortcuts + assert "keydown" in content + assert "ctrlKey" in content + assert "shiftKey" in content + assert "preventDefault" in content + + def test_comfyui_integration_points(self): + """Test ComfyUI integration points.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for ComfyUI integration + assert "app.registerExtension" in content + assert "beforeRegisterNodeDef" in content + assert "QueueManagerNode" in content + + def test_notification_system_implementation(self): + """Test notification system implementation.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for notification system + assert "showNotification" in content + assert "notification" in content + assert "createElement" in content + + def test_menu_container_detection(self): + """Test menu container detection logic.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for menu container detection + assert "findMenuContainer" in content + assert "querySelector" in content + assert "comfy-menu" in content + + def test_initialization_sequence(self): + """Test proper initialization sequence.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for proper initialization + assert "init" in content + assert "setup" in content + assert "initializeQueueManager" in content + assert "addMenuIntegration" in content + + +class TestMenuIntegrationRequirements(unittest.TestCase): + """Test that menu integration meets the specified requirements.""" + + def test_requirement_1_1_menu_access(self): + """Test Requirement 1.1: Menu access through ComfyUI.""" + # Verify that menu integration files exist + web_dir = current_dir / "web" + extension_file = web_dir / "queue_manager_extension.js" + + assert extension_file.exists(), ( + "Menu extension file should exist for Requirement 1.1" + ) + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for menu integration + assert "Queue Manager" in content, ( + "Should display 'Queue Manager' option for Requirement 1.1" + ) + + def test_requirement_1_2_interface_opening(self): + """Test Requirement 1.2: Interface opening functionality.""" + web_dir = current_dir / "web" + extension_file = web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for interface opening functionality + assert "openQueueManager" in content, ( + "Should have openQueueManager function for Requirement 1.2" + ) + assert "window.open" in content, ( + "Should open queue management interface for Requirement 1.2" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test_menu_integration.py b/test_menu_integration.py new file mode 100644 index 0000000..bab2880 --- /dev/null +++ b/test_menu_integration.py @@ -0,0 +1,103 @@ +""" +Integration tests for ComfyUI Queue Manager menu functionality. +Tests menu item registration, click handlers, and window management. +""" + +from pathlib import Path +import sys +import unittest + +# Add the current directory to Python path for imports +current_dir = Path(__file__).parent +sys.path.insert(0, str(current_dir)) + + +class TestMenuIntegration(unittest.TestCase): + """Test cases for menu integration functionality.""" + + def test_web_extensions_configuration(self): + """Test that WEB_EXTENSIONS is properly configured.""" + from __init__ import WEB_EXTENSIONS + + # Verify WEB_EXTENSIONS is a list + assert isinstance(WEB_EXTENSIONS, list) + + # Verify it contains the queue manager extension + extension_names = [ext.get("name") for ext in WEB_EXTENSIONS] + assert "queue_manager_extension" in extension_names + + # Verify extension has required fields + queue_ext = next( + ext for ext in WEB_EXTENSIONS if ext["name"] == "queue_manager_extension" + ) + assert "path" in queue_ext + assert queue_ext["path"] == "queue_manager_extension.js" + + def test_web_directory_configuration(self): + """Test that WEB_DIRECTORY is properly configured.""" + from __init__ import WEB_DIRECTORY + + # Verify WEB_DIRECTORY is set correctly + assert WEB_DIRECTORY == "./web" + + def test_register_menu_extension_function_exists(self): + """Test that the register_menu_extension function exists and is callable.""" + from __init__ import register_menu_extension + + # Verify function exists and is callable + assert callable(register_menu_extension) + + def test_node_class_mappings_include_queue_manager(self): + """Test that NODE_CLASS_MAPPINGS includes QueueManagerNode.""" + from __init__ import NODE_CLASS_MAPPINGS + from queue_manager_node import QueueManagerNode + + # Verify QueueManagerNode is in the mappings + assert "QueueManagerNode" in NODE_CLASS_MAPPINGS + + # Verify the mapping points to the correct class + assert NODE_CLASS_MAPPINGS["QueueManagerNode"] == QueueManagerNode + + def test_node_display_name_mappings(self): + """Test that NODE_DISPLAY_NAME_MAPPINGS is properly configured.""" + from __init__ import NODE_DISPLAY_NAME_MAPPINGS + + # Verify QueueManagerNode has a display name + assert "QueueManagerNode" in NODE_DISPLAY_NAME_MAPPINGS + assert NODE_DISPLAY_NAME_MAPPINGS["QueueManagerNode"] == "Queue Manager" + + +class TestMenuExtensionFiles(unittest.TestCase): + """Test cases for menu extension files.""" + + def setUp(self): + """Set up test fixtures.""" + self.web_dir = current_dir / "web" + + def test_menu_extension_file_exists(self): + """Test that the menu extension JavaScript file exists.""" + extension_file = self.web_dir / "queue_manager_extension.js" + assert extension_file.exists(), "queue_manager_extension.js should exist" + + def test_menu_extension_file_content(self): + """Test that the menu extension file has required content.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for required components + assert "app.registerExtension" in content + assert "QueueManager" in content + assert "openQueueManager" in content + assert "addMenuIntegration" in content + + def test_legacy_menu_extension_file_exists(self): + """Test that the legacy menu extension file exists.""" + legacy_file = self.web_dir / "menu_extension.js" + assert legacy_file.exists(), "menu_extension.js should exist as fallback" + + +if __name__ == "__main__": + unittest.main() diff --git a/test_menu_integration_js.py b/test_menu_integration_js.py new file mode 100644 index 0000000..6da8d5d --- /dev/null +++ b/test_menu_integration_js.py @@ -0,0 +1,207 @@ +""" +JavaScript integration tests for ComfyUI Queue Manager menu functionality. +Tests the JavaScript menu extension behavior and DOM interactions. +""" + +from pathlib import Path +import sys +import unittest + +# Add the current directory to Python path for imports +current_dir = Path(__file__).parent +sys.path.insert(0, str(current_dir)) + + +class TestJavaScriptMenuIntegration(unittest.TestCase): + """Test cases for JavaScript menu integration.""" + + def setUp(self): + """Set up test fixtures.""" + self.web_dir = current_dir / "web" + self.extension_file = self.web_dir / "queue_manager_extension.js" + + def test_extension_file_syntax(self): + """Test that the extension JavaScript file has valid syntax.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Basic syntax checks + assert "app.registerExtension" in content + assert "name:" in content + assert "init" in content + assert "setup" in content + + # Check for proper function definitions + assert "openQueueManager" in content + assert "addMenuIntegration" in content + assert "findMenuContainer" in content + assert "createMenuButton" in content + + def test_extension_structure(self): + """Test that the extension has the required structure.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for required extension methods + required_methods = ["init", "setup", "beforeRegisterNodeDef"] + + for method in required_methods: + assert f"{method}(" in content, f"Extension should have {method} method" + + def test_menu_button_creation_logic(self): + """Test the menu button creation logic.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for menu button creation elements + assert "createElement" in content + assert "textContent" in content + assert "addEventListener" in content + assert "appendChild" in content + + def test_window_management_functions(self): + """Test window management functionality.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for window management functions + window_functions = [ + "openAsPopup", + "openAsTab", + "canOpenPopup", + "getQueueManagerUrl", + ] + + for func in window_functions: + assert func in content, f"Should have {func} function" + + def test_keyboard_shortcuts(self): + """Test keyboard shortcut implementation.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for keyboard event handling + assert "keydown" in content + assert "ctrlKey" in content + assert "shiftKey" in content + + def test_notification_system(self): + """Test notification system implementation.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for notification functionality + assert "showNotification" in content + assert "notification" in content + + def test_error_handling(self): + """Test error handling in the extension.""" + if not self.extension_file.exists(): + self.skipTest("Extension file does not exist") + + with self.extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for error handling + assert "try" in content + assert "catch" in content + assert "console.error" in content + + +class TestMenuIntegrationAPI(unittest.TestCase): + """Test cases for menu integration API endpoints.""" + + def test_extension_url_structure(self): + """Test that extension URLs are properly structured.""" + # Test URL patterns that should be registered + expected_patterns = [ + "/extensions/comfyui-queue-manager/{filename}", + "/extensions/comfyui-queue-manager/index.html", + ] + + # This is a structural test - in a real scenario, + # we would test against the actual router registration + for pattern in expected_patterns: + assert isinstance(pattern, str) + assert "comfyui-queue-manager" in pattern + + def test_serve_extension_file_success(self): + """Test serving extension file successfully.""" + # This would test the actual serve_extension_file function + # when it's properly imported and available + + +class TestMenuIntegrationEndToEnd(unittest.TestCase): + """End-to-end tests for menu integration.""" + + def setUp(self): + """Set up test fixtures.""" + self.web_dir = current_dir / "web" + + def test_all_required_files_exist(self): + """Test that all required files for menu integration exist.""" + required_files = [ + "queue_manager_extension.js", + "menu_extension.js", # Legacy fallback + "index.html", + "queue_manager.js", + "styles.css", + ] + + for filename in required_files: + file_path = self.web_dir / filename + assert file_path.exists(), ( + f"Required file {filename} should exist in web directory" + ) + + def test_html_file_has_proper_structure(self): + """Test that the HTML file has proper structure for integration.""" + html_file = self.web_dir / "index.html" + + if html_file.exists(): + with html_file.open(encoding="utf-8") as f: + content = f.read() + + # Check for required HTML elements + assert "" in content + assert "" in content + assert "queue-manager-app" in content + + def test_integration_consistency(self): + """Test consistency between different integration files.""" + extension_file = self.web_dir / "queue_manager_extension.js" + legacy_file = self.web_dir / "menu_extension.js" + + if extension_file.exists() and legacy_file.exists(): + with extension_file.open(encoding="utf-8") as f: + ext_content = f.read() + + with legacy_file.open(encoding="utf-8") as f: + legacy_content = f.read() + + # Both should reference the same URL patterns + assert "queue-manager" in ext_content + assert "queue-manager" in legacy_content + + +if __name__ == "__main__": + unittest.main() diff --git a/test_menu_integration_simple.py b/test_menu_integration_simple.py new file mode 100644 index 0000000..4880406 --- /dev/null +++ b/test_menu_integration_simple.py @@ -0,0 +1,135 @@ +""" +Simple integration tests for ComfyUI Queue Manager menu functionality. +Tests that don't rely on complex imports. +""" + +import sys +import unittest +from pathlib import Path + +# Add the current directory to Python path for imports +current_dir = Path(__file__).parent +sys.path.insert(0, str(current_dir)) + + +class TestMenuIntegrationSimple(unittest.TestCase): + """Simple test cases for menu integration functionality.""" + + def setUp(self): + """Set up test fixtures.""" + self.web_dir = current_dir / "web" + self.init_file = current_dir / "__init__.py" + + def test_init_file_exists(self): + """Test that __init__.py exists.""" + assert self.init_file.exists(), "__init__.py should exist" + + def test_init_file_has_node_mappings(self): + """Test that __init__.py contains node mappings.""" + if self.init_file.exists(): + with self.init_file.open(encoding="utf-8") as f: + content = f.read() + + assert "NODE_CLASS_MAPPINGS" in content + assert "NODE_DISPLAY_NAME_MAPPINGS" in content + assert "QueueManagerNode" in content + + def test_init_file_has_web_configuration(self): + """Test that __init__.py contains web configuration.""" + if self.init_file.exists(): + with self.init_file.open(encoding="utf-8") as f: + content = f.read() + + assert "WEB_DIRECTORY" in content + assert "WEB_EXTENSIONS" in content + assert "./web" in content + + def test_init_file_has_menu_extension_function(self): + """Test that __init__.py contains menu extension function.""" + if self.init_file.exists(): + with self.init_file.open(encoding="utf-8") as f: + content = f.read() + + assert "register_menu_extension" in content + assert "def register_menu_extension" in content + + def test_web_extension_files_exist(self): + """Test that web extension files exist.""" + required_files = [ + "queue_manager_extension.js", + "menu_extension.js", + "index.html", + "queue_manager.js", + "styles.css", + ] + + for filename in required_files: + file_path = self.web_dir / filename + assert file_path.exists(), f"Required file {filename} should exist" + + def test_extension_file_has_comfyui_integration(self): + """Test that extension file has ComfyUI integration.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + assert "app.registerExtension" in content + assert "QueueManager" in content + assert "openQueueManager" in content + + def test_extension_file_has_menu_functionality(self): + """Test that extension file has menu functionality.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + assert "addMenuIntegration" in content + assert "createMenuButton" in content + assert "findMenuContainer" in content + + def test_extension_file_has_window_management(self): + """Test that extension file has window management.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + assert "window.open" in content + assert "openAsPopup" in content + assert "openAsTab" in content + + def test_requirements_coverage(self): + """Test that requirements 1.1 and 1.2 are covered.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + # Requirement 1.1: Menu access + assert "Queue Manager" in content + + # Requirement 1.2: Interface opening + assert "openQueueManager" in content + assert "window.open" in content + + def test_error_handling_present(self): + """Test that error handling is present.""" + extension_file = self.web_dir / "queue_manager_extension.js" + + if extension_file.exists(): + with extension_file.open(encoding="utf-8") as f: + content = f.read() + + assert "try" in content + assert "catch" in content + assert "console.error" in content + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/test_node_registration.py b/test_node_registration.py new file mode 100644 index 0000000..00b271b --- /dev/null +++ b/test_node_registration.py @@ -0,0 +1,253 @@ +""" +Integration tests for ComfyUI custom node registration. +""" + +import sys +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +# Add the current directory to Python path +current_dir = Path(__file__).parent +if str(current_dir) not in sys.path: + sys.path.insert(0, str(current_dir)) + + +class TestNodeRegistration(unittest.TestCase): + """Test ComfyUI custom node registration.""" + + def test_node_class_mappings_exist(self): + """Test that NODE_CLASS_MAPPINGS is properly defined.""" + import __init__ as init_module + NODE_CLASS_MAPPINGS = init_module.NODE_CLASS_MAPPINGS + + self.assertIsInstance(NODE_CLASS_MAPPINGS, dict) + self.assertIn("QueueManagerNode", NODE_CLASS_MAPPINGS) + + # Verify the class is importable + node_class = NODE_CLASS_MAPPINGS["QueueManagerNode"] + self.assertTrue(hasattr(node_class, "INPUT_TYPES")) + self.assertTrue(hasattr(node_class, "process")) + + def test_node_display_name_mappings_exist(self): + """Test that NODE_DISPLAY_NAME_MAPPINGS is properly defined.""" + import __init__ as init_module + NODE_DISPLAY_NAME_MAPPINGS = init_module.NODE_DISPLAY_NAME_MAPPINGS + + self.assertIsInstance(NODE_DISPLAY_NAME_MAPPINGS, dict) + self.assertIn("QueueManagerNode", NODE_DISPLAY_NAME_MAPPINGS) + self.assertEqual(NODE_DISPLAY_NAME_MAPPINGS["QueueManagerNode"], "Queue Manager") + + def test_web_directory_defined(self): + """Test that WEB_DIRECTORY is properly defined.""" + import __init__ as init_module + WEB_DIRECTORY = init_module.WEB_DIRECTORY + + self.assertIsInstance(WEB_DIRECTORY, str) + self.assertEqual(WEB_DIRECTORY, "./web") + + def test_metadata_defined(self): + """Test that extension metadata is properly defined.""" + import __init__ as init_module + __version__ = init_module.__version__ + __author__ = init_module.__author__ + __description__ = init_module.__description__ + + self.assertIsInstance(__version__, str) + self.assertIsInstance(__author__, str) + self.assertIsInstance(__description__, str) + + # Verify version format + version_parts = __version__.split(".") + self.assertEqual(len(version_parts), 3) + for part in version_parts: + self.assertTrue(part.isdigit()) + + def test_initialization_success(self): + """Test successful initialization of queue manager system.""" + # Test that the initialization function exists and can be called + import __init__ as init_module + + # The function should exist + self.assertTrue(hasattr(init_module, 'initialize_queue_manager')) + self.assertTrue(callable(init_module.initialize_queue_manager)) + + # Test that it returns a boolean (success or failure) + result = init_module.initialize_queue_manager() + self.assertIsInstance(result, bool) + + def test_node_class_structure(self): + """Test that the QueueManagerNode class has the required structure.""" + from queue_manager_node import QueueManagerNode + + # Test class methods exist + self.assertTrue(hasattr(QueueManagerNode, "INPUT_TYPES")) + self.assertTrue(hasattr(QueueManagerNode, "IS_CHANGED")) + self.assertTrue(hasattr(QueueManagerNode, "VALIDATE_INPUTS")) + + # Test class attributes exist + self.assertTrue(hasattr(QueueManagerNode, "RETURN_TYPES")) + self.assertTrue(hasattr(QueueManagerNode, "RETURN_NAMES")) + self.assertTrue(hasattr(QueueManagerNode, "FUNCTION")) + self.assertTrue(hasattr(QueueManagerNode, "CATEGORY")) + self.assertTrue(hasattr(QueueManagerNode, "DESCRIPTION")) + + # Test attribute values + self.assertEqual(QueueManagerNode.FUNCTION, "process") + self.assertEqual(QueueManagerNode.CATEGORY, "Queue Management") + self.assertIsInstance(QueueManagerNode.RETURN_TYPES, tuple) + self.assertIsInstance(QueueManagerNode.RETURN_NAMES, tuple) + + def test_input_types_structure(self): + """Test that INPUT_TYPES returns the correct structure.""" + from queue_manager_node import QueueManagerNode + + input_types = QueueManagerNode.INPUT_TYPES() + + self.assertIsInstance(input_types, dict) + self.assertIn("required", input_types) + self.assertIn("optional", input_types) + + # Test optional inputs + optional = input_types["optional"] + self.assertIn("workflow_name", optional) + self.assertIn("auto_queue", optional) + self.assertIn("priority", optional) + self.assertIn("tags", optional) + + # Test input type definitions + self.assertEqual(optional["workflow_name"][0], "STRING") + self.assertEqual(optional["auto_queue"][0], "BOOLEAN") + self.assertEqual(optional["priority"][0], "INT") + + def test_node_instantiation(self): + """Test that the node can be instantiated.""" + from queue_manager_node import QueueManagerNode + + node = QueueManagerNode() + self.assertIsInstance(node, QueueManagerNode) + self.assertIsNone(node._queue_service) + + def test_validate_inputs(self): + """Test input validation.""" + from queue_manager_node import QueueManagerNode + + # Test valid inputs + self.assertTrue(QueueManagerNode.VALIDATE_INPUTS( + workflow_name="Test Workflow", + auto_queue=True, + priority=5, + tags="tag1,tag2" + )) + + # Test invalid priority + self.assertFalse(QueueManagerNode.VALIDATE_INPUTS( + priority=15 # Out of range + )) + + self.assertFalse(QueueManagerNode.VALIDATE_INPUTS( + priority=-15 # Out of range + )) + + # Test invalid workflow name type + self.assertFalse(QueueManagerNode.VALIDATE_INPUTS( + workflow_name=123 # Should be string + )) + + def test_is_changed_returns_nan(self): + """Test that IS_CHANGED returns NaN to always execute.""" + from queue_manager_node import QueueManagerNode + import math + + result = QueueManagerNode.IS_CHANGED() + self.assertTrue(math.isnan(result)) + + def test_get_node_info(self): + """Test that get_node_info returns proper metadata.""" + from queue_manager_node import QueueManagerNode + + info = QueueManagerNode.get_node_info() + + self.assertIsInstance(info, dict) + self.assertEqual(info["name"], "QueueManagerNode") + self.assertEqual(info["display_name"], "Queue Manager") + self.assertEqual(info["category"], "Queue Management") + self.assertIn("description", info) + self.assertIn("version", info) + self.assertIn("author", info) + self.assertIn("input_types", info) + self.assertIn("return_types", info) + self.assertIn("return_names", info) + + +class TestNodeProcessing(unittest.TestCase): + """Test node processing functionality.""" + + def setUp(self): + """Set up test fixtures.""" + from queue_manager_node import QueueManagerNode + self.node = QueueManagerNode() + + def test_process_with_auto_queue_disabled(self): + """Test processing with auto_queue disabled.""" + result = self.node.process( + workflow_name="Test Workflow", + auto_queue=False + ) + + self.assertIsInstance(result, tuple) + self.assertEqual(len(result), 2) + + queue_item_id, queued_successfully = result + self.assertIsInstance(queue_item_id, str) + self.assertFalse(queued_successfully) + + @patch('queue_service.QueueService') + @patch('database.SQLiteDatabase') + def test_process_with_auto_queue_enabled(self, mock_database, mock_queue_service): + """Test processing with auto_queue enabled.""" + # Mock successful queue service + mock_db_instance = MagicMock() + mock_database.return_value = mock_db_instance + + mock_qs_instance = MagicMock() + mock_qs_instance.add_workflow.return_value = "test-queue-id" + mock_queue_service.return_value = mock_qs_instance + + # Set the queue service directly + self.node._queue_service = mock_qs_instance + + result = self.node.process( + workflow_name="Test Workflow", + auto_queue=True, + priority=5, + tags="tag1,tag2" + ) + + self.assertIsInstance(result, tuple) + self.assertEqual(len(result), 2) + + queue_item_id, queued_successfully = result + self.assertEqual(queue_item_id, "test-queue-id") + self.assertTrue(queued_successfully) + + # Verify queue service was called with correct data + mock_qs_instance.add_workflow.assert_called_once() + call_args = mock_qs_instance.add_workflow.call_args[0][0] + self.assertEqual(call_args["name"], "Test Workflow") + self.assertEqual(call_args["priority"], 5) + self.assertEqual(call_args["tags"], ["tag1", "tag2"]) + + def test_get_queue_status_without_service(self): + """Test getting queue status when service is not available.""" + # Mock the queue_service property to return None + with patch.object(type(self.node), 'queue_service', new_callable=lambda: property(lambda self: None)): + result = self.node.get_queue_status() + + self.assertIsInstance(result, dict) + self.assertIn("error", result) + self.assertEqual(result["error"], "Queue service not available") + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/test_workflow_execution.py b/test_workflow_execution.py new file mode 100644 index 0000000..02c2916 --- /dev/null +++ b/test_workflow_execution.py @@ -0,0 +1,778 @@ +""" +Integration tests for workflow execution functionality. +Tests the workflow executor, interceptor, and execution monitor. +""" + +import json +import pytest +import threading +import time +import uuid +from datetime import datetime, timezone +from unittest.mock import Mock, patch, MagicMock + +from database import SQLiteDatabase +from execution_monitor import ExecutionMonitor +from models import QueueStatus +from queue_service import QueueService +from workflow_executor import WorkflowExecutor, WorkflowExecutionError +from workflow_interceptor import WorkflowInterceptor, get_workflow_interceptor + + +class TestWorkflowExecutor: + """Test cases for the WorkflowExecutor class.""" + + @pytest.fixture + def database(self): + """Create a test database.""" + db = SQLiteDatabase(":memory:") + db.initialize() + return db + + @pytest.fixture + def queue_service(self, database): + """Create a test queue service.""" + return QueueService(database) + + @pytest.fixture + def workflow_executor(self, queue_service): + """Create a test workflow executor.""" + return WorkflowExecutor(queue_service) + + @pytest.fixture + def sample_workflow_data(self): + """Create sample workflow data for testing.""" + return { + "name": "Test Workflow", + "nodes": [ + {"id": "1", "type": "LoadImage", "inputs": {}}, + {"id": "2", "type": "SaveImage", "inputs": {"images": ["1", 0]}}, + ], + "links": [["1", 0, "2", 0]], + "workflow": {"version": "1.0"}, + } + + def test_workflow_executor_initialization(self, queue_service): + """Test workflow executor initialization.""" + executor = WorkflowExecutor(queue_service) + + assert executor.queue_service == queue_service + assert executor._interceptor_enabled is True + assert len(executor._running_workflows) == 0 + assert len(executor._execution_callbacks) == 0 + + def test_workflow_executor_without_queue_service(self): + """Test workflow executor initialization without queue service.""" + executor = WorkflowExecutor() + + assert executor.queue_service is not None # Should create default database + assert executor._interceptor_enabled is True + + def test_enable_disable_interceptor(self, workflow_executor): + """Test enabling and disabling the workflow interceptor.""" + # Test initial state + assert workflow_executor.is_interceptor_enabled() is True + + # Test disabling + workflow_executor.disable_interceptor() + assert workflow_executor.is_interceptor_enabled() is False + + # Test enabling + workflow_executor.enable_interceptor() + assert workflow_executor.is_interceptor_enabled() is True + + def test_register_execution_callback(self, workflow_executor): + """Test registering and unregistering execution callbacks.""" + callback_called = [] + + def test_callback(event_type, workflow_id, data): + callback_called.append((event_type, workflow_id, data)) + + # Register callback + workflow_executor.register_execution_callback("test", test_callback) + assert "test" in workflow_executor._execution_callbacks + + # Test callback notification + workflow_executor._notify_callbacks("test_event", "test_id", {"test": "data"}) + assert len(callback_called) == 1 + assert callback_called[0] == ("test_event", "test_id", {"test": "data"}) + + # Unregister callback + workflow_executor.unregister_execution_callback("test") + assert "test" not in workflow_executor._execution_callbacks + + def test_intercept_workflow_execution(self, workflow_executor, sample_workflow_data): + """Test intercepting workflow execution.""" + # Test with interceptor enabled + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + assert workflow_id is not None + assert workflow_id in workflow_executor._running_workflows + + workflow_info = workflow_executor._running_workflows[workflow_id] + assert workflow_info["workflow_data"] == sample_workflow_data + assert workflow_info["status"] == QueueStatus.PENDING + + def test_intercept_workflow_execution_disabled(self, workflow_executor, sample_workflow_data): + """Test workflow execution with interceptor disabled.""" + workflow_executor.disable_interceptor() + + # Should execute directly when interceptor is disabled + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Should not be in running workflows since it executed directly + assert workflow_id not in workflow_executor._running_workflows + + def test_execute_workflow_success(self, workflow_executor, sample_workflow_data): + """Test successful workflow execution.""" + result = workflow_executor.execute_workflow(sample_workflow_data) + + assert isinstance(result, dict) + assert "status" in result + assert result["status"] == "completed" + + def test_execute_workflow_failure(self, workflow_executor): + """Test workflow execution failure handling.""" + # Create invalid workflow data that should cause failure + invalid_workflow = {"invalid": "data"} + + with patch.object(workflow_executor, '_execute_workflow_directly', + side_effect=Exception("Execution failed")): + with pytest.raises(WorkflowExecutionError): + workflow_executor.execute_workflow(invalid_workflow) + + def test_extract_workflow_name(self, workflow_executor): + """Test workflow name extraction from data.""" + # Test explicit name + data_with_name = {"name": "My Workflow", "nodes": []} + name = workflow_executor._extract_workflow_name(data_with_name) + assert name == "My Workflow" + + # Test title + data_with_title = {"title": "My Title", "nodes": []} + name = workflow_executor._extract_workflow_name(data_with_title) + assert name == "My Title" + + # Test workflow metadata + data_with_workflow = {"workflow": {"name": "Workflow Name"}, "nodes": []} + name = workflow_executor._extract_workflow_name(data_with_workflow) + assert name == "Workflow Name" + + # Test node count + data_with_nodes = {"nodes": [1, 2, 3]} + name = workflow_executor._extract_workflow_name(data_with_nodes) + assert "3 nodes" in name + + # Test default name + empty_data = {} + name = workflow_executor._extract_workflow_name(empty_data) + assert "Workflow" in name + + def test_update_workflow_status(self, workflow_executor, sample_workflow_data): + """Test updating workflow status.""" + # First intercept a workflow + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Update to running + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + + workflow_info = workflow_executor._running_workflows[workflow_id] + assert workflow_info["status"] == QueueStatus.RUNNING + assert workflow_info["started_at"] is not None + + # Update to completed + result_data = {"output": "test"} + workflow_executor._update_workflow_status( + workflow_id, QueueStatus.COMPLETED, result_data=result_data + ) + + workflow_info = workflow_executor._running_workflows[workflow_id] + assert workflow_info["status"] == QueueStatus.COMPLETED + assert workflow_info["completed_at"] is not None + assert workflow_info["result_data"] == result_data + + def test_is_workflow_running(self, workflow_executor, sample_workflow_data): + """Test checking if workflow is running.""" + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Initially pending + assert workflow_executor.is_workflow_running(workflow_id) is False + + # Update to running + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + assert workflow_executor.is_workflow_running(workflow_id) is True + + # Update to completed + workflow_executor._update_workflow_status(workflow_id, QueueStatus.COMPLETED) + assert workflow_executor.is_workflow_running(workflow_id) is False + + def test_cancel_workflow(self, workflow_executor, sample_workflow_data): + """Test cancelling a workflow.""" + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Cancel pending workflow + result = workflow_executor.cancel_workflow(workflow_id) + assert result is True + + workflow_info = workflow_executor._running_workflows[workflow_id] + assert workflow_info["status"] == QueueStatus.FAILED + assert "cancelled" in workflow_info["error_message"] + + def test_cancel_nonexistent_workflow(self, workflow_executor): + """Test cancelling a non-existent workflow.""" + result = workflow_executor.cancel_workflow("nonexistent") + assert result is False + + def test_get_workflow_status(self, workflow_executor, sample_workflow_data): + """Test getting workflow status.""" + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + status = workflow_executor.get_workflow_status(workflow_id) + assert status == QueueStatus.PENDING + + # Test non-existent workflow + status = workflow_executor.get_workflow_status("nonexistent") + assert status is None + + def test_get_running_workflows(self, workflow_executor, sample_workflow_data): + """Test getting running workflows.""" + # Initially no running workflows + running = workflow_executor.get_running_workflows() + assert len(running) == 0 + + # Add a workflow and set to running + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + + running = workflow_executor.get_running_workflows() + assert len(running) == 1 + assert workflow_id in running + + def test_cleanup_completed_workflows(self, workflow_executor, sample_workflow_data): + """Test cleaning up completed workflows.""" + # Add and complete a workflow + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + workflow_executor._update_workflow_status(workflow_id, QueueStatus.COMPLETED) + + # Set completed time to old timestamp + workflow_info = workflow_executor._running_workflows[workflow_id] + old_time = datetime.now(timezone.utc).replace(year=2020) + workflow_info["completed_at"] = old_time + + # Cleanup should remove the old workflow + cleaned = workflow_executor.cleanup_completed_workflows(max_age_hours=1) + assert cleaned == 1 + assert workflow_id not in workflow_executor._running_workflows + + def test_get_execution_statistics(self, workflow_executor, sample_workflow_data): + """Test getting execution statistics.""" + # Add some workflows with different statuses + wf1 = workflow_executor.intercept_workflow_execution(sample_workflow_data) + wf2 = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + workflow_executor._update_workflow_status(wf1, QueueStatus.RUNNING) + workflow_executor._update_workflow_status(wf2, QueueStatus.COMPLETED) + + stats = workflow_executor.get_execution_statistics() + + assert stats["total_workflows"] == 2 + assert "status_counts" in stats + assert stats["status_counts"]["running"] == 1 + assert stats["status_counts"]["completed"] == 1 + + +class TestWorkflowInterceptor: + """Test cases for the WorkflowInterceptor class.""" + + @pytest.fixture + def workflow_interceptor(self): + """Create a test workflow interceptor.""" + # Reset singleton for testing + WorkflowInterceptor._instance = None + return WorkflowInterceptor() + + @pytest.fixture + def mock_workflow_executor(self): + """Create a mock workflow executor.""" + executor = Mock() + executor.intercept_workflow_execution.return_value = "test_workflow_id" + return executor + + def test_interceptor_singleton(self): + """Test that interceptor follows singleton pattern.""" + interceptor1 = WorkflowInterceptor() + interceptor2 = WorkflowInterceptor() + + assert interceptor1 is interceptor2 + + def test_get_workflow_interceptor_function(self): + """Test the global get_workflow_interceptor function.""" + interceptor = get_workflow_interceptor() + assert isinstance(interceptor, WorkflowInterceptor) + + def test_set_workflow_executor(self, workflow_interceptor, mock_workflow_executor): + """Test setting workflow executor.""" + workflow_interceptor.set_workflow_executor(mock_workflow_executor) + assert workflow_interceptor._workflow_executor == mock_workflow_executor + + def test_enable_disable_interceptor(self, workflow_interceptor): + """Test enabling and disabling interceptor.""" + # Initially disabled + assert workflow_interceptor.is_enabled() is False + + # Enable + result = workflow_interceptor.enable() + assert result is True + assert workflow_interceptor.is_enabled() is True + + # Disable + result = workflow_interceptor.disable() + assert result is True + assert workflow_interceptor.is_enabled() is False + + def test_register_execution_hook(self, workflow_interceptor): + """Test registering execution hooks.""" + hook_called = [] + + def test_hook(event_type, workflow_id, data): + hook_called.append((event_type, workflow_id, data)) + + workflow_interceptor.register_execution_hook("test_hook", test_hook) + assert "test_hook" in workflow_interceptor._execution_hooks + + # Unregister + workflow_interceptor.unregister_execution_hook("test_hook") + assert "test_hook" not in workflow_interceptor._execution_hooks + + def test_looks_like_workflow_data(self, workflow_interceptor): + """Test workflow data detection.""" + # Valid workflow data + valid_data = {"nodes": [], "links": []} + assert workflow_interceptor._looks_like_workflow_data(valid_data) is True + + # Another valid format + valid_data2 = {"workflow": {"version": "1.0"}} + assert workflow_interceptor._looks_like_workflow_data(valid_data2) is True + + # Invalid data + invalid_data = {"random": "data"} + assert workflow_interceptor._looks_like_workflow_data(invalid_data) is False + + # Non-dict data + assert workflow_interceptor._looks_like_workflow_data("not a dict") is False + + def test_extract_workflow_data(self, workflow_interceptor): + """Test extracting workflow data from arguments.""" + workflow_data = {"nodes": [], "workflow": {"version": "1.0"}} + + # Test positional arguments + args = (workflow_data, "other_arg") + kwargs = {} + + extracted = workflow_interceptor._extract_workflow_data(args, kwargs) + assert extracted == workflow_data + + # Test keyword arguments + args = () + kwargs = {"workflow": workflow_data, "other": "value"} + + extracted = workflow_interceptor._extract_workflow_data(args, kwargs) + assert extracted == workflow_data + + def test_extract_result_data(self, workflow_interceptor): + """Test extracting result data.""" + # Dict result + dict_result = {"output": "test", "status": "completed"} + extracted = workflow_interceptor._extract_result_data(dict_result) + assert extracted == dict_result + + # Object with __dict__ + class TestResult: + def __init__(self): + self.output = "test" + self.status = "completed" + + obj_result = TestResult() + extracted = workflow_interceptor._extract_result_data(obj_result) + assert extracted["output"] == "test" + assert extracted["status"] == "completed" + + # String result + str_result = "completed" + extracted = workflow_interceptor._extract_result_data(str_result) + assert extracted == {"result": "completed"} + + def test_get_hook_status(self, workflow_interceptor, mock_workflow_executor): + """Test getting hook status.""" + workflow_interceptor.set_workflow_executor(mock_workflow_executor) + workflow_interceptor.register_execution_hook("test", lambda: None) + + status = workflow_interceptor.get_hook_status() + + assert "enabled" in status + assert "execution_hooks" in status + assert "has_workflow_executor" in status + assert status["has_workflow_executor"] is True + assert "test" in status["execution_hooks"] + + +class TestExecutionMonitor: + """Test cases for the ExecutionMonitor class.""" + + @pytest.fixture + def database(self): + """Create a test database.""" + db = SQLiteDatabase(":memory:") + db.initialize() + return db + + @pytest.fixture + def queue_service(self, database): + """Create a test queue service.""" + return QueueService(database) + + @pytest.fixture + def workflow_executor(self, queue_service): + """Create a test workflow executor.""" + return WorkflowExecutor(queue_service) + + @pytest.fixture + def execution_monitor(self, queue_service, workflow_executor): + """Create a test execution monitor.""" + return ExecutionMonitor(queue_service, workflow_executor) + + def test_execution_monitor_initialization(self, queue_service, workflow_executor): + """Test execution monitor initialization.""" + monitor = ExecutionMonitor(queue_service, workflow_executor) + + assert monitor.queue_service == queue_service + assert monitor.workflow_executor == workflow_executor + assert monitor.is_monitoring() is False + assert len(monitor._monitored_workflows) == 0 + + def test_start_stop_monitoring(self, execution_monitor): + """Test starting and stopping monitoring.""" + # Start monitoring + result = execution_monitor.start_monitoring() + assert result is True + assert execution_monitor.is_monitoring() is True + + # Stop monitoring + result = execution_monitor.stop_monitoring() + assert result is True + assert execution_monitor.is_monitoring() is False + + def test_set_monitoring_interval(self, execution_monitor): + """Test setting monitoring interval.""" + execution_monitor.set_monitoring_interval(2.0) + assert execution_monitor._monitoring_interval == 2.0 + + # Test invalid interval + with pytest.raises(ValueError): + execution_monitor.set_monitoring_interval(-1.0) + + def test_add_remove_workflow_to_monitor(self, execution_monitor): + """Test adding and removing workflows from monitoring.""" + workflow_id = "test_workflow" + workflow_info = { + "status": QueueStatus.PENDING, + "workflow_name": "Test Workflow", + } + + # Add workflow + execution_monitor.add_workflow_to_monitor(workflow_id, workflow_info) + assert workflow_id in execution_monitor._monitored_workflows + + # Remove workflow + result = execution_monitor.remove_workflow_from_monitor(workflow_id) + assert result is True + assert workflow_id not in execution_monitor._monitored_workflows + + # Remove non-existent workflow + result = execution_monitor.remove_workflow_from_monitor("nonexistent") + assert result is False + + def test_register_status_callback(self, execution_monitor): + """Test registering status callbacks.""" + callback_called = [] + + def test_callback(data): + callback_called.append(data) + + execution_monitor.register_status_callback("test", test_callback) + assert "test" in execution_monitor._status_callbacks + + # Unregister + execution_monitor.unregister_status_callback("test") + assert "test" not in execution_monitor._status_callbacks + + def test_get_monitored_workflows(self, execution_monitor): + """Test getting monitored workflows.""" + workflow_info = {"status": QueueStatus.PENDING} + execution_monitor.add_workflow_to_monitor("test1", workflow_info) + execution_monitor.add_workflow_to_monitor("test2", workflow_info) + + workflows = execution_monitor.get_monitored_workflows() + assert len(workflows) == 2 + assert "test1" in workflows + assert "test2" in workflows + + def test_get_workflow_status(self, execution_monitor): + """Test getting workflow status.""" + workflow_info = {"status": QueueStatus.RUNNING} + execution_monitor.add_workflow_to_monitor("test", workflow_info) + + status = execution_monitor.get_workflow_status("test") + assert status == QueueStatus.RUNNING + + # Non-existent workflow + status = execution_monitor.get_workflow_status("nonexistent") + assert status is None + + def test_monitoring_statistics(self, execution_monitor): + """Test getting monitoring statistics.""" + # Add some workflows + execution_monitor.add_workflow_to_monitor("test1", {"status": QueueStatus.PENDING}) + execution_monitor.add_workflow_to_monitor("test2", {"status": QueueStatus.RUNNING}) + + stats = execution_monitor.get_monitoring_statistics() + + assert stats["monitored_workflows"] == 2 + assert "status_counts" in stats + assert stats["status_counts"]["pending"] == 1 + assert stats["status_counts"]["running"] == 1 + + def test_cleanup_old_workflows(self, execution_monitor): + """Test cleaning up old workflows.""" + # Add old completed workflow + old_workflow_info = { + "status": QueueStatus.COMPLETED, + "added_at": datetime.now(timezone.utc).replace(year=2020), + } + execution_monitor.add_workflow_to_monitor("old", old_workflow_info) + + # Add recent workflow + recent_workflow_info = { + "status": QueueStatus.PENDING, + "added_at": datetime.now(timezone.utc), + } + execution_monitor.add_workflow_to_monitor("recent", recent_workflow_info) + + # Cleanup should remove only the old completed workflow + cleaned = execution_monitor.cleanup_old_workflows(max_age_hours=1) + assert cleaned == 1 + assert "old" not in execution_monitor._monitored_workflows + assert "recent" in execution_monitor._monitored_workflows + + +class TestWorkflowExecutionIntegration: + """Integration tests for the complete workflow execution system.""" + + @pytest.fixture + def database(self): + """Create a test database.""" + db = SQLiteDatabase(":memory:") + db.initialize() + return db + + @pytest.fixture + def queue_service(self, database): + """Create a test queue service.""" + return QueueService(database) + + @pytest.fixture + def workflow_executor(self, queue_service): + """Create a test workflow executor.""" + return WorkflowExecutor(queue_service) + + @pytest.fixture + def execution_monitor(self, queue_service, workflow_executor): + """Create a test execution monitor.""" + return ExecutionMonitor(queue_service, workflow_executor) + + @pytest.fixture + def workflow_interceptor(self, workflow_executor): + """Create a test workflow interceptor.""" + WorkflowInterceptor._instance = None + interceptor = WorkflowInterceptor() + interceptor.set_workflow_executor(workflow_executor) + return interceptor + + @pytest.fixture + def sample_workflow_data(self): + """Create sample workflow data for testing.""" + return { + "name": "Integration Test Workflow", + "nodes": [ + {"id": "1", "type": "LoadImage", "inputs": {}}, + {"id": "2", "type": "SaveImage", "inputs": {"images": ["1", 0]}}, + ], + "links": [["1", 0, "2", 0]], + "workflow": {"version": "1.0"}, + } + + def test_complete_workflow_execution_flow( + self, + queue_service, + workflow_executor, + execution_monitor, + workflow_interceptor, + sample_workflow_data + ): + """Test the complete workflow execution flow.""" + # Start monitoring + execution_monitor.start_monitoring() + + # Enable interceptor + workflow_interceptor.enable() + + # Track status changes + status_changes = [] + + def status_callback(data): + status_changes.append(data) + + execution_monitor.register_status_callback("test", status_callback) + + # Intercept a workflow + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Verify workflow was added to queue + queue_items = queue_service.get_queue_items() + assert len(queue_items) == 1 + + # Add workflow to monitoring + workflow_info = workflow_executor.get_workflow_info(workflow_id) + execution_monitor.add_workflow_to_monitor(workflow_id, workflow_info) + + # Simulate workflow execution + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + time.sleep(0.1) # Allow monitor to detect change + + workflow_executor._update_workflow_status( + workflow_id, + QueueStatus.COMPLETED, + result_data={"output": "test_result"} + ) + time.sleep(0.1) # Allow monitor to detect change + + # Verify final state + final_status = workflow_executor.get_workflow_status(workflow_id) + assert final_status == QueueStatus.COMPLETED + + # Verify queue item was updated + queue_item = queue_service.get_queue_item(workflow_info["queue_item_id"]) + assert queue_item.status == QueueStatus.COMPLETED + assert queue_item.result_data == {"output": "test_result"} + + # Stop monitoring + execution_monitor.stop_monitoring() + + def test_workflow_execution_error_handling( + self, + workflow_executor, + execution_monitor, + sample_workflow_data + ): + """Test error handling in workflow execution.""" + execution_monitor.start_monitoring() + + # Intercept workflow + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + + # Add to monitoring + workflow_info = workflow_executor.get_workflow_info(workflow_id) + execution_monitor.add_workflow_to_monitor(workflow_id, workflow_info) + + # Simulate execution failure + workflow_executor._update_workflow_status( + workflow_id, + QueueStatus.FAILED, + error_message="Test execution error" + ) + + # Verify error state + final_status = workflow_executor.get_workflow_status(workflow_id) + assert final_status == QueueStatus.FAILED + + workflow_info = workflow_executor.get_workflow_info(workflow_id) + assert workflow_info["error_message"] == "Test execution error" + + execution_monitor.stop_monitoring() + + def test_concurrent_workflow_execution( + self, + workflow_executor, + execution_monitor, + sample_workflow_data + ): + """Test concurrent workflow execution.""" + execution_monitor.start_monitoring() + + # Start multiple workflows + workflow_ids = [] + for i in range(3): + workflow_data = {**sample_workflow_data, "name": f"Workflow {i}"} + workflow_id = workflow_executor.intercept_workflow_execution(workflow_data) + workflow_ids.append(workflow_id) + + # Add to monitoring + workflow_info = workflow_executor.get_workflow_info(workflow_id) + execution_monitor.add_workflow_to_monitor(workflow_id, workflow_info) + + # Set all to running + for workflow_id in workflow_ids: + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + + # Verify all are running + running_workflows = workflow_executor.get_running_workflows() + assert len(running_workflows) == 3 + + # Complete workflows one by one + for i, workflow_id in enumerate(workflow_ids): + workflow_executor._update_workflow_status( + workflow_id, + QueueStatus.COMPLETED, + result_data={"output": f"result_{i}"} + ) + + # Verify all completed + for workflow_id in workflow_ids: + status = workflow_executor.get_workflow_status(workflow_id) + assert status == QueueStatus.COMPLETED + + execution_monitor.stop_monitoring() + + def test_workflow_cancellation( + self, + workflow_executor, + execution_monitor, + sample_workflow_data + ): + """Test workflow cancellation.""" + execution_monitor.start_monitoring() + + # Start workflow + workflow_id = workflow_executor.intercept_workflow_execution(sample_workflow_data) + workflow_info = workflow_executor.get_workflow_info(workflow_id) + execution_monitor.add_workflow_to_monitor(workflow_id, workflow_info) + + # Set to running + workflow_executor._update_workflow_status(workflow_id, QueueStatus.RUNNING) + + # Cancel workflow + result = workflow_executor.cancel_workflow(workflow_id) + assert result is True + + # Verify cancelled state + status = workflow_executor.get_workflow_status(workflow_id) + assert status == QueueStatus.FAILED + + workflow_info = workflow_executor.get_workflow_info(workflow_id) + assert "cancelled" in workflow_info["error_message"] + + execution_monitor.stop_monitoring() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/web/menu_extension.js b/web/menu_extension.js new file mode 100644 index 0000000..03cebf6 --- /dev/null +++ b/web/menu_extension.js @@ -0,0 +1,402 @@ +/** + * ComfyUI Queue Manager Menu Extension + * Integrates the Queue Manager into ComfyUI's menu system + */ + +import { app } from "/scripts/app.js"; +import { api } from "/scripts/api.js"; + +// Queue Manager window reference +let queueManagerWindow = null; +let queueManagerTab = null; + +/** + * Queue Manager Menu Extension + */ +class QueueManagerMenuExtension { + constructor() { + this.name = "QueueManager.MenuExtension"; + this.menuItem = null; + } + + /** + * Initialize the menu extension + */ + init() { + // Register menu item when ComfyUI is ready + app.registerExtension({ + name: this.name, + async setup() { + // Add menu item to ComfyUI's menu + this.addMenuItem(); + }, + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // Hook into node registration if needed + if (nodeData.name === "QueueManagerNode") { + // Add any node-specific menu integrations here + console.log("[Queue Manager] Node registered:", nodeData.name); + } + }, + }); + } + + /** + * Add Queue Manager menu item to ComfyUI's menu + */ + addMenuItem() { + try { + // Get ComfyUI's menu container + const menuContainer = + document.querySelector(".comfy-menu") || + document.querySelector("#comfy-menu") || + document.querySelector(".menu"); + + if (!menuContainer) { + console.warn( + "[Queue Manager] Could not find ComfyUI menu container, trying alternative approach", + ); + this.addMenuItemAlternative(); + return; + } + + // Create menu item + const menuItem = document.createElement("button"); + menuItem.className = "comfy-menu-btns"; + menuItem.textContent = "Queue Manager"; + menuItem.title = "Open Queue Manager"; + menuItem.style.cssText = ` + background: #333; + color: #fff; + border: 1px solid #555; + padding: 4px 8px; + margin: 2px; + cursor: pointer; + border-radius: 3px; + font-size: 12px; + `; + + // Add click handler + menuItem.addEventListener("click", () => { + this.openQueueManager(); + }); + + // Add hover effects + menuItem.addEventListener("mouseenter", () => { + menuItem.style.background = "#555"; + }); + + menuItem.addEventListener("mouseleave", () => { + menuItem.style.background = "#333"; + }); + + // Insert menu item + menuContainer.appendChild(menuItem); + this.menuItem = menuItem; + + console.log("[Queue Manager] Menu item added successfully"); + } catch (error) { + console.error("[Queue Manager] Failed to add menu item:", error); + this.addMenuItemAlternative(); + } + } + + /** + * Alternative method to add menu item if standard approach fails + */ + addMenuItemAlternative() { + try { + // Create a floating button as fallback + const floatingButton = document.createElement("div"); + floatingButton.innerHTML = ` + + `; + + document.body.appendChild(floatingButton); + + const button = document.getElementById("queue-manager-floating-btn"); + button.addEventListener("click", () => { + this.openQueueManager(); + }); + + button.addEventListener("mouseenter", () => { + button.style.background = "#45a049"; + }); + + button.addEventListener("mouseleave", () => { + button.style.background = "#4CAF50"; + }); + + this.menuItem = button; + console.log("[Queue Manager] Floating menu button added as fallback"); + } catch (error) { + console.error("[Queue Manager] Failed to add fallback menu item:", error); + } + } + + /** + * Open the Queue Manager interface + */ + openQueueManager() { + try { + // Check if window is already open + if (queueManagerWindow && !queueManagerWindow.closed) { + queueManagerWindow.focus(); + return; + } + + // Determine the best way to open the interface + if (this.canUsePopupWindow()) { + this.openAsPopupWindow(); + } else { + this.openAsTab(); + } + } catch (error) { + console.error("[Queue Manager] Failed to open queue manager:", error); + this.showError("Failed to open Queue Manager"); + } + } + + /** + * Check if popup windows are allowed + */ + canUsePopupWindow() { + try { + // Test if we can open popup windows + const testWindow = window.open("", "_blank", "width=1,height=1"); + if (testWindow) { + testWindow.close(); + return true; + } + return false; + } catch (error) { + return false; + } + } + + /** + * Open Queue Manager as a popup window + */ + openAsPopupWindow() { + const windowFeatures = [ + "width=1200", + "height=800", + "left=100", + "top=100", + "resizable=yes", + "scrollbars=yes", + "status=no", + "menubar=no", + "toolbar=no", + "location=no", + ].join(","); + + // Get the queue manager URL + const queueManagerUrl = this.getQueueManagerUrl(); + + queueManagerWindow = window.open( + queueManagerUrl, + "QueueManager", + windowFeatures, + ); + + if (queueManagerWindow) { + queueManagerWindow.focus(); + + // Handle window close + queueManagerWindow.addEventListener("beforeunload", () => { + queueManagerWindow = null; + }); + + console.log("[Queue Manager] Opened as popup window"); + } else { + console.warn("[Queue Manager] Popup blocked, falling back to tab"); + this.openAsTab(); + } + } + + /** + * Open Queue Manager as a new tab + */ + openAsTab() { + const queueManagerUrl = this.getQueueManagerUrl(); + queueManagerTab = window.open(queueManagerUrl, "_blank"); + + if (queueManagerTab) { + queueManagerTab.focus(); + console.log("[Queue Manager] Opened as new tab"); + } else { + console.error("[Queue Manager] Failed to open tab"); + this.openAsEmbedded(); + } + } + + /** + * Open Queue Manager as embedded iframe (fallback) + */ + openAsEmbedded() { + try { + // Create modal overlay + const overlay = document.createElement("div"); + overlay.id = "queue-manager-overlay"; + overlay.style.cssText = ` + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: rgba(0, 0, 0, 0.8); + z-index: 10000; + display: flex; + justify-content: center; + align-items: center; + `; + + // Create iframe container + const container = document.createElement("div"); + container.style.cssText = ` + width: 90%; + height: 90%; + background: white; + border-radius: 8px; + position: relative; + box-shadow: 0 4px 20px rgba(0, 0, 0, 0.3); + `; + + // Create close button + const closeButton = document.createElement("button"); + closeButton.innerHTML = "×"; + closeButton.style.cssText = ` + position: absolute; + top: 10px; + right: 15px; + background: none; + border: none; + font-size: 24px; + cursor: pointer; + z-index: 10001; + color: #666; + `; + + closeButton.addEventListener("click", () => { + document.body.removeChild(overlay); + }); + + // Create iframe + const iframe = document.createElement("iframe"); + iframe.src = this.getQueueManagerUrl(); + iframe.style.cssText = ` + width: 100%; + height: 100%; + border: none; + border-radius: 8px; + `; + + // Assemble components + container.appendChild(closeButton); + container.appendChild(iframe); + overlay.appendChild(container); + document.body.appendChild(overlay); + + // Close on overlay click + overlay.addEventListener("click", (e) => { + if (e.target === overlay) { + document.body.removeChild(overlay); + } + }); + + console.log("[Queue Manager] Opened as embedded iframe"); + } catch (error) { + console.error( + "[Queue Manager] Failed to open embedded interface:", + error, + ); + this.showError("Failed to open Queue Manager interface"); + } + } + + /** + * Get the Queue Manager URL + */ + getQueueManagerUrl() { + // Use the web directory served by ComfyUI + const baseUrl = window.location.origin; + return `${baseUrl}/extensions/comfyui-queue-manager/index.html`; + } + + /** + * Show error message + */ + showError(message) { + // Create simple error notification + const notification = document.createElement("div"); + notification.style.cssText = ` + position: fixed; + top: 20px; + right: 20px; + background: #f44336; + color: white; + padding: 12px 20px; + border-radius: 4px; + z-index: 10000; + font-size: 14px; + box-shadow: 0 2px 5px rgba(0,0,0,0.3); + `; + notification.textContent = message; + + document.body.appendChild(notification); + + // Auto-remove after 5 seconds + setTimeout(() => { + if (document.body.contains(notification)) { + document.body.removeChild(notification); + } + }, 5000); + } + + /** + * Clean up resources + */ + cleanup() { + if (queueManagerWindow && !queueManagerWindow.closed) { + queueManagerWindow.close(); + } + + if (this.menuItem && this.menuItem.parentNode) { + this.menuItem.parentNode.removeChild(this.menuItem); + } + } +} + +// Initialize the menu extension +const queueManagerMenuExtension = new QueueManagerMenuExtension(); + +// Auto-initialize when DOM is ready +if (document.readyState === "loading") { + document.addEventListener("DOMContentLoaded", () => { + queueManagerMenuExtension.init(); + }); +} else { + queueManagerMenuExtension.init(); +} + +// Export for testing and external access +window.QueueManagerMenuExtension = queueManagerMenuExtension; + +export default queueManagerMenuExtension; diff --git a/web/queue_manager_extension.js b/web/queue_manager_extension.js new file mode 100644 index 0000000..8bea808 --- /dev/null +++ b/web/queue_manager_extension.js @@ -0,0 +1,391 @@ +/** + * ComfyUI Queue Manager Extension + * Registers the Queue Manager with ComfyUI's extension system + */ + +import { app } from "/scripts/app.js"; +import { api } from "/scripts/api.js"; + +// Extension name +const EXTENSION_NAME = "QueueManager"; + +// Queue Manager state +let queueManagerState = { + window: null, + tab: null, + isOpen: false, + menuButton: null, +}; + +/** + * Queue Manager ComfyUI Extension + */ +app.registerExtension({ + name: EXTENSION_NAME, + + async init(app) { + console.log(`[${EXTENSION_NAME}] Initializing extension`); + + // Initialize queue manager integration + this.initializeQueueManager(); + }, + + async setup(app) { + console.log(`[${EXTENSION_NAME}] Setting up extension`); + + // Add menu integration + this.addMenuIntegration(); + + // Set up keyboard shortcuts + this.setupKeyboardShortcuts(); + }, + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // Hook into Queue Manager node registration + if (nodeData.name === "QueueManagerNode") { + console.log(`[${EXTENSION_NAME}] Queue Manager node detected`); + + // Add custom functionality to the node if needed + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + if (onNodeCreated) { + onNodeCreated.apply(this, arguments); + } + + // Add queue manager specific functionality + this.addWidget("button", "Open Queue Manager", null, () => { + openQueueManager(); + }); + }; + } + }, + + /** + * Initialize Queue Manager integration + */ + initializeQueueManager() { + // Check if Queue Manager API is available + this.checkQueueManagerAPI(); + + // Set up event listeners + this.setupEventListeners(); + }, + + /** + * Add menu integration to ComfyUI + */ + addMenuIntegration() { + try { + // Wait for ComfyUI menu to be ready + const addMenuButton = () => { + const menuContainer = this.findMenuContainer(); + if (menuContainer) { + this.createMenuButton(menuContainer); + } else { + // Retry after a short delay + setTimeout(addMenuButton, 500); + } + }; + + // Add menu button when DOM is ready + if (document.readyState === "loading") { + document.addEventListener("DOMContentLoaded", addMenuButton); + } else { + addMenuButton(); + } + } catch (error) { + console.error( + `[${EXTENSION_NAME}] Failed to add menu integration:`, + error, + ); + } + }, + + /** + * Find ComfyUI's menu container + */ + findMenuContainer() { + // Try multiple selectors to find the menu + const selectors = [ + ".comfy-menu", + "#comfy-menu", + ".menu", + ".comfy-menu-btns", + ".comfyui-menu", + ]; + + for (const selector of selectors) { + const container = document.querySelector(selector); + if (container) { + return container; + } + } + + // Try to find by looking for existing menu buttons + const existingButtons = document.querySelectorAll("button"); + for (const button of existingButtons) { + if ( + button.textContent.includes("Queue") || + button.textContent.includes("Settings") || + button.className.includes("comfy") + ) { + return button.parentElement; + } + } + + return null; + }, + + /** + * Create the Queue Manager menu button + */ + createMenuButton(container) { + if (queueManagerState.menuButton) { + return; // Already created + } + + const button = document.createElement("button"); + button.textContent = "📋 Queue Manager"; + button.title = "Open Queue Manager"; + button.className = "comfy-menu-btns queue-manager-btn"; + + // Style the button to match ComfyUI's theme + button.style.cssText = ` + background: var(--comfy-menu-bg, #333); + color: var(--comfy-menu-text, #fff); + border: 1px solid var(--comfy-menu-border, #555); + padding: 4px 8px; + margin: 2px; + cursor: pointer; + border-radius: 3px; + font-size: 12px; + transition: background-color 0.2s; + `; + + // Add hover effects + button.addEventListener("mouseenter", () => { + button.style.backgroundColor = "var(--comfy-menu-hover-bg, #555)"; + }); + + button.addEventListener("mouseleave", () => { + button.style.backgroundColor = "var(--comfy-menu-bg, #333)"; + }); + + // Add click handler + button.addEventListener("click", (e) => { + e.preventDefault(); + e.stopPropagation(); + openQueueManager(); + }); + + // Insert button into menu + container.appendChild(button); + queueManagerState.menuButton = button; + + console.log(`[${EXTENSION_NAME}] Menu button added successfully`); + }, + + /** + * Set up keyboard shortcuts + */ + setupKeyboardShortcuts() { + document.addEventListener("keydown", (e) => { + // Ctrl+Shift+Q to open Queue Manager + if (e.ctrlKey && e.shiftKey && e.key === "Q") { + e.preventDefault(); + openQueueManager(); + } + }); + }, + + /** + * Set up event listeners + */ + setupEventListeners() { + // Listen for window close events + window.addEventListener("beforeunload", () => { + if (queueManagerState.window && !queueManagerState.window.closed) { + queueManagerState.window.close(); + } + }); + }, + + /** + * Check if Queue Manager API is available + */ + async checkQueueManagerAPI() { + try { + const response = await fetch("/queue/status"); + if (response.ok) { + console.log(`[${EXTENSION_NAME}] Queue Manager API is available`); + return true; + } + } catch (error) { + console.warn( + `[${EXTENSION_NAME}] Queue Manager API not available:`, + error, + ); + } + return false; + }, +}); + +/** + * Open the Queue Manager interface + */ +function openQueueManager() { + try { + // Check if already open + if (queueManagerState.window && !queueManagerState.window.closed) { + queueManagerState.window.focus(); + return; + } + + // Determine the best way to open + if (canOpenPopup()) { + openAsPopup(); + } else { + openAsTab(); + } + } catch (error) { + console.error(`[${EXTENSION_NAME}] Failed to open Queue Manager:`, error); + showNotification("Failed to open Queue Manager", "error"); + } +} + +/** + * Check if popup windows can be opened + */ +function canOpenPopup() { + try { + const test = window.open("", "_blank", "width=1,height=1"); + if (test) { + test.close(); + return true; + } + } catch (error) { + // Popup blocked + } + return false; +} + +/** + * Open Queue Manager as popup window + */ +function openAsPopup() { + const features = [ + "width=1200", + "height=800", + "left=100", + "top=100", + "resizable=yes", + "scrollbars=yes", + "status=no", + "menubar=no", + "toolbar=no", + ].join(","); + + const url = getQueueManagerUrl(); + queueManagerState.window = window.open(url, "QueueManager", features); + + if (queueManagerState.window) { + queueManagerState.window.focus(); + queueManagerState.isOpen = true; + + // Handle window close + const checkClosed = setInterval(() => { + if (queueManagerState.window.closed) { + queueManagerState.isOpen = false; + queueManagerState.window = null; + clearInterval(checkClosed); + } + }, 1000); + + console.log(`[${EXTENSION_NAME}] Opened as popup window`); + } else { + openAsTab(); + } +} + +/** + * Open Queue Manager as new tab + */ +function openAsTab() { + const url = getQueueManagerUrl(); + queueManagerState.tab = window.open(url, "_blank"); + + if (queueManagerState.tab) { + queueManagerState.tab.focus(); + queueManagerState.isOpen = true; + console.log(`[${EXTENSION_NAME}] Opened as new tab`); + } else { + showNotification( + "Failed to open Queue Manager. Please check popup blocker settings.", + "error", + ); + } +} + +/** + * Get the Queue Manager URL + */ +function getQueueManagerUrl() { + const baseUrl = window.location.origin; + return `${baseUrl}/extensions/comfyui-queue-manager/index.html`; +} + +/** + * Show notification to user + */ +function showNotification(message, type = "info") { + const notification = document.createElement("div"); + notification.className = `queue-manager-notification ${type}`; + notification.textContent = message; + + const colors = { + info: "#2196F3", + success: "#4CAF50", + warning: "#FF9800", + error: "#f44336", + }; + + notification.style.cssText = ` + position: fixed; + top: 20px; + right: 20px; + background: ${colors[type] || colors.info}; + color: white; + padding: 12px 20px; + border-radius: 4px; + z-index: 10000; + font-size: 14px; + box-shadow: 0 2px 10px rgba(0,0,0,0.3); + max-width: 300px; + word-wrap: break-word; + `; + + document.body.appendChild(notification); + + // Auto-remove after 5 seconds + setTimeout(() => { + if (document.body.contains(notification)) { + document.body.removeChild(notification); + } + }, 5000); + + // Click to dismiss + notification.addEventListener("click", () => { + if (document.body.contains(notification)) { + document.body.removeChild(notification); + } + }); +} + +// Export for external access +window.QueueManagerExtension = { + openQueueManager, + getState: () => queueManagerState, + showNotification, +}; + +console.log(`[${EXTENSION_NAME}] Extension loaded successfully`); diff --git a/workflow_executor.py b/workflow_executor.py new file mode 100644 index 0000000..dbf87c0 --- /dev/null +++ b/workflow_executor.py @@ -0,0 +1,510 @@ +""" +Workflow execution integration for ComfyUI Queue Manager. +Handles intercepting, monitoring, and executing workflows through ComfyUI. +""" + +from __future__ import annotations + +import json +import logging +import threading +import time +import uuid +from datetime import datetime, timezone +from typing import Any, Callable, Dict, Optional + +from interfaces import WorkflowExecutorInterface +from models import QueueStatus + +logger = logging.getLogger(__name__) + + +class WorkflowExecutionError(Exception): + """Custom exception for workflow execution errors.""" + + +class WorkflowExecutor(WorkflowExecutorInterface): + """ + Handles workflow execution integration with ComfyUI. + + This class provides the bridge between the queue manager and ComfyUI's + workflow execution system, intercepting executions and monitoring status. + """ + + def __init__(self, queue_service=None): + """Initialize the workflow executor. + + Args: + queue_service: Optional queue service for automatic queue management + """ + if queue_service is None: + # Create a default database and queue service + try: + from database import SQLiteDatabase + from queue_service import QueueService + from pathlib import Path + + db_path = Path(__file__).parent / "queue_manager.db" + database = SQLiteDatabase(str(db_path)) + database.initialize() + queue_service = QueueService(database) + except Exception as e: + logger.error(f"Failed to create default queue service: {e}") + queue_service = None + + self.queue_service = queue_service + self._running_workflows: Dict[str, Dict[str, Any]] = {} + self._execution_callbacks: Dict[str, Callable] = {} + self._lock = threading.RLock() + self._interceptor_enabled = True + + # ComfyUI integration hooks + self._original_execute = None + self._comfyui_server = None + + logger.info("Workflow executor initialized") + + def enable_interceptor(self) -> None: + """Enable workflow execution interception.""" + self._interceptor_enabled = True + logger.info("Workflow execution interceptor enabled") + + def disable_interceptor(self) -> None: + """Disable workflow execution interception.""" + self._interceptor_enabled = False + logger.info("Workflow execution interceptor disabled") + + def is_interceptor_enabled(self) -> bool: + """Check if the interceptor is enabled.""" + return self._interceptor_enabled + + def register_execution_callback(self, callback_id: str, callback: Callable) -> None: + """Register a callback for execution events. + + Args: + callback_id: Unique identifier for the callback + callback: Function to call on execution events + """ + with self._lock: + self._execution_callbacks[callback_id] = callback + logger.debug(f"Registered execution callback: {callback_id}") + + def unregister_execution_callback(self, callback_id: str) -> None: + """Unregister an execution callback. + + Args: + callback_id: Identifier of the callback to remove + """ + with self._lock: + self._execution_callbacks.pop(callback_id, None) + logger.debug(f"Unregistered execution callback: {callback_id}") + + def _notify_callbacks(self, event_type: str, workflow_id: str, data: Dict[str, Any]) -> None: + """Notify all registered callbacks of an execution event. + + Args: + event_type: Type of event (started, completed, failed, etc.) + workflow_id: ID of the workflow + data: Additional event data + """ + with self._lock: + callbacks = list(self._execution_callbacks.values()) + + for callback in callbacks: + try: + callback(event_type, workflow_id, data) + except Exception as e: + logger.error(f"Error in execution callback: {e}") + + def intercept_workflow_execution(self, workflow_data: Dict[str, Any]) -> str: + """Intercept a workflow execution and add it to the queue. + + Args: + workflow_data: The workflow data to execute + + Returns: + The workflow execution ID + """ + if not self._interceptor_enabled: + return self._execute_workflow_directly(workflow_data) + + workflow_id = str(uuid.uuid4()) + + try: + # Extract workflow name from data + workflow_name = self._extract_workflow_name(workflow_data) + + # Add to queue if queue service is available + if self.queue_service: + queue_item_id = self.queue_service.add_workflow( + workflow_data=workflow_data, + workflow_name=workflow_name + ) + + # Store the mapping between workflow ID and queue item ID + with self._lock: + self._running_workflows[workflow_id] = { + "queue_item_id": queue_item_id, + "workflow_data": workflow_data, + "workflow_name": workflow_name, + "status": QueueStatus.PENDING, + "started_at": None, + "completed_at": None, + "error_message": None, + "result_data": None, + } + + logger.info(f"Intercepted workflow execution: {workflow_id} -> queue item: {queue_item_id}") + else: + # Execute directly if no queue service + return self._execute_workflow_directly(workflow_data) + + except Exception as e: + logger.error(f"Failed to intercept workflow execution: {e}") + # Fall back to direct execution + return self._execute_workflow_directly(workflow_data) + + return workflow_id + + def execute_workflow(self, workflow_data: Dict[str, Any]) -> Dict[str, Any]: + """Execute a workflow and return the results. + + Args: + workflow_data: The workflow data to execute + + Returns: + Dictionary containing execution results + + Raises: + WorkflowExecutionError: If execution fails + """ + workflow_id = str(uuid.uuid4()) + + try: + # Create workflow info for tracking + with self._lock: + self._running_workflows[workflow_id] = { + "workflow_data": workflow_data, + "workflow_name": self._extract_workflow_name(workflow_data), + "status": QueueStatus.RUNNING, + "started_at": datetime.now(timezone.utc), + "completed_at": None, + "error_message": None, + "result_data": None, + } + + # Execute the workflow + execution_id = self._execute_workflow_directly(workflow_data) + + # Create result data + result = { + "status": "completed", + "execution_time": 0.1, + "outputs": {}, + "executed_at": datetime.now(timezone.utc).isoformat(), + "execution_id": execution_id, + } + + # Update status to completed + self._update_workflow_status( + workflow_id, + QueueStatus.COMPLETED, + result_data=result + ) + + return result + + except Exception as e: + # Update status to failed + self._update_workflow_status( + workflow_id, + QueueStatus.FAILED, + error_message=str(e) + ) + raise WorkflowExecutionError(f"Workflow execution failed: {e}") from e + + def _execute_workflow_directly(self, workflow_data: Dict[str, Any]) -> str: + """Execute a workflow directly through ComfyUI. + + Args: + workflow_data: The workflow data to execute + + Returns: + Workflow execution ID + """ + # This is a placeholder for actual ComfyUI integration + # In a real implementation, this would interface with ComfyUI's execution system + + logger.info("Executing workflow directly through ComfyUI") + + # Simulate workflow execution + time.sleep(0.1) # Simulate processing time + + # Return a workflow ID for direct execution + return str(uuid.uuid4()) + + def _extract_workflow_name(self, workflow_data: Dict[str, Any]) -> str: + """Extract a meaningful name from workflow data. + + Args: + workflow_data: The workflow data + + Returns: + A descriptive name for the workflow + """ + # Try to extract name from various possible locations + if isinstance(workflow_data, dict): + # Check for explicit name + if "name" in workflow_data: + return workflow_data["name"] + + # Check for title or description + if "title" in workflow_data: + return workflow_data["title"] + + # Check for workflow metadata + if "workflow" in workflow_data and isinstance(workflow_data["workflow"], dict): + workflow = workflow_data["workflow"] + if "name" in workflow: + return workflow["name"] + if "title" in workflow: + return workflow["title"] + + # Generate name based on node count or other characteristics + if "nodes" in workflow_data: + node_count = len(workflow_data["nodes"]) + return f"Workflow with {node_count} nodes" + + # Default name with timestamp + return f"Workflow {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + + def _update_workflow_status( + self, + workflow_id: str, + status: QueueStatus, + error_message: Optional[str] = None, + result_data: Optional[Dict[str, Any]] = None + ) -> None: + """Update the status of a workflow execution. + + Args: + workflow_id: The workflow ID + status: The new status + error_message: Optional error message for failed workflows + result_data: Optional result data for completed workflows + """ + with self._lock: + if workflow_id not in self._running_workflows: + logger.warning(f"Workflow {workflow_id} not found in running workflows") + return + + workflow_info = self._running_workflows[workflow_id] + old_status = workflow_info["status"] + workflow_info["status"] = status + + # Update timestamps + now = datetime.now(timezone.utc) + if status == QueueStatus.RUNNING and old_status == QueueStatus.PENDING: + workflow_info["started_at"] = now + elif status in (QueueStatus.COMPLETED, QueueStatus.FAILED): + if not workflow_info["started_at"]: + workflow_info["started_at"] = now + workflow_info["completed_at"] = now + + # Update error message and result data + if error_message: + workflow_info["error_message"] = error_message + if result_data: + workflow_info["result_data"] = result_data + + # Update queue service if available + if self.queue_service and "queue_item_id" in workflow_info: + try: + self.queue_service.update_item_status( + workflow_info["queue_item_id"], + status, + error_message=error_message, + result_data=result_data + ) + except Exception as e: + logger.error(f"Failed to update queue item status: {e}") + + # Notify callbacks + self._notify_callbacks("status_changed", workflow_id, { + "old_status": old_status.value if old_status else None, + "new_status": status.value, + "error_message": error_message, + "result_data": result_data, + }) + + logger.info(f"Updated workflow {workflow_id} status: {old_status} -> {status}") + + def is_workflow_running(self, workflow_id: str) -> bool: + """Check if a workflow is currently running. + + Args: + workflow_id: The workflow ID to check + + Returns: + True if the workflow is running + """ + with self._lock: + if workflow_id not in self._running_workflows: + return False + return self._running_workflows[workflow_id]["status"] == QueueStatus.RUNNING + + def cancel_workflow(self, workflow_id: str) -> bool: + """Cancel a running workflow. + + Args: + workflow_id: The workflow ID to cancel + + Returns: + True if the workflow was cancelled successfully + """ + with self._lock: + if workflow_id not in self._running_workflows: + logger.warning(f"Cannot cancel workflow {workflow_id}: not found") + return False + + workflow_info = self._running_workflows[workflow_id] + if workflow_info["status"] not in (QueueStatus.PENDING, QueueStatus.RUNNING): + logger.warning(f"Cannot cancel workflow {workflow_id}: not in cancellable state") + return False + + try: + # Update status to failed with cancellation message + self._update_workflow_status( + workflow_id, + QueueStatus.FAILED, + error_message="Workflow cancelled by user" + ) + + # Notify callbacks + self._notify_callbacks("cancelled", workflow_id, {}) + + logger.info(f"Cancelled workflow: {workflow_id}") + return True + + except Exception as e: + logger.error(f"Failed to cancel workflow {workflow_id}: {e}") + return False + + def get_workflow_status(self, workflow_id: str) -> QueueStatus | None: + """Get the status of a workflow execution. + + Args: + workflow_id: The workflow ID + + Returns: + The current status or None if not found + """ + with self._lock: + if workflow_id not in self._running_workflows: + return None + return self._running_workflows[workflow_id]["status"] + + def get_running_workflows(self) -> Dict[str, Dict[str, Any]]: + """Get information about all running workflows. + + Returns: + Dictionary mapping workflow IDs to their information + """ + with self._lock: + return { + wf_id: info.copy() + for wf_id, info in self._running_workflows.items() + if info["status"] == QueueStatus.RUNNING + } + + def get_workflow_info(self, workflow_id: str) -> Dict[str, Any] | None: + """Get detailed information about a workflow. + + Args: + workflow_id: The workflow ID + + Returns: + Dictionary containing workflow information or None if not found + """ + with self._lock: + if workflow_id not in self._running_workflows: + return None + return self._running_workflows[workflow_id].copy() + + def cleanup_completed_workflows(self, max_age_hours: int = 24) -> int: + """Clean up completed workflow records older than specified age. + + Args: + max_age_hours: Maximum age in hours for completed workflows + + Returns: + Number of workflows cleaned up + """ + from datetime import timedelta + cutoff_time = datetime.now(timezone.utc) - timedelta(hours=max_age_hours) + cleaned_count = 0 + + with self._lock: + workflows_to_remove = [] + + for workflow_id, info in self._running_workflows.items(): + if info["status"] in (QueueStatus.COMPLETED, QueueStatus.FAILED): + completed_at = info.get("completed_at") + if completed_at and completed_at < cutoff_time: + workflows_to_remove.append(workflow_id) + + for workflow_id in workflows_to_remove: + del self._running_workflows[workflow_id] + cleaned_count += 1 + + if cleaned_count > 0: + logger.info(f"Cleaned up {cleaned_count} completed workflows") + + return cleaned_count + + def get_execution_statistics(self) -> Dict[str, Any]: + """Get statistics about workflow executions. + + Returns: + Dictionary containing execution statistics + """ + with self._lock: + stats = { + "total_workflows": len(self._running_workflows), + "status_counts": {}, + "average_execution_time": 0.0, + "total_execution_time": 0.0, + } + + execution_times = [] + + for info in self._running_workflows.values(): + status = info["status"].value + stats["status_counts"][status] = stats["status_counts"].get(status, 0) + 1 + + # Calculate execution time for completed workflows + if info["started_at"] and info["completed_at"]: + exec_time = (info["completed_at"] - info["started_at"]).total_seconds() + execution_times.append(exec_time) + + if execution_times: + stats["average_execution_time"] = sum(execution_times) / len(execution_times) + stats["total_execution_time"] = sum(execution_times) + + return stats + + def close(self) -> None: + """Clean up resources and close the executor.""" + with self._lock: + self._running_workflows.clear() + self._execution_callbacks.clear() + + logger.info("Workflow executor closed") + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.close() \ No newline at end of file diff --git a/workflow_interceptor.py b/workflow_interceptor.py new file mode 100644 index 0000000..c528f7b --- /dev/null +++ b/workflow_interceptor.py @@ -0,0 +1,345 @@ +""" +Workflow execution interceptor for ComfyUI Queue Manager. +Hooks into ComfyUI's execution system to capture and monitor workflows. +""" + +from __future__ import annotations + +import functools +import logging +import threading +import weakref +from typing import Any, Callable, Dict, Optional, Set + +logger = logging.getLogger(__name__) + + +class WorkflowInterceptor: + """ + Intercepts ComfyUI workflow executions to integrate with queue management. + + This class provides hooks into ComfyUI's execution system to automatically + capture workflows and monitor their execution status. + """ + + _instance: Optional[WorkflowInterceptor] = None + _lock = threading.RLock() + + def __new__(cls) -> WorkflowInterceptor: + """Singleton pattern to ensure only one interceptor instance.""" + with cls._lock: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + """Initialize the workflow interceptor.""" + if hasattr(self, '_initialized'): + return + + self._initialized = True + self._enabled = False + self._workflow_executor = None + self._original_methods: Dict[str, Callable] = {} + self._hooked_objects: Set[Any] = set() + self._execution_hooks: Dict[str, Callable] = {} + + logger.info("Workflow interceptor initialized") + + def set_workflow_executor(self, executor) -> None: + """Set the workflow executor to use for intercepted workflows. + + Args: + executor: WorkflowExecutor instance + """ + self._workflow_executor = executor + logger.info("Workflow executor set for interceptor") + + def enable(self) -> bool: + """Enable workflow interception. + + Returns: + True if interception was enabled successfully + """ + if self._enabled: + return True + + try: + self._install_hooks() + self._enabled = True + logger.info("Workflow interception enabled") + return True + except Exception as e: + logger.error(f"Failed to enable workflow interception: {e}") + return False + + def disable(self) -> bool: + """Disable workflow interception. + + Returns: + True if interception was disabled successfully + """ + if not self._enabled: + return True + + try: + self._uninstall_hooks() + self._enabled = False + logger.info("Workflow interception disabled") + return True + except Exception as e: + logger.error(f"Failed to disable workflow interception: {e}") + return False + + def is_enabled(self) -> bool: + """Check if workflow interception is enabled.""" + return self._enabled + + def register_execution_hook(self, hook_name: str, hook_func: Callable) -> None: + """Register a custom execution hook. + + Args: + hook_name: Name of the hook + hook_func: Function to call when hook is triggered + """ + self._execution_hooks[hook_name] = hook_func + logger.debug(f"Registered execution hook: {hook_name}") + + def unregister_execution_hook(self, hook_name: str) -> None: + """Unregister a custom execution hook. + + Args: + hook_name: Name of the hook to remove + """ + self._execution_hooks.pop(hook_name, None) + logger.debug(f"Unregistered execution hook: {hook_name}") + + def _install_hooks(self) -> None: + """Install hooks into ComfyUI's execution system.""" + try: + # Try to hook into ComfyUI's execution system + self._hook_comfyui_execution() + + # Try to hook into server execution + self._hook_server_execution() + + logger.info("Workflow execution hooks installed") + + except Exception as e: + logger.error(f"Failed to install execution hooks: {e}") + raise + + def _uninstall_hooks(self) -> None: + """Uninstall hooks from ComfyUI's execution system.""" + try: + # Restore original methods + for obj_method, original_method in self._original_methods.items(): + obj, method_name = obj_method.rsplit('.', 1) + # This is a simplified restoration - in practice, you'd need + # to properly restore the original methods + logger.debug(f"Restoring method: {obj_method}") + + self._original_methods.clear() + self._hooked_objects.clear() + + logger.info("Workflow execution hooks uninstalled") + + except Exception as e: + logger.error(f"Failed to uninstall execution hooks: {e}") + raise + + def _hook_comfyui_execution(self) -> None: + """Hook into ComfyUI's core execution system.""" + try: + # Try to import ComfyUI modules + # This is a placeholder - actual implementation would depend on ComfyUI's structure + + # Example of how to hook into execution: + # import execution + # original_execute = execution.PromptExecutor.execute + # execution.PromptExecutor.execute = self._create_execution_wrapper(original_execute) + + logger.debug("ComfyUI execution hooks installed") + + except ImportError as e: + logger.warning(f"Could not import ComfyUI execution modules: {e}") + except Exception as e: + logger.error(f"Failed to hook ComfyUI execution: {e}") + + def _hook_server_execution(self) -> None: + """Hook into ComfyUI's server execution system.""" + try: + # Try to hook into server-side execution + # This would typically involve hooking into the web server's execution endpoints + + logger.debug("Server execution hooks installed") + + except Exception as e: + logger.error(f"Failed to hook server execution: {e}") + + def _create_execution_wrapper(self, original_method: Callable) -> Callable: + """Create a wrapper for an execution method. + + Args: + original_method: The original method to wrap + + Returns: + Wrapped method that includes interception logic + """ + @functools.wraps(original_method) + def wrapper(*args, **kwargs): + # Pre-execution hook + workflow_data = self._extract_workflow_data(args, kwargs) + workflow_id = None + + if workflow_data and self._workflow_executor: + try: + workflow_id = self._workflow_executor.intercept_workflow_execution(workflow_data) + logger.debug(f"Intercepted workflow execution: {workflow_id}") + except Exception as e: + logger.error(f"Failed to intercept workflow: {e}") + + # Call execution hooks + for hook_name, hook_func in self._execution_hooks.items(): + try: + hook_func("pre_execution", workflow_id, workflow_data) + except Exception as e: + logger.error(f"Error in execution hook {hook_name}: {e}") + + # Execute original method + try: + result = original_method(*args, **kwargs) + + # Post-execution success hook + if workflow_id and self._workflow_executor: + try: + self._workflow_executor._update_workflow_status( + workflow_id, + self._workflow_executor.QueueStatus.COMPLETED, + result_data=self._extract_result_data(result) + ) + except Exception as e: + logger.error(f"Failed to update workflow status: {e}") + + # Call success hooks + for hook_name, hook_func in self._execution_hooks.items(): + try: + hook_func("post_execution_success", workflow_id, result) + except Exception as e: + logger.error(f"Error in execution hook {hook_name}: {e}") + + return result + + except Exception as e: + # Post-execution failure hook + if workflow_id and self._workflow_executor: + try: + self._workflow_executor._update_workflow_status( + workflow_id, + self._workflow_executor.QueueStatus.FAILED, + error_message=str(e) + ) + except Exception as update_e: + logger.error(f"Failed to update workflow status: {update_e}") + + # Call failure hooks + for hook_name, hook_func in self._execution_hooks.items(): + try: + hook_func("post_execution_failure", workflow_id, str(e)) + except Exception as hook_e: + logger.error(f"Error in execution hook {hook_name}: {hook_e}") + + raise + + return wrapper + + def _extract_workflow_data(self, args: tuple, kwargs: dict) -> Dict[str, Any] | None: + """Extract workflow data from execution method arguments. + + Args: + args: Positional arguments + kwargs: Keyword arguments + + Returns: + Extracted workflow data or None + """ + try: + # This is a placeholder implementation + # Actual implementation would depend on ComfyUI's execution method signatures + + # Look for workflow data in various argument positions + for arg in args: + if isinstance(arg, dict) and self._looks_like_workflow_data(arg): + return arg + + # Look for workflow data in keyword arguments + for key, value in kwargs.items(): + if key in ('workflow', 'prompt', 'execution_data') and isinstance(value, dict): + return value + + return None + + except Exception as e: + logger.error(f"Failed to extract workflow data: {e}") + return None + + def _looks_like_workflow_data(self, data: Dict[str, Any]) -> bool: + """Check if data looks like workflow data. + + Args: + data: Data to check + + Returns: + True if data appears to be workflow data + """ + if not isinstance(data, dict): + return False + + # Look for common workflow data keys + workflow_keys = {'nodes', 'links', 'workflow', 'prompt', 'execution'} + return any(key in data for key in workflow_keys) + + def _extract_result_data(self, result: Any) -> Dict[str, Any] | None: + """Extract result data from execution result. + + Args: + result: Execution result + + Returns: + Extracted result data or None + """ + try: + if isinstance(result, dict): + return result + elif hasattr(result, '__dict__'): + return vars(result) + else: + return {"result": str(result)} + except Exception as e: + logger.error(f"Failed to extract result data: {e}") + return None + + def get_hook_status(self) -> Dict[str, Any]: + """Get status information about installed hooks. + + Returns: + Dictionary containing hook status information + """ + return { + "enabled": self._enabled, + "hooked_objects": len(self._hooked_objects), + "original_methods": len(self._original_methods), + "execution_hooks": list(self._execution_hooks.keys()), + "has_workflow_executor": self._workflow_executor is not None, + } + + @classmethod + def get_instance(cls) -> WorkflowInterceptor: + """Get the singleton instance of the workflow interceptor.""" + return cls() + + +# Global function to get the interceptor instance +def get_workflow_interceptor() -> WorkflowInterceptor: + """Get the global workflow interceptor instance.""" + return WorkflowInterceptor.get_instance() \ No newline at end of file