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

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)


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")