mirror of
https://github.com/jmagar/unraid-mcp.git
synced 2026-03-01 16:04:24 -08:00
lintfree
This commit is contained in:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user