From 675a466d02de85ca020269beaa030c0904e750d2 Mon Sep 17 00:00:00 2001 From: Jacob Magar Date: Sun, 15 Mar 2026 18:56:14 -0400 Subject: [PATCH] feat(live): add unraid_live tool with 11 subscription snapshot actions Creates unraid_mcp/tools/live.py with SNAPSHOT_ACTIONS (9 one-shot reads) and COLLECT_ACTIONS (2 streaming collectors), plus tests/test_live.py with 6 passing tests. Registers register_live_tool in server.py, bringing the total to 12 tools. --- tests/test_live.py | 95 ++++++++++++++++++++++ unraid_mcp/server.py | 2 + unraid_mcp/tools/live.py | 167 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 264 insertions(+) create mode 100644 tests/test_live.py create mode 100644 unraid_mcp/tools/live.py diff --git a/tests/test_live.py b/tests/test_live.py new file mode 100644 index 0000000..0624bcb --- /dev/null +++ b/tests/test_live.py @@ -0,0 +1,95 @@ +# tests/test_live.py +"""Tests for unraid_live subscription snapshot tool.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from fastmcp import FastMCP + + +@pytest.fixture +def mcp(): + return FastMCP("test") + + +def _make_live_tool(mcp): + from unraid_mcp.tools.live import register_live_tool + + register_live_tool(mcp) + local_provider = mcp.providers[0] + tool = local_provider._components["tool:unraid_live@"] + return tool.fn + + +@pytest.fixture +def _mock_subscribe_once(): + with patch("unraid_mcp.tools.live.subscribe_once") as m: + yield m + + +@pytest.fixture +def _mock_subscribe_collect(): + with patch("unraid_mcp.tools.live.subscribe_collect") as m: + yield m + + +@pytest.mark.asyncio +async def test_cpu_returns_snapshot(mcp, _mock_subscribe_once): + _mock_subscribe_once.return_value = {"systemMetricsCpu": {"percentTotal": 23.5, "cpus": []}} + tool_fn = _make_live_tool(mcp) + result = await tool_fn(action="cpu") + assert result["success"] is True + assert result["data"]["systemMetricsCpu"]["percentTotal"] == 23.5 + + +@pytest.mark.asyncio +async def test_memory_returns_snapshot(mcp, _mock_subscribe_once): + _mock_subscribe_once.return_value = { + "systemMetricsMemory": {"total": 32000000000, "used": 10000000000, "percentTotal": 31.2} + } + tool_fn = _make_live_tool(mcp) + result = await tool_fn(action="memory") + assert result["success"] is True + + +@pytest.mark.asyncio +async def test_log_tail_requires_path(mcp, _mock_subscribe_collect): + _mock_subscribe_collect.return_value = [] + tool_fn = _make_live_tool(mcp) + from unraid_mcp.core.exceptions import ToolError + + with pytest.raises(ToolError, match="path"): + await tool_fn(action="log_tail") + + +@pytest.mark.asyncio +async def test_log_tail_with_path(mcp, _mock_subscribe_collect): + _mock_subscribe_collect.return_value = [ + {"logFile": {"path": "/var/log/syslog", "content": "line1\nline2", "totalLines": 2}} + ] + tool_fn = _make_live_tool(mcp) + result = await tool_fn(action="log_tail", path="/var/log/syslog", collect_for=1.0) + assert result["success"] is True + assert result["event_count"] == 1 + + +@pytest.mark.asyncio +async def test_notification_feed_collects_events(mcp, _mock_subscribe_collect): + _mock_subscribe_collect.return_value = [ + {"notificationAdded": {"id": "1", "title": "Alert"}}, + {"notificationAdded": {"id": "2", "title": "Info"}}, + ] + tool_fn = _make_live_tool(mcp) + result = await tool_fn(action="notification_feed", collect_for=2.0) + assert result["event_count"] == 2 + + +@pytest.mark.asyncio +async def test_invalid_action_raises(mcp): + from unraid_mcp.core.exceptions import ToolError + + tool_fn = _make_live_tool(mcp) + with pytest.raises(ToolError, match="Invalid action"): + await tool_fn(action="nonexistent") # type: ignore[arg-type] diff --git a/unraid_mcp/server.py b/unraid_mcp/server.py index 9c151c7..13e5bc7 100644 --- a/unraid_mcp/server.py +++ b/unraid_mcp/server.py @@ -24,6 +24,7 @@ from .tools.docker import register_docker_tool from .tools.health import register_health_tool from .tools.info import register_info_tool from .tools.keys import register_keys_tool +from .tools.live import register_live_tool from .tools.notifications import register_notifications_tool from .tools.rclone import register_rclone_tool from .tools.settings import register_settings_tool @@ -64,6 +65,7 @@ def register_all_modules() -> None: register_keys_tool, register_health_tool, register_settings_tool, + register_live_tool, ] for registrar in registrars: registrar(mcp) diff --git a/unraid_mcp/tools/live.py b/unraid_mcp/tools/live.py new file mode 100644 index 0000000..caed254 --- /dev/null +++ b/unraid_mcp/tools/live.py @@ -0,0 +1,167 @@ +"""Real-time subscription snapshot tool. + +Provides the `unraid_live` tool with 11 actions — one per GraphQL +subscription. Each action opens a transient WebSocket, receives one event +(or collects events for `collect_for` seconds), then closes. + +Use `subscribe_once` actions for current-state reads (cpu, memory, array_state). +Use `subscribe_collect` actions for event streams (notification_feed, log_tail). +""" + +from __future__ import annotations + +from typing import Any, Literal, get_args + +from fastmcp import FastMCP + +from ..config.logging import logger +from ..core.exceptions import ToolError, tool_error_handler +from ..subscriptions.snapshot import subscribe_collect, subscribe_once + + +SNAPSHOT_ACTIONS = { + "cpu": """ + subscription { systemMetricsCpu { id percentTotal cpus { percentTotal percentUser percentSystem percentIdle } } } + """, + "memory": """ + subscription { systemMetricsMemory { id total used free available active buffcache percentTotal swapTotal swapUsed swapFree percentSwapTotal } } + """, + "cpu_telemetry": """ + subscription { systemMetricsCpuTelemetry { id totalPower power temp } } + """, + "array_state": """ + subscription { arraySubscription { id state capacity { kilobytes { free used total } } parityCheckStatus { status progress speed errors } } } + """, + "parity_progress": """ + subscription { parityHistorySubscription { date status progress speed errors correcting paused running } } + """, + "ups_status": """ + subscription { upsUpdates { id name model status battery { chargeLevel estimatedRuntime health } power { inputVoltage outputVoltage loadPercentage } } } + """, + "notifications_overview": """ + subscription { notificationsOverview { unread { info warning alert total } archive { info warning alert total } } } + """, + "owner": """ + subscription { ownerSubscription { username url avatar } } + """, + "server_status": """ + subscription { serversSubscription { id name status guid wanip lanip localurl remoteurl } } + """, +} + +COLLECT_ACTIONS = { + "notification_feed": """ + subscription { notificationAdded { id title subject description importance type timestamp } } + """, + "log_tail": """ + subscription LogTail($path: String!) { logFile(path: $path) { path content totalLines startLine } } + """, +} + +ALL_LIVE_ACTIONS = set(SNAPSHOT_ACTIONS) | set(COLLECT_ACTIONS) + +LIVE_ACTIONS = Literal[ + "array_state", + "cpu", + "cpu_telemetry", + "log_tail", + "memory", + "notification_feed", + "notifications_overview", + "owner", + "parity_progress", + "server_status", + "ups_status", +] + +if set(get_args(LIVE_ACTIONS)) != ALL_LIVE_ACTIONS: + _missing = ALL_LIVE_ACTIONS - set(get_args(LIVE_ACTIONS)) + _extra = set(get_args(LIVE_ACTIONS)) - ALL_LIVE_ACTIONS + raise RuntimeError( + f"LIVE_ACTIONS and ALL_LIVE_ACTIONS are out of sync. " + f"Missing: {_missing or 'none'}. Extra: {_extra or 'none'}" + ) + + +def register_live_tool(mcp: FastMCP) -> None: + """Register the unraid_live tool with the FastMCP instance.""" + + @mcp.tool() + async def unraid_live( + action: LIVE_ACTIONS, + path: str | None = None, + collect_for: float = 5.0, + timeout: float = 10.0, + ) -> dict[str, Any]: + """Get real-time data from Unraid via WebSocket subscriptions. + + Each action opens a transient WebSocket, receives data, then closes. + + Snapshot actions (return current state): + cpu - Real-time CPU utilization (all cores) + memory - Real-time memory and swap utilization + cpu_telemetry - CPU power draw and temperature per package + array_state - Live array state and parity status + parity_progress - Live parity check progress + ups_status - Real-time UPS battery and power state + notifications_overview - Live notification counts by severity + owner - Live owner info + server_status - Live server connection state + + Collection actions (collect events for `collect_for` seconds): + notification_feed - Collect new notification events (default: 5s window) + log_tail - Tail a log file (requires path; default: 5s window) + + Parameters: + path - Log file path for log_tail action (required) + collect_for - Seconds to collect events for collect actions (default: 5.0) + timeout - WebSocket connection/handshake timeout in seconds (default: 10.0) + """ + if action not in ALL_LIVE_ACTIONS: + raise ToolError( + f"Invalid action '{action}'. Must be one of: {sorted(ALL_LIVE_ACTIONS)}" + ) + + with tool_error_handler("live", action, logger): + logger.info(f"Executing unraid_live action={action} timeout={timeout}") + + if action in SNAPSHOT_ACTIONS: + data = await subscribe_once(SNAPSHOT_ACTIONS[action], timeout=timeout) + return {"success": True, "action": action, "data": data} + + # Collect actions + if action == "log_tail": + if not path: + raise ToolError("path is required for 'log_tail' action") + events = await subscribe_collect( + COLLECT_ACTIONS["log_tail"], + variables={"path": path}, + collect_for=collect_for, + timeout=timeout, + ) + return { + "success": True, + "action": action, + "path": path, + "collect_for": collect_for, + "event_count": len(events), + "events": events, + } + + if action == "notification_feed": + events = await subscribe_collect( + COLLECT_ACTIONS["notification_feed"], + collect_for=collect_for, + timeout=timeout, + ) + return { + "success": True, + "action": action, + "collect_for": collect_for, + "event_count": len(events), + "events": events, + } + + raise ToolError(f"Unhandled action '{action}' — this is a bug") + + logger.info("Live tool registered successfully")