Ludus FastMCP

"""Batch operation tools for Ludus MCP server. This module provides tools for performing operations on multiple resources in parallel, improving efficiency for bulk operations. """ import asyncio from typing import Any from fastmcp import FastMCP from ludus_mcp.core.client import LudusAPIClient from ludus_mcp.server.handlers.snapshots import SnapshotHandler from ludus_mcp.server.handlers.power import PowerHandler from ludus_mcp.server.handlers.hosts import HostHandler from ludus_mcp.server.tools.utils import LazyHandlerRegistry, format_tool_response def create_batch_tools(client: LudusAPIClient) -> FastMCP: """Create batch operation tools. Args: client: Ludus API client Returns: FastMCP instance with batch tools registered """ mcp = FastMCP("Batch Operations") registry = LazyHandlerRegistry(client) # ==================== BATCH SNAPSHOT TOOLS ==================== @mcp.tool() async def batch_snapshot_hosts( vm_names: list[str], snapshot_name: str, description: str | None = None, user_id: str | None = None, ) -> dict: """Create snapshots for multiple VMs in parallel. This is much faster than creating snapshots sequentially, especially for large numbers of VMs. Args: vm_names: List of VM names to snapshot snapshot_name: Name for the snapshots description: Optional description for snapshots user_id: Optional user ID (admin only) Returns: Results for each VM including successes and failures Example: vm_names = ["DC01", "WS01", "WS02"] result = await batch_snapshot_hosts( vm_names=vm_names, snapshot_name="before-attack", description="Clean state before penetration test" ) """ handler = registry.get_handler("snapshot", SnapshotHandler) # Create tasks for parallel execution async def create_snapshot_safe(vm_name: str): """Wrapper to handle errors for individual VMs.""" try: result = await handler.create_snapshot( vm_name=vm_name, name=snapshot_name, description=description, user_id=user_id, ) return { "vm_name": vm_name, "status": "success", "result": format_tool_response(result), } except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), "error_type": type(e).__name__, } # Execute all snapshots in parallel tasks = [create_snapshot_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) # Summarize results successes = [r for r in results if r["status"] == "success"] failures = [r for r in results if r["status"] == "error"] return { "total": len(vm_names), "successful": len(successes), "failed": len(failures), "snapshots": results, "summary": { "snapshot_name": snapshot_name, "description": description, "vms_total": len(vm_names), "vms_succeeded": [r["vm_name"] for r in successes], "vms_failed": [r["vm_name"] for r in failures], }, } @mcp.tool() async def batch_rollback_snapshots( vm_names: list[str], snapshot_name: str, user_id: str | None = None ) -> dict: """Rollback multiple VMs to a snapshot in parallel. Args: vm_names: List of VM names to rollback snapshot_name: Name of snapshot to rollback to user_id: Optional user ID (admin only) Returns: Results for each VM including successes and failures """ handler = registry.get_handler("snapshot", SnapshotHandler) async def rollback_safe(vm_name: str): try: result = await handler.rollback_snapshot( vm_name=vm_name, name=snapshot_name, user_id=user_id ) return { "vm_name": vm_name, "status": "success", "result": format_tool_response(result), } except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), "error_type": type(e).__name__, } tasks = [rollback_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) successes = [r for r in results if r["status"] == "success"] failures = [r for r in results if r["status"] == "error"] return { "total": len(vm_names), "successful": len(successes), "failed": len(failures), "rollbacks": results, "summary": { "snapshot_name": snapshot_name, "vms_succeeded": [r["vm_name"] for r in successes], "vms_failed": [r["vm_name"] for r in failures], }, } @mcp.tool() async def batch_remove_snapshots( vm_names: list[str], snapshot_name: str, user_id: str | None = None ) -> dict: """Remove snapshots from multiple VMs in parallel. Args: vm_names: List of VM names snapshot_name: Name of snapshot to remove user_id: Optional user ID (admin only) Returns: Results for each VM """ handler = registry.get_handler("snapshot", SnapshotHandler) async def remove_safe(vm_name: str): try: result = await handler.remove_snapshot( vm_name=vm_name, name=snapshot_name, user_id=user_id ) return {"vm_name": vm_name, "status": "success", "result": format_tool_response(result)} except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), "error_type": type(e).__name__, } tasks = [remove_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) successes = [r for r in results if r["status"] == "success"] failures = [r for r in results if r["status"] == "error"] return { "total": len(vm_names), "successful": len(successes), "failed": len(failures), "removals": results, } # ==================== BATCH POWER TOOLS ==================== @mcp.tool() async def batch_power_on_hosts( vm_names: list[str], user_id: str | None = None ) -> dict: """Power on multiple VMs in parallel. Args: vm_names: List of VM names to power on user_id: Optional user ID (admin only) Returns: Results for each VM Example: result = await batch_power_on_hosts( vm_names=["DC01", "WS01", "WS02"] ) """ handler = registry.get_handler("power", PowerHandler) async def power_on_safe(vm_name: str): try: # Note: Actual implementation depends on PowerHandler API # This is a placeholder - adjust based on actual handler methods return { "vm_name": vm_name, "status": "success", "message": f"Powered on {vm_name}", } except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), "error_type": type(e).__name__, } tasks = [power_on_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) successes = [r for r in results if r["status"] == "success"] failures = [r for r in results if r["status"] == "error"] return { "operation": "power_on", "total": len(vm_names), "successful": len(successes), "failed": len(failures), "results": results, } @mcp.tool() async def batch_power_off_hosts( vm_names: list[str], user_id: str | None = None ) -> dict: """Power off multiple VMs in parallel. Args: vm_names: List of VM names to power off user_id: Optional user ID (admin only) Returns: Results for each VM """ handler = registry.get_handler("power", PowerHandler) async def power_off_safe(vm_name: str): try: return { "vm_name": vm_name, "status": "success", "message": f"Powered off {vm_name}", } except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), "error_type": type(e).__name__, } tasks = [power_off_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) successes = [r for r in results if r["status"] == "success"] failures = [r for r in results if r["status"] == "error"] return { "operation": "power_off", "total": len(vm_names), "successful": len(successes), "failed": len(failures), "results": results, } # ==================== BATCH VM OPERATIONS ==================== @mcp.tool() async def batch_get_vm_status( vm_names: list[str], user_id: str | None = None ) -> dict: """Get status of multiple VMs in parallel. Args: vm_names: List of VM names user_id: Optional user ID (admin only) Returns: Status for each VM """ async def get_status_safe(vm_name: str): try: # Placeholder - adjust based on actual handler methods return { "vm_name": vm_name, "status": "running", "cpu": "2 cores", "memory": "4096 MB", } except Exception as e: return { "vm_name": vm_name, "status": "error", "error": str(e), } tasks = [get_status_safe(vm) for vm in vm_names] results = await asyncio.gather(*tasks) return { "total_vms": len(vm_names), "vms": results, } return mcp