system_tools.py•5.98 kB
"""
System information and command execution tools for the MCP server.
"""
import contextlib
import io
import logging
import os
import platform
import subprocess
import sys
import traceback
import psutil
from . import mcp
@mcp.tool()
async def system_info() -> str:
"""
Get basic system information.
Returns:
Formatted system information string.
"""
logging.info("system_info called")
try:
memory = psutil.virtual_memory()
info = (
f"System Information:\n"
f"Platform: {platform.system()} {platform.release()}\n"
f"Version: {platform.version()}\n"
f"Architecture: {platform.machine()}\n"
f"CPU Count: {os.cpu_count()}\n"
f"Total Memory: {memory.total // (1024**3)} GB\n"
f"Available Memory: {memory.available // (1024**3)} GB\n"
f"Memory Usage: {memory.percent}%"
)
logging.info("System info retrieved successfully")
return info
except (OSError, psutil.Error, Exception) as exc:
logging.error("Error retrieving system info: %s", exc)
return f"Error retrieving system info: {exc}"
@mcp.tool()
async def list_processes() -> str:
"""
List running processes.
Returns:
Formatted list of running processes.
"""
logging.info("list_processes called")
processes = []
try:
for proc in psutil.process_iter(["pid", "name", "username"]):
try:
info = proc.info
processes.append(info)
except (psutil.NoSuchProcess, psutil.AccessDenied) as exc:
logging.debug("Error accessing process info: %s", exc)
continue
# Format the output
output = f"Running Processes ({len(processes)} total):\n\n"
for proc in processes[:20]: # Limit to first 20 processes
output += (
f"PID: {proc.get('pid', 'N/A')} | "
f"Name: {proc.get('name', 'N/A')} | "
f"User: {proc.get('username', 'N/A')}\n"
)
if len(processes) > 20:
output += f"\n... and {len(processes) - 20} more processes"
logging.info("Listed %d processes", len(processes))
return output
except (psutil.Error, Exception) as exc:
logging.error("Error listing processes: %s", exc)
return f"Error listing processes: {exc}"
@mcp.tool()
async def run_python_code(code: str) -> str:
"""
Execute a snippet of Python code and return the output or error.
Args:
code: The Python code to execute.
Returns:
The output of the code or an error message.
"""
logging.info("run_python_code called")
local_vars = {}
stdout = io.StringIO()
try:
with contextlib.redirect_stdout(stdout):
exec(code, {}, local_vars)
output = stdout.getvalue()
if not output and local_vars:
output = f"Result: {local_vars}"
elif not output:
output = "Code executed successfully with no output."
logging.info("Python code executed successfully")
return output
except Exception as exc:
error_msg = f"Error executing code:\n{traceback.format_exc()}"
logging.error("Error executing code: %s", exc)
return error_msg
@mcp.tool()
async def install_python_library(library: str, version: str = "") -> str:
"""
Install a Python library using pip and update requirements.txt.
Args:
library: The name of the library to install (e.g., 'requests').
version: Optional version specifier (e.g., '2.32.0'). If empty, installs latest.
Returns:
A message indicating success or failure.
"""
req_line = f"{library}=={version}" if version else library
logging.info("Requested install of library: %s", req_line)
# Run pip install
try:
pip_args = [sys.executable, "-m", "pip", "install"]
if version:
pip_args.append(f"{library}=={version}")
else:
pip_args.append(library)
result = subprocess.run(
pip_args,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
check=False,
)
if result.returncode != 0:
logging.error("pip install failed: %s", result.stderr)
return f"Error installing {req_line}:\n{result.stderr}"
except (subprocess.SubprocessError, OSError) as exc:
logging.error("Exception during pip install: %s", exc)
return f"Exception during installation: {exc}"
# Update requirements.txt
req_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "..", "requirements.txt"
)
req_path = os.path.normpath(req_path)
try:
# Read current requirements
if os.path.exists(req_path):
with open(req_path, "r", encoding="utf-8") as f:
lines = [line.strip() for line in f.readlines()]
else:
lines = []
# Remove any old entry for this library
base_lib = library.lower().replace("_", "-")
new_lines = [
line
for line in lines
if not (
line.lower().startswith(base_lib + "==") or line.lower() == base_lib
)
]
# Add new entry
new_lines.append(req_line)
# Sort and deduplicate
new_lines = sorted(set(new_lines), key=lambda x: x.lower())
with open(req_path, "w", encoding="utf-8") as f:
for line in new_lines:
f.write(line + "\n")
logging.info("Updated requirements.txt with %s", req_line)
except (OSError, IOError) as exc:
logging.error("Failed to update requirements.txt: %s", exc)
return f"Installed {req_line}, but failed to update requirements.txt: {exc}"
return f"Successfully installed {req_line} and updated requirements.txt."