diff --git a/unraid_mcp/core/guards.py b/unraid_mcp/core/guards.py index 6a9088c..c0c929e 100644 --- a/unraid_mcp/core/guards.py +++ b/unraid_mcp/core/guards.py @@ -10,9 +10,8 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from fastmcp import Context -from fastmcp.exceptions import ToolError - from ..config.logging import logger +from .exceptions import ToolError async def elicit_destructive_confirmation( diff --git a/unraid_mcp/tools/array.py b/unraid_mcp/tools/array.py index 5d2d10c..8568573 100644 --- a/unraid_mcp/tools/array.py +++ b/unraid_mcp/tools/array.py @@ -11,7 +11,7 @@ from fastmcp import Context, FastMCP from ..config.logging import logger from ..core.client import make_graphql_request from ..core.exceptions import ToolError, tool_error_handler -from ..core.guards import elicit_destructive_confirmation +from ..core.guards import gate_destructive_action QUERIES: dict[str, str] = { @@ -159,18 +159,17 @@ def register_array_tool(mcp: FastMCP) -> None: if action not in ALL_ACTIONS: raise ToolError(f"Invalid action '{action}'. Must be one of: {sorted(ALL_ACTIONS)}") - if action in DESTRUCTIVE_ACTIONS and not confirm: - desc_map = { + await gate_destructive_action( + ctx, + action, + DESTRUCTIVE_ACTIONS, + confirm, + { "remove_disk": f"Remove disk **{disk_id}** from the array. The array must be stopped first.", "clear_disk_stats": f"Clear all I/O statistics for disk **{disk_id}**. This cannot be undone.", "stop_array": "Stop the Unraid array. Running containers and VMs may lose access to array shares.", - } - confirmed = await elicit_destructive_confirmation(ctx, action, desc_map[action]) - if not confirmed: - raise ToolError( - f"Action '{action}' was not confirmed. " - "Re-run with confirm=True to bypass elicitation." - ) + }, + ) with tool_error_handler("array", action, logger): logger.info(f"Executing unraid_array action={action}") diff --git a/unraid_mcp/tools/keys.py b/unraid_mcp/tools/keys.py index e5ae701..e0238bc 100644 --- a/unraid_mcp/tools/keys.py +++ b/unraid_mcp/tools/keys.py @@ -11,7 +11,7 @@ from fastmcp import Context, FastMCP from ..config.logging import logger from ..core.client import make_graphql_request from ..core.exceptions import ToolError, tool_error_handler -from ..core.guards import elicit_destructive_confirmation +from ..core.guards import gate_destructive_action QUERIES: dict[str, str] = { @@ -104,14 +104,13 @@ def register_keys_tool(mcp: FastMCP) -> None: if action not in ALL_ACTIONS: raise ToolError(f"Invalid action '{action}'. Must be one of: {sorted(ALL_ACTIONS)}") - if action in DESTRUCTIVE_ACTIONS and not confirm: - _desc = f"Delete API key **{key_id}**. Any clients using this key will lose access." - confirmed = await elicit_destructive_confirmation(ctx, action, _desc) - if not confirmed: - raise ToolError( - f"Action '{action}' was not confirmed. " - "Re-run with confirm=True to bypass elicitation." - ) + await gate_destructive_action( + ctx, + action, + DESTRUCTIVE_ACTIONS, + confirm, + f"Delete API key **{key_id}**. Any clients using this key will lose access.", + ) with tool_error_handler("keys", action, logger): logger.info(f"Executing unraid_keys action={action}") diff --git a/unraid_mcp/tools/plugins.py b/unraid_mcp/tools/plugins.py index 7b8bf53..836de07 100644 --- a/unraid_mcp/tools/plugins.py +++ b/unraid_mcp/tools/plugins.py @@ -10,7 +10,7 @@ from fastmcp import Context, FastMCP from ..config.logging import logger from ..core.client import make_graphql_request from ..core.exceptions import ToolError, tool_error_handler -from ..core.guards import elicit_destructive_confirmation +from ..core.guards import gate_destructive_action QUERIES: dict[str, str] = { @@ -75,14 +75,13 @@ def register_plugins_tool(mcp: FastMCP) -> None: if action not in ALL_ACTIONS: raise ToolError(f"Invalid action '{action}'. Must be one of: {sorted(ALL_ACTIONS)}") - if action in DESTRUCTIVE_ACTIONS and not confirm: - _desc = f"Remove plugin(s) **{names}** from the Unraid API. This cannot be undone without re-installing." - confirmed = await elicit_destructive_confirmation(ctx, action, _desc) - if not confirmed: - raise ToolError( - f"Action '{action}' was not confirmed. " - "Re-run with confirm=True to bypass elicitation." - ) + await gate_destructive_action( + ctx, + action, + DESTRUCTIVE_ACTIONS, + confirm, + f"Remove plugin(s) **{names}** from the Unraid API. This cannot be undone without re-installing.", + ) with tool_error_handler("plugins", action, logger): logger.info(f"Executing unraid_plugins action={action}")