unraid-mcp/unraid_mcp/tools/virtualization.py

"""Virtual machine management.

Provides the `unraid_vm` tool with 9 actions for VM lifecycle management
including start, stop, pause, resume, force stop, reboot, and reset.
"""

from typing import Any, Literal, get_args

from fastmcp import FastMCP

from ..config.logging import logger
from ..core.client import make_graphql_request
from ..core.exceptions import ToolError, tool_error_handler


QUERIES: dict[str, str] = {
    "list": """
        query ListVMs {
          vms { id domains { id name state uuid } }
        }
    """,
    # NOTE: The Unraid GraphQL API does not expose a single-VM query.
    # The details query is identical to list; client-side filtering is required.
    "details": """
        query ListVMs {
          vms { id domains { id name state uuid } }
        }
    """,
}

MUTATIONS: dict[str, str] = {
    "start": """
        mutation StartVM($id: PrefixedID!) { vm { start(id: $id) } }
    """,
    "stop": """
        mutation StopVM($id: PrefixedID!) { vm { stop(id: $id) } }
    """,
    "pause": """
        mutation PauseVM($id: PrefixedID!) { vm { pause(id: $id) } }
    """,
    "resume": """
        mutation ResumeVM($id: PrefixedID!) { vm { resume(id: $id) } }
    """,
    "force_stop": """
        mutation ForceStopVM($id: PrefixedID!) { vm { forceStop(id: $id) } }
    """,
    "reboot": """
        mutation RebootVM($id: PrefixedID!) { vm { reboot(id: $id) } }
    """,
    "reset": """
        mutation ResetVM($id: PrefixedID!) { vm { reset(id: $id) } }
    """,
}

# Map action names to GraphQL field names (only where they differ)
_MUTATION_FIELDS: dict[str, str] = {
    "force_stop": "forceStop",
}

DESTRUCTIVE_ACTIONS = {"force_stop", "reset"}

VM_ACTIONS = Literal[
    "list",
    "details",
    "start",
    "stop",
    "pause",
    "resume",
    "force_stop",
    "reboot",
    "reset",
]

ALL_ACTIONS = set(QUERIES) | set(MUTATIONS)

if set(get_args(VM_ACTIONS)) != ALL_ACTIONS:
    _missing = ALL_ACTIONS - set(get_args(VM_ACTIONS))
    _extra = set(get_args(VM_ACTIONS)) - ALL_ACTIONS
    raise RuntimeError(
        f"VM_ACTIONS and ALL_ACTIONS are out of sync. "
        f"Missing from Literal: {_missing or 'none'}. Extra in Literal: {_extra or 'none'}"
    )


def register_vm_tool(mcp: FastMCP) -> None:
    """Register the unraid_vm tool with the FastMCP instance."""

    @mcp.tool()
    async def unraid_vm(
        action: VM_ACTIONS,
        vm_id: str | None = None,
        confirm: bool = False,
    ) -> dict[str, Any]:
        """Manage Unraid virtual machines.

        Actions:
          list - List all VMs with state
          details - Detailed info for a VM (requires vm_id: UUID, PrefixedID, or name)
          start - Start a VM (requires vm_id)
          stop - Gracefully stop a VM (requires vm_id)
          pause - Pause a VM (requires vm_id)
          resume - Resume a paused VM (requires vm_id)
          force_stop - Force stop a VM (requires vm_id, confirm=True)
          reboot - Reboot a VM (requires vm_id)
          reset - Reset a VM (requires vm_id, confirm=True)
        """
        if action not in ALL_ACTIONS:
            raise ToolError(f"Invalid action '{action}'. Must be one of: {sorted(ALL_ACTIONS)}")

        if action != "list" and not vm_id:
            raise ToolError(f"vm_id is required for '{action}' action")

        if action in DESTRUCTIVE_ACTIONS and not confirm:
            raise ToolError(f"Action '{action}' is destructive. Set confirm=True to proceed.")

        with tool_error_handler("vm", action, logger):
            try:
                logger.info(f"Executing unraid_vm action={action}")

                if action == "list":
                    data = await make_graphql_request(QUERIES["list"])
                    if data.get("vms"):
                        vms = data["vms"].get("domains") or data["vms"].get("domain") or []
                        if isinstance(vms, dict):
                            vms = [vms]
                        return {"vms": vms}
                    return {"vms": []}

                if action == "details":
                    data = await make_graphql_request(QUERIES["details"])
                    if not data.get("vms"):
                        raise ToolError("No VM data returned from server")
                    vms = data["vms"].get("domains") or data["vms"].get("domain") or []
                    if isinstance(vms, dict):
                        vms = [vms]
                    for vm in vms:
                        if (
                            vm.get("uuid") == vm_id
                            or vm.get("id") == vm_id
                            or vm.get("name") == vm_id
                        ):
                            return dict(vm)
                    available = [f"{v.get('name')} (UUID: {v.get('uuid')})" for v in vms]
                    raise ToolError(f"VM '{vm_id}' not found. Available: {', '.join(available)}")

                # Mutations
                if action in MUTATIONS:
                    data = await make_graphql_request(MUTATIONS[action], {"id": vm_id})
                    field = _MUTATION_FIELDS.get(action, action)
                    if data.get("vm") and field in data["vm"]:
                        return {
                            "success": data["vm"][field],
                            "action": action,
                            "vm_id": vm_id,
                        }
                    raise ToolError(f"Failed to {action} VM or unexpected response")

                raise ToolError(f"Unhandled action '{action}' — this is a bug")

            except ToolError:
                raise
            except Exception as e:
                if "VMs are not available" in str(e):
                    raise ToolError(
                        "VMs not available on this server. Check VM support is enabled."
                    ) from e
                raise

    logger.info("VM tool registered successfully")