[feat] Implement comprehensive batch tracking and OpenAPI-driven data models

Enhances ComfyUI Manager with robust batch execution tracking and unified data model architecture: - Implemented automatic batch history serialization with before/after system state snapshots - Added comprehensive state management capturing installed nodes, models, and ComfyUI version info - Enhanced task queue with proper client ID handling and WebSocket notifications - Migrated all data models to OpenAPI-generated Pydantic models for consistency - Added documentation for new TaskQueue methods (done_count, total_count, finalize) - Fixed 64 linting errors with proper imports and code cleanup Technical improvements: - All models now auto-generated from openapi.yaml ensuring API/implementation consistency - Batch tracking captures complete system state at operation start and completion - Enhanced REST endpoints with comprehensive documentation - Removed manual model files in favor of single source of truth - Added helper methods for system state capture and batch lifecycle management
2025-06-08 01:18:14 -07:00
parent 35eddc2965
commit 49549ddcb8
6 changed files with 1679 additions and 629 deletions
--- a/comfyui_manager/data_models/README.md
+++ b/comfyui_manager/data_models/README.md
@@ -0,0 +1,67 @@
+# Data Models
+
+This directory contains Pydantic models for ComfyUI Manager, providing type safety, validation, and serialization for the API and internal data structures.
+
+## Overview
+
+- `generated_models.py` - All models auto-generated from OpenAPI spec
+- `__init__.py` - Package exports for all models
+
+**Note**: All models are now auto-generated from the OpenAPI specification. Manual model files (`task_queue.py`, `state_management.py`) have been deprecated in favor of a single source of truth.
+
+## Generating Types from OpenAPI
+
+The state management models are automatically generated from the OpenAPI specification using `datamodel-codegen`. This ensures type safety and consistency between the API specification and the Python code.
+
+### Prerequisites
+
+Install the code generator:
+```bash
+pipx install datamodel-code-generator
+```
+
+### Generation Command
+
+To regenerate all models after updating the OpenAPI spec:
+
+```bash
+datamodel-codegen \
+  --use-subclass-enum \
+  --field-constraints \
+  --strict-types bytes \
+  --input openapi.yaml \
+  --output comfyui_manager/data_models/generated_models.py \
+  --output-model-type pydantic_v2.BaseModel
+```
+
+### When to Regenerate
+
+You should regenerate the models when:
+
+1. **Adding new API endpoints** that return new data structures
+2. **Modifying existing schemas** in the OpenAPI specification  
+3. **Adding new state management features** that require new models
+
+### Important Notes
+
+- **Single source of truth**: All models are now generated from `openapi.yaml`
+- **No manual models**: All previously manual models have been migrated to the OpenAPI spec
+- **OpenAPI requirements**: New schemas must be referenced in API paths to be generated by datamodel-codegen
+- **Validation**: Always validate the OpenAPI spec before generation:
+  ```bash
+  python3 -c "import yaml; yaml.safe_load(open('openapi.yaml'))"
+  ```
+
+### Example: Adding New State Models
+
+1. Add your schema to `openapi.yaml` under `components/schemas/`
+2. Reference the schema in an API endpoint response
+3. Run the generation command above
+4. Update `__init__.py` to export the new models
+5. Import and use the models in your code
+
+### Troubleshooting
+
+- **Models not generated**: Ensure schemas are under `components/schemas/` (not `parameters/`)
+- **Missing models**: Verify schemas are referenced in at least one API path
+- **Import errors**: Check that new models are added to `__init__.py` exports
--- a/comfyui_manager/data_models/init.py
+++ b/comfyui_manager/data_models/init.py
@@ -3,24 +3,105 @@ Data models for ComfyUI Manager.

 This package contains Pydantic models used throughout the ComfyUI Manager
 for data validation, serialization, and type safety.
+
+All models are auto-generated from the OpenAPI specification to ensure
+consistency between the API and implementation.
 """

-from .task_queue import (
+from .generated_models import (
+    # Core Task Queue Models
    QueueTaskItem,
    TaskHistoryItem,
    TaskStateMessage,
+    TaskExecutionStatus,
+    
+    # WebSocket Message Models
    MessageTaskDone,
    MessageTaskStarted,
+    MessageTaskFailed,
    MessageUpdate,
    ManagerMessageName,
+    
+    # State Management Models
+    BatchExecutionRecord,
+    ComfyUISystemState,
+    BatchOperation,
+    InstalledNodeInfo,
+    InstalledModelInfo,
+    ComfyUIVersionInfo,
+    
+    # Other models
+    Kind,
+    StatusStr,
+    ManagerPackInfo,
+    ManagerPackInstalled,
+    SelectedVersion,
+    ManagerChannel,
+    ManagerDatabaseSource,
+    ManagerPackState,
+    ManagerPackInstallType,
+    ManagerPack,
+    InstallPackParams,
+    UpdateAllPacksParams,
+    QueueStatus,
+    ManagerMappings,
+    ModelMetadata,
+    NodePackageMetadata,
+    SnapshotItem,
+    Error,
+    InstalledPacksResponse,
+    HistoryResponse,
+    HistoryListResponse,
+    InstallType,
+    OperationType,
+    Result,
 )

 __all__ = [
+    # Core Task Queue Models
    "QueueTaskItem",
    "TaskHistoryItem",
    "TaskStateMessage",
+    "TaskExecutionStatus",
+    
+    # WebSocket Message Models
    "MessageTaskDone",
    "MessageTaskStarted",
+    "MessageTaskFailed",
    "MessageUpdate",
    "ManagerMessageName",
-]
+    
+    # State Management Models
+    "BatchExecutionRecord",
+    "ComfyUISystemState",
+    "BatchOperation",
+    "InstalledNodeInfo",
+    "InstalledModelInfo",
+    "ComfyUIVersionInfo",
+    
+    # Other models
+    "Kind",
+    "StatusStr",
+    "ManagerPackInfo",
+    "ManagerPackInstalled",
+    "SelectedVersion",
+    "ManagerChannel",
+    "ManagerDatabaseSource",
+    "ManagerPackState",
+    "ManagerPackInstallType",
+    "ManagerPack",
+    "InstallPackParams",
+    "UpdateAllPacksParams",
+    "QueueStatus",
+    "ManagerMappings",
+    "ModelMetadata",
+    "NodePackageMetadata",
+    "SnapshotItem",
+    "Error",
+    "InstalledPacksResponse",
+    "HistoryResponse",
+    "HistoryListResponse",
+    "InstallType",
+    "OperationType",
+    "Result",
+]
--- a/comfyui_manager/data_models/generated_models.py
+++ b/comfyui_manager/data_models/generated_models.py
@@ -0,0 +1,417 @@
+# generated by datamodel-codegen:
+#   filename:  openapi.yaml
+#   timestamp: 2025-06-08T08:07:38+00:00
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, RootModel
+
+
+class Kind(str, Enum):
+    install = 'install'
+    uninstall = 'uninstall'
+    update = 'update'
+    update_all = 'update-all'
+    update_comfyui = 'update-comfyui'
+    fix = 'fix'
+    disable = 'disable'
+    enable = 'enable'
+    install_model = 'install-model'
+
+
+class QueueTaskItem(BaseModel):
+    ui_id: str = Field(..., description='Unique identifier for the task')
+    client_id: str = Field(..., description='Client identifier that initiated the task')
+    kind: Kind = Field(..., description='Type of task being performed')
+
+
+class StatusStr(str, Enum):
+    success = 'success'
+    error = 'error'
+    skip = 'skip'
+
+
+class TaskExecutionStatus(BaseModel):
+    status_str: StatusStr = Field(..., description='Overall task execution status')
+    completed: bool = Field(..., description='Whether the task completed')
+    messages: List[str] = Field(..., description='Additional status messages')
+
+
+class ManagerMessageName(str, Enum):
+    cm_task_completed = 'cm-task-completed'
+    cm_task_started = 'cm-task-started'
+    cm_queue_status = 'cm-queue-status'
+
+
+class ManagerPackInfo(BaseModel):
+    id: str = Field(
+        ...,
+        description='Either github-author/github-repo or name of pack from the registry',
+    )
+    version: str = Field(..., description='Semantic version or Git commit hash')
+    ui_id: Optional[str] = Field(None, description='Task ID - generated internally')
+
+
+class ManagerPackInstalled(BaseModel):
+    ver: str = Field(
+        ...,
+        description='The version of the pack that is installed (Git commit hash or semantic version)',
+    )
+    cnr_id: Optional[str] = Field(
+        None, description='The name of the pack if installed from the registry'
+    )
+    aux_id: Optional[str] = Field(
+        None,
+        description='The name of the pack if installed from github (author/repo-name format)',
+    )
+    enabled: bool = Field(..., description='Whether the pack is enabled')
+
+
+class SelectedVersion(str, Enum):
+    latest = 'latest'
+    nightly = 'nightly'
+
+
+class ManagerChannel(str, Enum):
+    default = 'default'
+    recent = 'recent'
+    legacy = 'legacy'
+    forked = 'forked'
+    dev = 'dev'
+    tutorial = 'tutorial'
+
+
+class ManagerDatabaseSource(str, Enum):
+    remote = 'remote'
+    local = 'local'
+    cache = 'cache'
+
+
+class ManagerPackState(str, Enum):
+    installed = 'installed'
+    disabled = 'disabled'
+    not_installed = 'not_installed'
+    import_failed = 'import_failed'
+    needs_update = 'needs_update'
+
+
+class ManagerPackInstallType(str, Enum):
+    git_clone = 'git-clone'
+    copy = 'copy'
+    cnr = 'cnr'
+
+
+class UpdateState(str, Enum):
+    false = 'false'
+    true = 'true'
+
+
+class ManagerPack(ManagerPackInfo):
+    author: Optional[str] = Field(
+        None, description="Pack author name or 'Unclaimed' if added via GitHub crawl"
+    )
+    files: Optional[List[str]] = Field(None, description='Files included in the pack')
+    reference: Optional[str] = Field(
+        None, description='The type of installation reference'
+    )
+    title: Optional[str] = Field(None, description='The display name of the pack')
+    cnr_latest: Optional[SelectedVersion] = None
+    repository: Optional[str] = Field(None, description='GitHub repository URL')
+    state: Optional[ManagerPackState] = None
+    update_state: Optional[UpdateState] = Field(
+        None, alias='update-state', description='Update availability status'
+    )
+    stars: Optional[int] = Field(None, description='GitHub stars count')
+    last_update: Optional[datetime] = Field(None, description='Last update timestamp')
+    health: Optional[str] = Field(None, description='Health status of the pack')
+    description: Optional[str] = Field(None, description='Pack description')
+    trust: Optional[bool] = Field(None, description='Whether the pack is trusted')
+    install_type: Optional[ManagerPackInstallType] = None
+
+
+class InstallPackParams(ManagerPackInfo):
+    selected_version: Union[str, SelectedVersion] = Field(
+        ..., description='Semantic version, Git commit hash, latest, or nightly'
+    )
+    repository: Optional[str] = Field(
+        None,
+        description='GitHub repository URL (required if selected_version is nightly)',
+    )
+    pip: Optional[List[str]] = Field(None, description='PyPi dependency names')
+    mode: ManagerDatabaseSource
+    channel: ManagerChannel
+    skip_post_install: Optional[bool] = Field(
+        None, description='Whether to skip post-installation steps'
+    )
+
+
+class UpdateAllPacksParams(BaseModel):
+    mode: Optional[ManagerDatabaseSource] = None
+    ui_id: Optional[str] = Field(None, description='Task ID - generated internally')
+
+
+class QueueStatus(BaseModel):
+    total_count: int = Field(
+        ..., description='Total number of tasks (pending + running)'
+    )
+    done_count: int = Field(..., description='Number of completed tasks')
+    in_progress_count: int = Field(..., description='Number of tasks currently running')
+    pending_count: Optional[int] = Field(
+        None, description='Number of tasks waiting to be executed'
+    )
+    is_processing: bool = Field(..., description='Whether the task worker is active')
+    client_id: Optional[str] = Field(
+        None, description='Client ID (when filtered by client)'
+    )
+
+
+class ManagerMapping(BaseModel):
+    title_aux: Optional[str] = Field(None, description='The display name of the pack')
+
+
+class ManagerMappings(
+    RootModel[Optional[Dict[str, List[Union[List[str], ManagerMapping]]]]]
+):
+    root: Optional[Dict[str, List[Union[List[str], ManagerMapping]]]] = None
+
+
+class ModelMetadata(BaseModel):
+    name: str = Field(..., description='Name of the model')
+    type: str = Field(..., description='Type of model')
+    base: Optional[str] = Field(None, description='Base model type')
+    save_path: Optional[str] = Field(None, description='Path for saving the model')
+    url: str = Field(..., description='Download URL')
+    filename: str = Field(..., description='Target filename')
+    ui_id: Optional[str] = Field(None, description='ID for UI reference')
+
+
+class InstallType(str, Enum):
+    git = 'git'
+    copy = 'copy'
+    pip = 'pip'
+
+
+class NodePackageMetadata(BaseModel):
+    title: Optional[str] = Field(None, description='Display name of the node package')
+    name: Optional[str] = Field(None, description='Repository/package name')
+    files: Optional[List[str]] = Field(None, description='Source URLs for the package')
+    description: Optional[str] = Field(
+        None, description='Description of the node package functionality'
+    )
+    install_type: Optional[InstallType] = Field(None, description='Installation method')
+    version: Optional[str] = Field(None, description='Version identifier')
+    id: Optional[str] = Field(
+        None, description='Unique identifier for the node package'
+    )
+    ui_id: Optional[str] = Field(None, description='ID for UI reference')
+    channel: Optional[str] = Field(None, description='Source channel')
+    mode: Optional[str] = Field(None, description='Source mode')
+
+
+class SnapshotItem(RootModel[str]):
+    root: str = Field(..., description='Name of the snapshot')
+
+
+class Error(BaseModel):
+    error: str = Field(..., description='Error message')
+
+
+class InstalledPacksResponse(RootModel[Optional[Dict[str, ManagerPackInstalled]]]):
+    root: Optional[Dict[str, ManagerPackInstalled]] = None
+
+
+class HistoryListResponse(BaseModel):
+    ids: Optional[List[str]] = Field(
+        None, description='List of available batch history IDs'
+    )
+
+
+class InstalledNodeInfo(BaseModel):
+    name: str = Field(..., description='Node package name')
+    version: str = Field(..., description='Installed version')
+    repository_url: Optional[str] = Field(None, description='Git repository URL')
+    install_method: str = Field(
+        ..., description='Installation method (cnr, git, pip, etc.)'
+    )
+    enabled: Optional[bool] = Field(
+        True, description='Whether the node is currently enabled'
+    )
+    install_date: Optional[datetime] = Field(
+        None, description='ISO timestamp of installation'
+    )
+
+
+class InstalledModelInfo(BaseModel):
+    name: str = Field(..., description='Model filename')
+    path: str = Field(..., description='Full path to model file')
+    type: str = Field(..., description='Model type (checkpoint, lora, vae, etc.)')
+    size_bytes: Optional[int] = Field(None, description='File size in bytes', ge=0)
+    hash: Optional[str] = Field(None, description='Model file hash for verification')
+    install_date: Optional[datetime] = Field(
+        None, description='ISO timestamp when added'
+    )
+
+
+class ComfyUIVersionInfo(BaseModel):
+    version: str = Field(..., description='ComfyUI version string')
+    commit_hash: Optional[str] = Field(None, description='Git commit hash')
+    branch: Optional[str] = Field(None, description='Git branch name')
+    is_stable: Optional[bool] = Field(
+        False, description='Whether this is a stable release'
+    )
+    last_updated: Optional[datetime] = Field(
+        None, description='ISO timestamp of last update'
+    )
+
+
+class OperationType(str, Enum):
+    install = 'install'
+    update = 'update'
+    uninstall = 'uninstall'
+    fix = 'fix'
+    disable = 'disable'
+    enable = 'enable'
+    install_model = 'install-model'
+
+
+class Result(str, Enum):
+    success = 'success'
+    failed = 'failed'
+    skipped = 'skipped'
+
+
+class BatchOperation(BaseModel):
+    operation_id: str = Field(..., description='Unique operation identifier')
+    operation_type: OperationType = Field(..., description='Type of operation')
+    target: str = Field(
+        ..., description='Target of the operation (node name, model name, etc.)'
+    )
+    target_version: Optional[str] = Field(
+        None, description='Target version for the operation'
+    )
+    result: Result = Field(..., description='Operation result')
+    error_message: Optional[str] = Field(
+        None, description='Error message if operation failed'
+    )
+    start_time: datetime = Field(
+        ..., description='ISO timestamp when operation started'
+    )
+    end_time: Optional[datetime] = Field(
+        None, description='ISO timestamp when operation completed'
+    )
+    client_id: Optional[str] = Field(
+        None, description='Client that initiated the operation'
+    )
+
+
+class ComfyUISystemState(BaseModel):
+    snapshot_time: datetime = Field(
+        ..., description='ISO timestamp when snapshot was taken'
+    )
+    comfyui_version: ComfyUIVersionInfo
+    frontend_version: Optional[str] = Field(
+        None, description='ComfyUI frontend version if available'
+    )
+    python_version: str = Field(..., description='Python interpreter version')
+    platform_info: str = Field(
+        ..., description='Operating system and platform information'
+    )
+    installed_nodes: Optional[Dict[str, InstalledNodeInfo]] = Field(
+        None, description='Map of installed node packages by name'
+    )
+    installed_models: Optional[Dict[str, InstalledModelInfo]] = Field(
+        None, description='Map of installed models by name'
+    )
+    manager_config: Optional[Dict[str, Any]] = Field(
+        None, description='ComfyUI Manager configuration settings'
+    )
+
+
+class BatchExecutionRecord(BaseModel):
+    batch_id: str = Field(..., description='Unique batch identifier')
+    start_time: datetime = Field(..., description='ISO timestamp when batch started')
+    end_time: Optional[datetime] = Field(
+        None, description='ISO timestamp when batch completed'
+    )
+    state_before: ComfyUISystemState
+    state_after: Optional[ComfyUISystemState] = Field(
+        None, description='System state after batch execution'
+    )
+    operations: Optional[List[BatchOperation]] = Field(
+        None, description='List of operations performed in this batch'
+    )
+    total_operations: Optional[int] = Field(
+        0, description='Total number of operations in batch', ge=0
+    )
+    successful_operations: Optional[int] = Field(
+        0, description='Number of successful operations', ge=0
+    )
+    failed_operations: Optional[int] = Field(
+        0, description='Number of failed operations', ge=0
+    )
+    skipped_operations: Optional[int] = Field(
+        0, description='Number of skipped operations', ge=0
+    )
+
+
+class TaskHistoryItem(BaseModel):
+    ui_id: str = Field(..., description='Unique identifier for the task')
+    client_id: str = Field(..., description='Client identifier that initiated the task')
+    kind: str = Field(..., description='Type of task that was performed')
+    timestamp: datetime = Field(..., description='ISO timestamp when task completed')
+    result: str = Field(..., description='Task result message or details')
+    status: Optional[TaskExecutionStatus] = None
+
+
+class TaskStateMessage(BaseModel):
+    history: Dict[str, TaskHistoryItem] = Field(
+        ..., description='Map of task IDs to their history items'
+    )
+    running_queue: List[QueueTaskItem] = Field(
+        ..., description='Currently executing tasks'
+    )
+    pending_queue: List[QueueTaskItem] = Field(
+        ..., description='Tasks waiting to be executed'
+    )
+
+
+class MessageTaskDone(BaseModel):
+    ui_id: str = Field(..., description='Task identifier')
+    result: str = Field(..., description='Task result message')
+    kind: str = Field(..., description='Type of task')
+    status: Optional[TaskExecutionStatus] = None
+    timestamp: datetime = Field(..., description='ISO timestamp when task completed')
+    state: TaskStateMessage
+
+
+class MessageTaskStarted(BaseModel):
+    ui_id: str = Field(..., description='Task identifier')
+    kind: str = Field(..., description='Type of task')
+    timestamp: datetime = Field(..., description='ISO timestamp when task started')
+    state: TaskStateMessage
+
+
+class MessageTaskFailed(BaseModel):
+    ui_id: str = Field(..., description='Task identifier')
+    error: str = Field(..., description='Error message')
+    kind: str = Field(..., description='Type of task')
+    timestamp: datetime = Field(..., description='ISO timestamp when task failed')
+    state: TaskStateMessage
+
+
+class MessageUpdate(
+    RootModel[Union[MessageTaskDone, MessageTaskStarted, MessageTaskFailed]]
+):
+    root: Union[MessageTaskDone, MessageTaskStarted, MessageTaskFailed] = Field(
+        ..., description='Union type for all possible WebSocket message updates'
+    )
+
+
+class HistoryResponse(BaseModel):
+    history: Optional[Dict[str, TaskHistoryItem]] = Field(
+        None, description='Map of task IDs to their history items'
+    )
--- a/comfyui_manager/data_models/task_queue.py
+++ b/comfyui_manager/data_models/task_queue.py
@@ -1,69 +0,0 @@
-"""
-Task queue data models for ComfyUI Manager.
-
-Contains Pydantic models for task queue management, WebSocket messaging,
-and task state tracking.
-"""
-
-from typing import Optional, Union, Dict
-from enum import Enum
-from pydantic import BaseModel
-
-
-class QueueTaskItem(BaseModel):
-    """Represents a task item in the queue."""
-
-    ui_id: str
-    client_id: str
-    kind: str
-
-
-class TaskHistoryItem(BaseModel):
-    """Represents a completed task in the history."""
-
-    ui_id: str
-    client_id: str
-    kind: str
-    timestamp: str
-    result: str
-    status: Optional[dict] = None
-
-
-class TaskStateMessage(BaseModel):
-    """Current state of the task queue system."""
-
-    history: Dict[str, TaskHistoryItem]
-    running_queue: list[QueueTaskItem]
-    pending_queue: list[QueueTaskItem]
-
-
-class MessageTaskDone(BaseModel):
-    """WebSocket message sent when a task completes."""
-
-    ui_id: str
-    result: str
-    kind: str
-    status: Optional[dict]
-    timestamp: str
-    state: TaskStateMessage
-
-
-class MessageTaskStarted(BaseModel):
-    """WebSocket message sent when a task starts."""
-
-    ui_id: str
-    kind: str
-    timestamp: str
-    state: TaskStateMessage
-
-
-# Union type for all possible WebSocket message updates
-MessageUpdate = Union[MessageTaskDone, MessageTaskStarted]
-
-
-class ManagerMessageName(Enum):
-    """WebSocket message type constants."""
-
-    TASK_DONE = "cm-task-completed"
-    TASK_STARTED = "cm-task-started"
-    STATUS = "cm-queue-status"