FabledSteward/steward/core/plugin_manager.py

# steward/core/plugin_manager.py
from __future__ import annotations
import logging
import os
import sys
from pathlib import Path
from typing import TYPE_CHECKING

import yaml
from packaging.version import Version

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# Track which plugins were successfully loaded so hot-reload knows
# whether to register a blueprint (new plugin) or skip it (already registered).
_LOADED_PLUGINS: set[str] = set()

# Track plugins that failed to load: name → human-readable reason.
_FAILED_PLUGINS: dict[str, str] = {}

# Nav entries contributed by plugins via the optional get_nav() export, so a
# plugin's UI gets a home in the sidebar. Each entry:
#   {"plugin": <name>, "label": str, "href": str, "section": str}
_PLUGIN_NAV: list[dict] = []

# Loaded plugin modules, by name, for the MCP server: a plugin contributes
# agent tools through an optional register_mcp_tools(register) export, which
# steward/mcp/server.py calls when it mounts (and on hot-reload, below).
_PLUGIN_MODULES: dict[str, object] = {}


def get_plugin_failures() -> dict[str, str]:
    """Return a copy of the failed-plugin registry (name → error message)."""
    return dict(_FAILED_PLUGINS)


def get_plugin_nav() -> list[dict]:
    """Plugin-contributed sidebar entries, sorted by section then label."""
    return sorted(_PLUGIN_NAV, key=lambda e: (e.get("section", ""), e.get("label", "")))


def _collect_plugin_nav(name: str, module) -> None:
    """Pull a plugin's optional get_nav() entries into the nav registry.

    Tolerant by design: a missing hook is fine, and a raising or malformed hook
    is logged and skipped — a bad nav contribution must never break loading.
    Idempotent per plugin (drops prior entries first) so hot-reload re-runs cleanly.
    """
    global _PLUGIN_NAV
    _PLUGIN_NAV = [e for e in _PLUGIN_NAV if e.get("plugin") != name]
    if not hasattr(module, "get_nav"):
        return
    try:
        items = module.get_nav() or []
    except Exception:
        logger.exception("Plugin %r: get_nav() raised, skipping its nav entries", name)
        return
    for item in items:
        try:
            label, href = str(item["label"]), str(item["href"])
        except (TypeError, KeyError):
            logger.warning("Plugin %r: malformed nav item %r, skipping", name, item)
            continue
        _PLUGIN_NAV.append({
            "plugin": name,
            "label": label,
            "href": href,
            "section": str(item.get("section", "Infrastructure")),
        })


def get_loaded_plugin_modules() -> dict[str, object]:
    """Loaded plugin modules by name (a copy), for hooks consumed after loading."""
    return dict(_PLUGIN_MODULES)


def loaded_plugin_submodule(name: str, sub: str):
    """A loaded plugin's submodule (e.g. ("docker", "models")), or None when
    the plugin isn't loaded. Another plugin reaches its models this way:
    importing plugins.<name>.models by path would define its tables a second
    time, because the plugin was loaded as _steward_plugin_<name> (#5775)."""
    import importlib
    loaded = _PLUGIN_MODULES.get(name)
    if loaded is None:
        return None
    try:
        return importlib.import_module(f"{loaded.__name__}.{sub}")
    except ImportError:
        return None


def _collect_plugin_module(app: "Quart", name: str, module) -> None:
    """Remember a loaded plugin module; hand it to a live MCP server if mounted.

    At boot the MCP server mounts after plugins load and reads them via
    get_loaded_plugin_modules(); a plugin hot-loaded later is registered here.
    """
    _PLUGIN_MODULES[name] = module
    mcp_server = getattr(app, "mcp_server", None)
    if mcp_server is not None:
        mcp_server.register_plugin(name, module)


def resolve_plugin_path(plugin_dirs: list[Path], name: str) -> Path | None:
    """Return the first plugin root that contains `name`, else None.

    Roots are searched in order, so a bundled (first-party) plugin shadows an
    external plugin of the same name. Pure function — no app/IO beyond exists().
    """
    for pd in plugin_dirs:
        candidate = pd / name
        if candidate.exists():
            return candidate
    return None


def _import_plugin(name: str, plugin_path: Path):
    """Load a plugin module by file path, avoiding sys.modules stdlib collisions.

    Using importlib.import_module(name) fails for plugins whose names shadow
    Python stdlib modules (e.g. a plugin named 'json' or 'http' vs the stdlib
    package of the same name).
    This helper loads from the filesystem path directly and registers the module
    under a namespaced key so relative imports within the plugin still work.
    """
    import importlib.util

    module_key = f"_steward_plugin_{name}"
    if module_key in sys.modules:
        return sys.modules[module_key]

    init_file = plugin_path / "__init__.py"
    spec = importlib.util.spec_from_file_location(
        module_key,
        str(init_file),
        submodule_search_locations=[str(plugin_path)],
    )
    if spec is None or spec.loader is None:
        raise ImportError(f"Could not create import spec for plugin {name!r}")

    module = importlib.util.module_from_spec(spec)
    module.__package__ = module_key
    sys.modules[module_key] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        sys.modules.pop(module_key, None)
        raise
    return module


def load_plugins(app: "Quart") -> None:
    """Import all enabled plugins and register their blueprints and scheduled tasks.

    For each plugin listed as enabled in config.yaml under 'plugins':
    1. Locate the plugin directory under PLUGIN_DIR
    2. Load and validate plugin.yaml (name must match dir, min_app_version check)
    3. Merge plugin.yaml config defaults with user overrides in app.config["PLUGINS"]
    4. Import the plugin Python package
    5. Validate required exports: setup() and get_scheduled_tasks()
    6. Call setup(app)
    7. Register blueprint at /plugins/<name>/ if get_blueprint() exists
    8. Append get_scheduled_tasks() results to app._task_registry
    """
    import steward

    plugin_dirs = [Path(d) for d in app.config["PLUGIN_DIRS"]]
    plugins_cfg: dict = app.config["PLUGINS"]

    # Ensure every plugin root is on sys.path so plugins are importable by name
    for pd in plugin_dirs:
        pd_str = str(pd.resolve())
        if pd_str not in sys.path:
            sys.path.insert(0, pd_str)

    for name, cfg in list(plugins_cfg.items()):
        if not cfg.get("enabled", False):
            continue

        plugin_path = resolve_plugin_path(plugin_dirs, name)
        if plugin_path is None:
            roots = ", ".join(str(d) for d in plugin_dirs)
            _FAILED_PLUGINS[name] = f"Plugin directory not found in: {roots}"
            logger.error("Plugin %r: not found in any plugin root (%s), skipping", name, roots)
            continue

        # Load and validate plugin.yaml
        yaml_path = plugin_path / "plugin.yaml"
        if not yaml_path.exists():
            _FAILED_PLUGINS[name] = "plugin.yaml not found"
            logger.error("Plugin %r: plugin.yaml not found, skipping", name)
            continue

        try:
            with yaml_path.open() as f:
                meta = yaml.safe_load(f) or {}
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"Failed to read plugin.yaml: {exc}"
            logger.exception("Plugin %r: failed to read plugin.yaml, skipping", name)
            continue

        if meta.get("name") != name:
            _FAILED_PLUGINS[name] = (
                f"plugin.yaml name {meta.get('name')!r} does not match directory name"
            )
            logger.error(
                "Plugin %r: plugin.yaml name=%r does not match directory name, skipping",
                name, meta.get("name"),
            )
            continue

        min_ver = meta.get("min_app_version")
        if min_ver:
            try:
                if Version(steward.__version__) < Version(min_ver):
                    _FAILED_PLUGINS[name] = (
                        f"Requires steward>={min_ver}, running {steward.__version__}"
                    )
                    logger.error(
                        "Plugin %r: requires steward>=%s but running %s, skipping",
                        name, min_ver, steward.__version__,
                    )
                    continue
            except Exception as exc:
                _FAILED_PLUGINS[name] = f"Invalid min_app_version {min_ver!r}: {exc}"
                logger.exception("Plugin %r: invalid min_app_version %r, skipping", name, min_ver)
                continue

        # Merge plugin.yaml config defaults with user overrides (non-"enabled" keys).
        # If a stored value's type doesn't match the default (e.g. a list default
        # was corrupted to a string in the DB), fall back to the yaml default.
        defaults = meta.get("config", {})
        user_overrides = {k: v for k, v in cfg.items() if k != "enabled"}
        merged = {"enabled": True, **defaults}
        for k, v in user_overrides.items():
            default_v = defaults.get(k)
            if default_v is not None and type(v) is not type(default_v):
                logger.warning(
                    "Plugin %r: config key %r has wrong type (%s), using yaml default",
                    name, k, type(v).__name__,
                )
            else:
                merged[k] = v
        plugins_cfg[name] = merged

        # Import the plugin package (file-path load avoids stdlib name collisions)
        try:
            module = _import_plugin(name, plugin_path)
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"Import error: {exc}"
            logger.exception("Plugin %r: import failed, skipping", name)
            continue

        # Validate required exports
        if not hasattr(module, "setup"):
            _FAILED_PLUGINS[name] = "Missing required export: setup()"
            logger.error("Plugin %r: missing required export 'setup', skipping", name)
            continue
        if not hasattr(module, "get_scheduled_tasks"):
            _FAILED_PLUGINS[name] = "Missing required export: get_scheduled_tasks()"
            logger.error(
                "Plugin %r: missing required export 'get_scheduled_tasks', skipping", name
            )
            continue

        # Call setup
        try:
            module.setup(app)
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"setup() failed: {exc}"
            logger.exception("Plugin %r: setup() raised, skipping", name)
            continue

        # Register blueprint (optional)
        if hasattr(module, "get_blueprint"):
            try:
                bp = module.get_blueprint()
                app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
                logger.info("Plugin %r: blueprint registered at /plugins/%s/", name, name)
            except Exception as exc:
                _FAILED_PLUGINS[name] = f"Blueprint registration failed: {exc}"
                logger.exception("Plugin %r: get_blueprint() raised", name)
                continue

        # Register scheduled tasks
        try:
            tasks = module.get_scheduled_tasks()
            app._task_registry.extend(tasks)
            logger.info("Plugin %r: registered %d scheduled task(s)", name, len(tasks))
        except Exception as exc:
            _FAILED_PLUGINS[name] = f"get_scheduled_tasks() failed: {exc}"
            logger.exception("Plugin %r: get_scheduled_tasks() raised", name)
            continue

        _FAILED_PLUGINS.pop(name, None)  # clear any stale failure from a previous reload attempt
        _LOADED_PLUGINS.add(name)
        _collect_plugin_nav(name, module)
        _collect_plugin_module(app, name, module)
        logger.info("Plugin %r loaded (v%s)", name, meta.get("version", "?"))


def hot_reload_plugin(app: "Quart", name: str) -> tuple[bool, str]:
    """Activate a plugin that was just enabled without restarting the app.

    This works for plugins that were NOT previously loaded. For plugins that
    are already registered (blueprint already mounted), a full restart is
    required to pick up code changes — use restart_app() in that case.

    Returns (success, message).
    """
    import steward

    if name in _LOADED_PLUGINS:
        return False, "Plugin already loaded — restart required to apply updates"

    plugin_dirs = [Path(d).resolve() for d in app.config["PLUGIN_DIRS"]]
    plugin_path = resolve_plugin_path(plugin_dirs, name)

    if plugin_path is None:
        return False, f"Plugin {name!r} not found in any plugin root"

    yaml_path = plugin_path / "plugin.yaml"
    if not yaml_path.exists():
        return False, "plugin.yaml not found"

    try:
        with yaml_path.open() as f:
            meta = yaml.safe_load(f) or {}
    except Exception as exc:
        return False, f"Could not read plugin.yaml: {exc}"

    if meta.get("name") != name:
        return False, f"plugin.yaml name mismatch: expected {name!r}, got {meta.get('name')!r}"

    min_ver = meta.get("min_app_version")
    if min_ver:
        try:
            if Version(steward.__version__) < Version(min_ver):
                return False, (
                    f"Plugin requires steward>={min_ver} "
                    f"but running {steward.__version__}"
                )
        except Exception:
            return False, f"Invalid min_app_version: {min_ver!r}"

    # Merge config defaults with any stored user config.
    # Discard stored values whose type doesn't match the yaml default (corruption guard).
    defaults = meta.get("config", {})
    stored_cfg = app.config.get("PLUGINS", {}).get(name, {})
    merged = {"enabled": True, **defaults}
    for k, v in stored_cfg.items():
        if k == "enabled":
            continue
        default_v = defaults.get(k)
        if default_v is not None and type(v) is not type(default_v):
            logger.warning("Plugin %r: config key %r has wrong type, using yaml default", name, k)
        else:
            merged[k] = v
    app.config["PLUGINS"][name] = merged

    # Run migrations if the plugin has them (safe to re-run — alembic is idempotent).
    # Pass ALL plugin migration dirs so Alembic can resolve any previously-stamped
    # revisions from other plugins that are already in alembic_version.
    mdir = plugin_path / "migrations"
    if mdir.exists():
        try:
            from steward.core.migration_runner import (
                run_plugin_migrations,
                discover_all_in,
            )
            all_dirs = discover_all_in(plugin_dirs)
            run_plugin_migrations(app.config["DATABASE_URL"], all_dirs)
        except Exception:
            logger.exception("Plugin %r: migration failed during hot-reload", name)
            return False, "Migration failed — check logs"

    # Import (file-path load avoids stdlib name collisions)
    try:
        module = _import_plugin(name, plugin_path)
    except Exception as exc:
        return False, f"Import failed: {exc}"

    if not hasattr(module, "setup") or not hasattr(module, "get_scheduled_tasks"):
        return False, "Plugin missing required exports (setup, get_scheduled_tasks)"

    try:
        module.setup(app)
    except Exception as exc:
        return False, f"setup() raised: {exc}"

    if hasattr(module, "get_blueprint"):
        try:
            bp = module.get_blueprint()
            app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
        except Exception as exc:
            return False, f"Blueprint registration failed: {exc}"

    try:
        tasks = module.get_scheduled_tasks()
        # De-duplicate: skip tasks whose name already exists in the registry
        existing_names = {t.name for t in app._task_registry}
        new_tasks = [t for t in tasks if t.name not in existing_names]
        app._task_registry.extend(new_tasks)
    except Exception as exc:
        return False, f"get_scheduled_tasks() raised: {exc}"

    _LOADED_PLUGINS.add(name)
    _collect_plugin_nav(name, module)
    _collect_plugin_module(app, name, module)
    logger.info("Plugin %r hot-reloaded (v%s)", name, meta.get("version", "?"))
    return True, f"Plugin activated (v{meta.get('version', '?')})"


def restart_app() -> None:
    """Replace the current process with a fresh copy (in-app restart).

    Used when a plugin update requires a full restart to take effect.
    In Docker the container will be relaunched by the restart policy.
    """
    logger.info("Initiating in-app restart via os.execv")
    os.execv(sys.executable, [sys.executable] + sys.argv)