From 088588b3b350d71d5e89459e796f0736f046d34f Mon Sep 17 00:00:00 2001 From: canvrno Date: Tue, 18 Feb 2025 22:58:59 -0700 Subject: [PATCH] Tools cleanup --- src/proxmox_mcp/server.py | 172 ++++++++++++---------------- src/proxmox_mcp/tools/vm_console.py | 15 +-- 2 files changed, 81 insertions(+), 106 deletions(-) diff --git a/src/proxmox_mcp/server.py b/src/proxmox_mcp/server.py index b5c35cb..593eec8 100644 --- a/src/proxmox_mcp/server.py +++ b/src/proxmox_mcp/server.py @@ -2,48 +2,31 @@ import json import logging import os +import sys +import signal from pathlib import Path -from typing import Dict, Any, Optional, List +from typing import Dict, Any, Optional, List, Annotated +from pydantic import BaseModel, Field +from urllib.parse import urljoin from mcp.server.fastmcp import FastMCP from mcp.server.fastmcp.tools import Tool from mcp.server.fastmcp.tools.base import Tool as BaseTool from mcp.types import CallToolResult as Response, TextContent as Content from proxmoxer import ProxmoxAPI -from pydantic import BaseModel -from urllib.parse import urljoin - -class CustomProxmoxAPI(ProxmoxAPI): - def __init__(self, host, **kwargs): - self._host = host # Store the host - super().__init__(host, **kwargs) - - def get(self, *args, **kwargs): - try: - # Always ensure base_url is set correctly - self._store['base_url'] = f'https://{self._host}:{self._store.get("port", "8006")}/api2/json' - print(f"Using base_url: {self._store['base_url']}") - - # Handle both dotted and string notation - if args and isinstance(args[0], str): - path = args[0] - print(f"Making request to path: {path}") - full_url = f"{self._store['base_url']}/{path}" - print(f"Full URL: {full_url}") - return self._backend.get(full_url) - - print("Using dotted notation") - return super().get(*args, **kwargs) - except Exception as e: - print(f"Error in CustomProxmoxAPI.get: {e}") - raise # Import from the same directory -import sys -import os.path sys.path.append(os.path.dirname(os.path.abspath(__file__))) from tools.vm_console import VMConsoleManager +class NodeStatus(BaseModel): + node: Annotated[str, Field(description="Name/ID of node to query (e.g. 'pve1', 'proxmox-node2')")] + +class VMCommand(BaseModel): + node: Annotated[str, Field(description="Host node name (e.g. 'pve1', 'proxmox-node2')")] + vmid: Annotated[str, Field(description="VM ID number (e.g. '100', '101')")] + command: Annotated[str, Field(description="Shell command to run (e.g. 'uname -a', 'systemctl status nginx')")] + class ProxmoxConfig(BaseModel): host: str port: int = 8006 @@ -76,20 +59,14 @@ class ProxmoxMCPServer: def _load_config(self, config_path: Optional[str]) -> Config: """Load configuration from file or environment variables.""" - print(f"Loading config from path: {config_path}") if not config_path: raise ValueError("PROXMOX_MCP_CONFIG environment variable must be set") try: with open(config_path) as f: config_data = json.load(f) - print(f"Raw config data: {config_data}") - - # Ensure host is not empty if not config_data.get('proxmox', {}).get('host'): raise ValueError("Proxmox host cannot be empty") - - print(f"Using host: {config_data['proxmox']['host']}") return Config(**config_data) except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON in config file: {e}") @@ -98,33 +75,36 @@ class ProxmoxMCPServer: def _setup_logging(self) -> None: """Configure logging based on settings.""" - import os - # Convert relative path to absolute log_file = self.config.logging.file if log_file and not os.path.isabs(log_file): log_file = os.path.join(os.getcwd(), log_file) + # Create handlers + file_handler = logging.FileHandler(log_file) + file_handler.setLevel(getattr(logging, self.config.logging.level.upper())) + + # Console handler for errors only + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.ERROR) + + # Configure formatters + formatter = logging.Formatter(self.config.logging.format) + file_handler.setFormatter(formatter) + console_handler.setFormatter(formatter) + # Configure root logger - logging.basicConfig( - level=getattr(logging, self.config.logging.level.upper()), - format=self.config.logging.format, - handlers=[ - logging.FileHandler(log_file), - logging.StreamHandler() # Also log to console - ] - ) + root_logger = logging.getLogger() + root_logger.setLevel(getattr(logging, self.config.logging.level.upper())) + root_logger.addHandler(file_handler) + root_logger.addHandler(console_handler) self.logger = logging.getLogger("proxmox-mcp") - self.logger.info(f"Logging initialized. File: {log_file}, Level: {self.config.logging.level}") def _setup_proxmox(self) -> ProxmoxAPI: """Initialize Proxmox API connection.""" try: self.logger.info(f"Connecting to Proxmox with config: {self.config.proxmox}") - print(f"Initializing ProxmoxAPI with host={self.config.proxmox.host}, port={self.config.proxmox.port}") - - print(f"Creating ProxmoxAPI with host={self.config.proxmox.host}") # Store the working configuration for reuse self.proxmox_config = { @@ -136,19 +116,9 @@ class ProxmoxMCPServer: 'service': 'PVE' } - # Create CustomProxmoxAPI instance with stored config - api = CustomProxmoxAPI(**self.proxmox_config) - print("ProxmoxAPI initialized successfully") - # Test the connection by making a simple request - print("Testing API connection...") - test_result = api.version.get() - print(f"Connection test result: {test_result}") - - # Test nodes endpoint specifically - print("Testing nodes endpoint...") - nodes_result = api.nodes.get() - print(f"Nodes test result: {nodes_result}") - + # Create and test CustomProxmoxAPI instance + api = ProxmoxAPI(**self.proxmox_config) + api.version.get() # Test connection return api except Exception as e: self.logger.error(f"Failed to connect to Proxmox: {e}") @@ -157,27 +127,26 @@ class ProxmoxMCPServer: def _setup_tools(self) -> None: """Register MCP tools.""" - @self.mcp.tool() + @self.mcp.tool( + description="List all nodes in the Proxmox cluster with their status, CPU, memory, and role information.\n\n" + "Example:\n" + '{"node": "pve1", "status": "online", "cpu_usage": 0.15, "memory": {"used": "8GB", "total": "32GB"}}') def get_nodes() -> List[Content]: - """List all nodes in the Proxmox cluster.""" try: - print(f"Using ProxmoxAPI instance with config: {self.proxmox_config}") - print("Getting nodes using dotted notation...") result = self.proxmox.nodes.get() - print(f"Raw nodes result: {result}") nodes = [{"node": node["node"], "status": node["status"]} for node in result] return [Content(type="text", text=json.dumps(nodes))] except Exception as e: self.logger.error(f"Failed to get nodes: {e}") raise - @self.mcp.tool() - def get_node_status(node: str) -> List[Content]: - """Get detailed status of a specific node. - - Args: - node: Name of the node to get status for - """ + @self.mcp.tool( + description="Get detailed status information for a specific Proxmox node.\n\n" + "Parameters:\n" + "node* - Name/ID of node to query (e.g. 'pve1')\n\n" + "Example:\n" + '{"cpu": {"usage": 0.15}, "memory": {"used": "8GB", "total": "32GB"}}') + def get_node_status(node: Annotated[str, Field(description="Name/ID of node to query (e.g. 'pve1')")]) -> List[Content]: try: result = self.proxmox.nodes(node).status.get() return [Content(type="text", text=json.dumps(result))] @@ -185,9 +154,11 @@ class ProxmoxMCPServer: self.logger.error(f"Failed to get node status: {e}") raise - @self.mcp.tool() + @self.mcp.tool( + description="List all virtual machines across the cluster with their status and resource usage.\n\n" + "Example:\n" + '{"vmid": "100", "name": "ubuntu", "status": "running", "cpu": 2, "memory": 4096}') def get_vms() -> List[Content]: - """List all VMs across the cluster.""" try: result = [] for node in self.proxmox.nodes.get(): @@ -203,9 +174,11 @@ class ProxmoxMCPServer: self.logger.error(f"Failed to get VMs: {e}") raise - @self.mcp.tool() + @self.mcp.tool( + description="List all LXC containers across the cluster with their status and configuration.\n\n" + "Example:\n" + '{"vmid": "200", "name": "nginx", "status": "running", "template": "ubuntu-20.04"}') def get_containers() -> List[Content]: - """List all LXC containers.""" try: result = [] for node in self.proxmox.nodes.get(): @@ -221,9 +194,11 @@ class ProxmoxMCPServer: self.logger.error(f"Failed to get containers: {e}") raise - @self.mcp.tool() + @self.mcp.tool( + description="List storage pools across the cluster with their usage and configuration.\n\n" + "Example:\n" + '{"storage": "local-lvm", "type": "lvm", "used": "500GB", "total": "1TB"}') def get_storage() -> List[Content]: - """List available storage.""" try: result = self.proxmox.storage.get() storage = [{"storage": storage["storage"], "type": storage["type"]} for storage in result] @@ -232,9 +207,11 @@ class ProxmoxMCPServer: self.logger.error(f"Failed to get storage: {e}") raise - @self.mcp.tool() + @self.mcp.tool( + description="Get overall Proxmox cluster health and configuration status.\n\n" + "Example:\n" + '{"name": "proxmox", "quorum": "ok", "nodes": 3, "ha_status": "active"}') def get_cluster_status() -> List[Content]: - """Get overall cluster status.""" try: result = self.proxmox.cluster.status.get() return [Content(type="text", text=json.dumps(result))] @@ -242,15 +219,19 @@ class ProxmoxMCPServer: self.logger.error(f"Failed to get cluster status: {e}") raise - @self.mcp.tool() - async def execute_vm_command(node: str, vmid: str, command: str) -> List[Content]: - """Execute a command in a VM's console. - - Args: - node: Name of the node where VM is running - vmid: ID of the VM - command: Command to execute - """ + @self.mcp.tool( + description="Execute commands in a VM via QEMU guest agent.\n\n" + "Parameters:\n" + "node* - Host node name (e.g. 'pve1')\n" + "vmid* - VM ID number (e.g. '100')\n" + "command* - Shell command to run (e.g. 'uname -a')\n\n" + "Example:\n" + '{"success": true, "output": "Linux vm1 5.4.0", "exit_code": 0}') + async def execute_vm_command( + node: Annotated[str, Field(description="Host node name (e.g. 'pve1')")], + vmid: Annotated[str, Field(description="VM ID number (e.g. '100')")], + command: Annotated[str, Field(description="Shell command to run (e.g. 'uname -a')")] + ) -> List[Content]: try: result = await self.vm_console.execute_command(node, vmid, command) return [Content(type="text", text=json.dumps(result))] @@ -261,11 +242,9 @@ class ProxmoxMCPServer: def start(self) -> None: """Start the MCP server.""" import anyio - import signal - import sys def signal_handler(signum, frame): - print("Received signal to shutdown...") + self.logger.info("Received signal to shutdown...") sys.exit(0) # Set up signal handlers @@ -273,10 +252,9 @@ class ProxmoxMCPServer: signal.signal(signal.SIGTERM, signal_handler) try: - print("Starting MCP server...") + self.logger.info("Starting MCP server...") anyio.run(self.mcp.run_stdio_async) except Exception as e: - print(f"Server error: {e}") self.logger.error(f"Server error: {e}") sys.exit(1) diff --git a/src/proxmox_mcp/tools/vm_console.py b/src/proxmox_mcp/tools/vm_console.py index 37dd8f2..c5d1de8 100644 --- a/src/proxmox_mcp/tools/vm_console.py +++ b/src/proxmox_mcp/tools/vm_console.py @@ -52,9 +52,9 @@ class VMConsoleManager: # Start command execution self.logger.info("Starting command execution...") try: - print(f"Executing command via agent: {command}") + self.logger.debug(f"Executing command via agent: {command}") exec_result = endpoint("exec").post(command=command) - print(f"Raw exec response: {exec_result}") + self.logger.debug(f"Raw exec response: {exec_result}") self.logger.info(f"Command started with result: {exec_result}") except Exception as e: self.logger.error(f"Failed to start command: {str(e)}") @@ -72,23 +72,20 @@ class VMConsoleManager: # Get command output using exec-status try: - print(f"Getting status for PID {pid}...") + self.logger.debug(f"Getting status for PID {pid}...") console = endpoint("exec-status").get(pid=pid) - print(f"Raw exec-status response: {console}") + self.logger.debug(f"Raw exec-status response: {console}") if not console: raise RuntimeError("No response from exec-status") except Exception as e: self.logger.error(f"Failed to get command status: {str(e)}") raise RuntimeError(f"Failed to get command status: {str(e)}") self.logger.info(f"Command completed with status: {console}") - print(f"Command completed with status: {console}") except Exception as e: self.logger.error(f"API call failed: {str(e)}") - print(f"API call error: {str(e)}") # Print error for immediate feedback raise RuntimeError(f"API call failed: {str(e)}") - self.logger.info(f"Raw API response type: {type(console)}") - self.logger.info(f"Raw API response: {console}") - print(f"Raw API response: {console}") # Print to stdout for immediate feedback + self.logger.debug(f"Raw API response type: {type(console)}") + self.logger.debug(f"Raw API response: {console}") # Handle different response structures if isinstance(console, dict):