This commit is contained in:
Jacob Magar
2025-08-12 11:35:00 -04:00
parent 8fbec924cd
commit 493a376640
34 changed files with 525 additions and 1564 deletions

View File

@@ -5,7 +5,7 @@ log files, physical disks with SMART data, and system storage operations
with custom timeout configurations for disk-intensive operations.
"""
from typing import Any, Dict, List, Optional
from typing import Any
import httpx
from fastmcp import FastMCP
@@ -15,15 +15,15 @@ from ..core.client import make_graphql_request
from ..core.exceptions import ToolError
def register_storage_tools(mcp: FastMCP):
def register_storage_tools(mcp: FastMCP) -> None:
"""Register all storage tools with the FastMCP instance.
Args:
mcp: FastMCP instance to register tools with
"""
@mcp.tool()
async def get_shares_info() -> List[Dict[str, Any]]:
async def get_shares_info() -> list[dict[str, Any]]:
"""Retrieves information about user shares."""
query = """
query GetSharesInfo {
@@ -50,13 +50,14 @@ def register_storage_tools(mcp: FastMCP):
try:
logger.info("Executing get_shares_info tool")
response_data = await make_graphql_request(query)
return response_data.get("shares", [])
shares = response_data.get("shares", [])
return list(shares) if isinstance(shares, list) else []
except Exception as e:
logger.error(f"Error in get_shares_info: {e}", exc_info=True)
raise ToolError(f"Failed to retrieve shares information: {str(e)}")
raise ToolError(f"Failed to retrieve shares information: {str(e)}") from e
@mcp.tool()
async def get_notifications_overview() -> Dict[str, Any]:
async def get_notifications_overview() -> dict[str, Any]:
"""Retrieves an overview of system notifications (unread and archive counts by severity)."""
query = """
query GetNotificationsOverview {
@@ -72,19 +73,20 @@ def register_storage_tools(mcp: FastMCP):
logger.info("Executing get_notifications_overview tool")
response_data = await make_graphql_request(query)
if response_data.get("notifications"):
return response_data["notifications"].get("overview", {})
overview = response_data["notifications"].get("overview", {})
return dict(overview) if isinstance(overview, dict) else {}
return {}
except Exception as e:
logger.error(f"Error in get_notifications_overview: {e}", exc_info=True)
raise ToolError(f"Failed to retrieve notifications overview: {str(e)}")
raise ToolError(f"Failed to retrieve notifications overview: {str(e)}") from e
@mcp.tool()
async def list_notifications(
type: str,
offset: int,
limit: int,
importance: Optional[str] = None
) -> List[Dict[str, Any]]:
type: str,
offset: int,
limit: int,
importance: str | None = None
) -> list[dict[str, Any]]:
"""Lists notifications with filtering. Type: UNREAD/ARCHIVE. Importance: INFO/WARNING/ALERT."""
query = """
query ListNotifications($filter: NotificationFilter!) {
@@ -114,19 +116,20 @@ def register_storage_tools(mcp: FastMCP):
# Remove null importance from variables if not provided, as GraphQL might be strict
if not importance:
del variables["filter"]["importance"]
try:
logger.info(f"Executing list_notifications: type={type}, offset={offset}, limit={limit}, importance={importance}")
response_data = await make_graphql_request(query, variables)
if response_data.get("notifications"):
return response_data["notifications"].get("list", [])
notifications_list = response_data["notifications"].get("list", [])
return list(notifications_list) if isinstance(notifications_list, list) else []
return []
except Exception as e:
logger.error(f"Error in list_notifications: {e}", exc_info=True)
raise ToolError(f"Failed to list notifications: {str(e)}")
raise ToolError(f"Failed to list notifications: {str(e)}") from e
@mcp.tool()
async def list_available_log_files() -> List[Dict[str, Any]]:
async def list_available_log_files() -> list[dict[str, Any]]:
"""Lists all available log files that can be queried."""
query = """
query ListLogFiles {
@@ -141,13 +144,14 @@ def register_storage_tools(mcp: FastMCP):
try:
logger.info("Executing list_available_log_files tool")
response_data = await make_graphql_request(query)
return response_data.get("logFiles", [])
log_files = response_data.get("logFiles", [])
return list(log_files) if isinstance(log_files, list) else []
except Exception as e:
logger.error(f"Error in list_available_log_files: {e}", exc_info=True)
raise ToolError(f"Failed to list available log files: {str(e)}")
raise ToolError(f"Failed to list available log files: {str(e)}") from e
@mcp.tool()
async def get_logs(log_file_path: str, tail_lines: int = 100) -> Dict[str, Any]:
async def get_logs(log_file_path: str, tail_lines: int = 100) -> dict[str, Any]:
"""Retrieves content from a specific log file, defaulting to the last 100 lines."""
# The Unraid GraphQL API Query.logFile takes 'lines' and 'startLine'.
# To implement 'tail_lines', we would ideally need to know the total lines first,
@@ -158,7 +162,7 @@ def register_storage_tools(mcp: FastMCP):
# If not, this tool might need to be smarter or the API might not directly support 'tail' easily.
# The SDL for LogFileContent implies it returns startLine, so it seems aware of ranges.
# Let's try fetching with just 'lines' to see if it acts as a tail,
# Let's try fetching with just 'lines' to see if it acts as a tail,
# or if we need Query.logFiles first to get totalLines for calculation.
# For robust tailing, one might need to fetch totalLines first, then calculate start_line for the tail.
# Simplified: query for the last 'tail_lines'. If the API doesn't support tailing this way, we may need adjustment.
@@ -178,16 +182,17 @@ def register_storage_tools(mcp: FastMCP):
try:
logger.info(f"Executing get_logs for {log_file_path}, tail_lines={tail_lines}")
response_data = await make_graphql_request(query, variables)
return response_data.get("logFile", {})
log_file = response_data.get("logFile", {})
return dict(log_file) if isinstance(log_file, dict) else {}
except Exception as e:
logger.error(f"Error in get_logs for {log_file_path}: {e}", exc_info=True)
raise ToolError(f"Failed to retrieve logs from {log_file_path}: {str(e)}")
raise ToolError(f"Failed to retrieve logs from {log_file_path}: {str(e)}") from e
@mcp.tool()
async def list_physical_disks() -> List[Dict[str, Any]]:
async def list_physical_disks() -> list[dict[str, Any]]:
"""Lists all physical disks recognized by the Unraid system."""
# Querying an extremely minimal set of fields for diagnostics
query = """
query = """
query ListPhysicalDisksMinimal {
disks {
id
@@ -199,15 +204,16 @@ def register_storage_tools(mcp: FastMCP):
try:
logger.info("Executing list_physical_disks tool with minimal query and increased timeout")
# Increased read timeout for this potentially slow query
long_timeout = httpx.Timeout(10.0, read=90.0, connect=5.0)
long_timeout = httpx.Timeout(10.0, read=90.0, connect=5.0)
response_data = await make_graphql_request(query, custom_timeout=long_timeout)
return response_data.get("disks", [])
disks = response_data.get("disks", [])
return list(disks) if isinstance(disks, list) else []
except Exception as e:
logger.error(f"Error in list_physical_disks: {e}", exc_info=True)
raise ToolError(f"Failed to list physical disks: {str(e)}")
raise ToolError(f"Failed to list physical disks: {str(e)}") from e
@mcp.tool()
async def get_disk_details(disk_id: str) -> Dict[str, Any]:
async def get_disk_details(disk_id: str) -> dict[str, Any]:
"""Retrieves detailed SMART information and partition data for a specific physical disk."""
# Enhanced query with more comprehensive disk information
query = """
@@ -227,19 +233,20 @@ def register_storage_tools(mcp: FastMCP):
logger.info(f"Executing get_disk_details for disk: {disk_id}")
response_data = await make_graphql_request(query, variables)
raw_disk = response_data.get("disk", {})
if not raw_disk:
raise ToolError(f"Disk '{disk_id}' not found")
# Process disk information for human-readable output
def format_bytes(bytes_value):
if bytes_value is None: return "N/A"
bytes_value = int(bytes_value)
def format_bytes(bytes_value: int | None) -> str:
if bytes_value is None:
return "N/A"
value = float(int(bytes_value))
for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:
if bytes_value < 1024.0:
return f"{bytes_value:.2f} {unit}"
bytes_value /= 1024.0
return f"{bytes_value:.2f} EB"
if value < 1024.0:
return f"{value:.2f} {unit}"
value /= 1024.0
return f"{value:.2f} EB"
summary = {
'disk_id': raw_disk.get('id'),
@@ -256,15 +263,15 @@ def register_storage_tools(mcp: FastMCP):
'partition_count': len(raw_disk.get('partitions', [])),
'total_partition_size': format_bytes(sum(p.get('size', 0) for p in raw_disk.get('partitions', []) if p.get('size')))
}
return {
'summary': summary,
'partitions': raw_disk.get('partitions', []),
'details': raw_disk
}
except Exception as e:
logger.error(f"Error in get_disk_details for {disk_id}: {e}", exc_info=True)
raise ToolError(f"Failed to retrieve disk details for {disk_id}: {str(e)}")
raise ToolError(f"Failed to retrieve disk details for {disk_id}: {str(e)}") from e
logger.info("Storage tools registered successfully")
logger.info("Storage tools registered successfully")