Tools cleanup
This commit is contained in:
canvrno
@@ -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)
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user