comfy-discord-web/user_state_registry.py

"""
user_state_registry.py
======================

Per-user workflow state + template registry.

Each web-UI user gets their own isolated WorkflowStateManager (persisted to
``user_settings/<user_label>.json``) and workflow template.

New users (no saved file) fall back to the global default workflow template
loaded at startup (WORKFLOW_FILE env var or last-used workflow from the
global Discord state manager).

Discord continues to use the shared global state/workflow manager — this
registry is only used by the web UI layer.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Dict, Optional

from workflow_state import WorkflowStateManager

logger = logging.getLogger(__name__)

_PROJECT_ROOT = Path(__file__).resolve().parent


class UserStateRegistry:
    """
    Per-user isolated workflow state and template store.

    Parameters
    ----------
    settings_dir : Path
        Directory where per-user state files are stored.  Created automatically.
    default_workflow : Optional[dict]
        The global default workflow template.  Used when a user has no saved
        ``last_workflow_file``, or the file no longer exists.
    """

    def __init__(
        self,
        settings_dir: Path,
        default_workflow: Optional[dict[str, Any]] = None,
    ) -> None:
        self._settings_dir = settings_dir
        self._settings_dir.mkdir(parents=True, exist_ok=True)
        self._default_workflow: Optional[dict[str, Any]] = default_workflow
        # user_label → WorkflowStateManager
        self._managers: Dict[str, WorkflowStateManager] = {}
        # user_label → workflow template dict (or None)
        self._templates: Dict[str, Optional[dict[str, Any]]] = {}

    # ------------------------------------------------------------------
    # Default workflow
    # ------------------------------------------------------------------

    def set_default_workflow(self, template: Optional[dict[str, Any]]) -> None:
        """Update the global fallback workflow (called when bot workflow changes)."""
        self._default_workflow = template

    # ------------------------------------------------------------------
    # User access
    # ------------------------------------------------------------------

    def get_state_manager(self, user_label: str) -> WorkflowStateManager:
        """Return (or lazily create) the WorkflowStateManager for a user."""
        if user_label not in self._managers:
            self._init_user(user_label)
        return self._managers[user_label]

    def get_workflow_template(self, user_label: str) -> Optional[dict[str, Any]]:
        """Return the workflow template for a user, or None if not loaded."""
        if user_label not in self._managers:
            self._init_user(user_label)
        return self._templates.get(user_label)

    def set_workflow(
        self, user_label: str, template: dict[str, Any], filename: str
    ) -> None:
        """
        Store a workflow template for a user and persist the filename.

        Clears existing overrides (matches the behaviour of loading a new
        workflow via the global state manager).
        """
        if user_label not in self._managers:
            self._init_user(user_label)
        sm = self._managers[user_label]
        sm.clear_overrides()
        sm.set_last_workflow_file(filename)
        self._templates[user_label] = template
        logger.debug(
            "UserStateRegistry: set workflow '%s' for user '%s'", filename, user_label
        )

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _init_user(self, user_label: str) -> None:
        """
        Initialise state manager and workflow template for a new user.

        1. Create WorkflowStateManager with per-user state file.
        2. If the file recorded a last_workflow_file, try to load it.
        3. Fall back to the global default template.
        """
        state_file = str(self._settings_dir / f"{user_label}.json")
        sm = WorkflowStateManager(state_file=state_file)
        self._managers[user_label] = sm

        # Try to restore the last workflow this user loaded
        last_wf = sm.get_last_workflow_file()
        template: Optional[dict[str, Any]] = None
        if last_wf:
            wf_path = _PROJECT_ROOT / "workflows" / last_wf
            if wf_path.exists():
                try:
                    with open(wf_path, "r", encoding="utf-8") as f:
                        template = json.load(f)
                    logger.debug(
                        "UserStateRegistry: restored workflow '%s' for user '%s'",
                        last_wf,
                        user_label,
                    )
                except Exception as exc:
                    logger.warning(
                        "UserStateRegistry: could not load '%s' for user '%s': %s",
                        last_wf,
                        user_label,
                        exc,
                    )
            else:
                logger.debug(
                    "UserStateRegistry: last workflow '%s' missing for user '%s'; using default",
                    last_wf,
                    user_label,
                )

        if template is None:
            template = self._default_workflow
        self._templates[user_label] = template