From 9d705a326d1667b98502694f1f5959352880132a Mon Sep 17 00:00:00 2001 From: "Owen W. Taylor" Date: Fri, 8 May 2026 13:26:00 -0400 Subject: [PATCH] feat: redo tool filtering for better tests Filtering tools at server startup made it hard to test the effects of CONFIG.toolset. Switch to a system where tools are completely filtered in the middleware. (Doing it at the session level would be possible, but the direction for future protocol versions of MCP is to remove sessions entirely) This allows making the tests for server.py more comprehensive and greatly reduces the amount of mocking necessary. * All tests for server.py are moved into a new test_server.py. * Instead of making "fixed" the default toolset for tools without a tag, fixed tools are tagged just like run_script tools are. * toolset.Toolset is modified to only have tags for inclusion rather than exclusion tags as well. * A new setup_client fixture is added that is more flexible than the old mcp_client fixture - it's callable and the desired toolset and mcp-app support are passed in. * Our mcp-app HTML resource is properly filtered along with the tools. --- src/linux_mcp_server/server.py | 194 +++++++++---- src/linux_mcp_server/tools/logs.py | 4 +- src/linux_mcp_server/tools/network.py | 6 +- src/linux_mcp_server/tools/processes.py | 4 +- src/linux_mcp_server/tools/run_script.py | 6 +- src/linux_mcp_server/tools/services.py | 6 +- src/linux_mcp_server/tools/storage.py | 8 +- src/linux_mcp_server/tools/system_info.py | 10 +- src/linux_mcp_server/toolset.py | 26 +- tests/conftest.py | 83 +++++- tests/test_middleware.py | 302 ------------------- tests/test_server.py | 336 ++++++++++++++++++++++ tests/test_toolset.py | 39 +-- tests/tools/test_run_script.py | 117 ++++---- 14 files changed, 649 insertions(+), 492 deletions(-) delete mode 100644 tests/test_middleware.py create mode 100644 tests/test_server.py diff --git a/src/linux_mcp_server/server.py b/src/linux_mcp_server/server.py index ca4ec39e..3d3d5ffa 100644 --- a/src/linux_mcp_server/server.py +++ b/src/linux_mcp_server/server.py @@ -3,16 +3,20 @@ import logging import sys +from dataclasses import dataclass from importlib import resources from pathlib import Path +from fastmcp import Context from fastmcp import FastMCP +from fastmcp.exceptions import NotFoundError from fastmcp.resources import ResourceContent from fastmcp.resources import ResourceResult from fastmcp.server.dependencies import get_access_token from fastmcp.server.middleware import Middleware from fastmcp.server.middleware import MiddlewareContext from fastmcp.server.middleware.middleware import CallNext +from fastmcp.utilities.components import FastMCPComponent from mcp.types import InitializeRequest from mcp.types import InitializeRequestParams from mcp.types import InitializeResult @@ -29,6 +33,7 @@ from linux_mcp_server.mcp_app import MCP_UI_EXTENSION from linux_mcp_server.mcp_app import RUN_SCRIPT_APP_URI from linux_mcp_server.toolset import get_toolset +from linux_mcp_server.toolset import Toolset as ToolsetInfo def monkeypatch_fastmcp_for_app_visibility(): @@ -153,40 +158,42 @@ def monkeypatch_fastmcp_for_app_visibility(): """ -match CONFIG.toolset: - case Toolset.FIXED: - instructions = INSTRUCTIONS_FIXED - case Toolset.RUN_SCRIPT: - instructions = INSTRUCTIONS_RUN_SCRIPT - case Toolset.BOTH: - instructions = INSTRUCTIONS_BOTH - case _: - assert False, f"Unknown toolset configuration: {CONFIG.toolset}" +def _get_instructions() -> str: + match CONFIG.toolset: + case Toolset.FIXED: + return INSTRUCTIONS_FIXED + case Toolset.RUN_SCRIPT: + return INSTRUCTIONS_RUN_SCRIPT + case Toolset.BOTH: + return INSTRUCTIONS_BOTH + case _: # pragma: no cover + assert False, f"Unknown toolset configuration: {CONFIG.toolset}" -toolset = get_toolset(CONFIG.toolset.value) -assert toolset is not None, f"Toolset not found in registry: {CONFIG.toolset}" -if CONFIG.toolset != Toolset.FIXED and CONFIG.gatekeeper_model is None: - logger.error("LINUX_MCP_GATEKEEPER_MODEL not set, this is needed for run_script tools") - sys.exit(1) +def _current_toolset(): + toolset = get_toolset(CONFIG.toolset.value) + assert toolset is not None, f"Toolset not found in registry: {CONFIG.toolset}" -# Create auth provider if configured -auth_provider = create_auth_provider() + return toolset + + +def _check_gatekeeper_model(): + if CONFIG.toolset != Toolset.FIXED and CONFIG.gatekeeper_model is None: + logger.error("LINUX_MCP_GATEKEEPER_MODEL not set, this is needed for run_script tools") + sys.exit(1) -mcp = FastMCP("linux-mcp-server", instructions=instructions, version=linux_mcp_server.__version__, auth=auth_provider) -if toolset.include_tags: - mcp.enable(tags=toolset.include_tags, only=True) -else: - mcp.enable() +_check_gatekeeper_model() + +# Create auth provider if configured +auth_provider = create_auth_provider() -if toolset.exclude_tags: - mcp.disable(tags=toolset.exclude_tags) +mcp = FastMCP("linux-mcp-server", version=linux_mcp_server.__version__, auth=auth_provider) @mcp.resource( RUN_SCRIPT_APP_URI, - tags={"run_script"}, + tags={"run_script", "mcp_apps_only"}, ) def run_script_app_html() -> ResourceResult: filename = "run-script-app.html" @@ -238,6 +245,64 @@ def _use_mcp_app_for_client(client_params: InitializeRequestParams): return MCP_APP_MIME_TYPE in mime_types +def _hide_app_tools_for_client(client_params: InitializeRequestParams): + # Versions of goose before 1.29.0 don't understand _meta.ui.visiblity, so would + # leak our app-only tools to the model. However, they also are happy to let the + # model call tools that aren't listed as well. So if we see such an old version + # of goose, we strip out the app-only tools. + client_info = getattr(client_params, "clientInfo", None) + if client_info and client_info.name and client_info.name.startswith("goose"): + try: + major, minor = client_info.version.split(".")[0:2] + if (int(major), int(minor)) < (1, 29): + return True + except ValueError: + return False + + return False + + +@dataclass +class ComponentFilter: + """ + Determines what components (tools/resources) should be visible for a client + given the current config. + """ + + mcp_apps: bool + toolset: ToolsetInfo + hide_app_tools: bool + + def includes(self, component: FastMCPComponent): + if not self.toolset.includes_tool(component.tags): + return False + + if self.mcp_apps: + if "mcp_apps_exclude" in component.tags: + return False + else: + if "mcp_apps_only" in component.tags: + return False + + if self.hide_app_tools and "hidden_from_model" in component.tags: + return False + + return True + + @staticmethod + def get(context: Context, *, is_list_tools=False): + client_params = context.session.client_params + assert client_params is not None + mcp_apps = _use_mcp_app_for_client(client_params) + hide_app_tools = mcp_apps and is_list_tools and _hide_app_tools_for_client(client_params) + + return ComponentFilter( + mcp_apps=mcp_apps, + toolset=_current_toolset(), + hide_app_tools=hide_app_tools, + ) + + # Middleware to enforce authorization policy class AuthorizationMiddleware(Middleware): async def on_call_tool(self, context: MiddlewareContext, call_next): @@ -245,15 +310,12 @@ async def on_call_tool(self, context: MiddlewareContext, call_next): tool_args = context.message.arguments or {} target_host = tool_args.get("host") - # Get tool from FastMCP else fail closed if not found - if context.fastmcp_context is None: - logger.error("FastMCP context not available in middleware") - raise ValueError("Internal error: FastMCP context not available") + assert context.fastmcp_context tool = await context.fastmcp_context.fastmcp.get_tool(context.message.name) - if tool is None: - logger.error(f"Tool not found: {context.message.name}") - raise ValueError(f"Tool not found: {context.message.name}") + if tool is None or not ComponentFilter.get(context.fastmcp_context).includes(tool): + logger.error(f"Tool not found: '{context.message.name}'") + raise NotFoundError(f"Tool not found: '{context.message.name}'") # For stdio without policy configured, allow everything if CONFIG.transport == Transport.stdio and CONFIG.policy_path is None: @@ -298,12 +360,26 @@ async def on_call_tool(self, context: MiddlewareContext, call_next): if action == PolicyAction.SSH_KEY: if not ssh_key_config: raise RuntimeError("Policy validation error: SSH_KEY action requires ssh_key configuration.") - logger.debug(f"SSH key override: p`ath={ssh_key_config.path}, user={ssh_key_config.user}") + logger.debug(f"SSH key override: path={ssh_key_config.path}, user={ssh_key_config.user}") return await call_next(context) class DynamicDiscoveryMiddleware(Middleware): + """ + Implement a dynamic server that presents the right instructions, tools, and resources + depending on the client and the current configuration. + + Our configuration is logically static, but treating it dynamic makes it much + easier to write the appropriate tests. + + The dynamic behavior is done *per call*, not per session, since that's more future-proof. + (Future MCP protocol versions will remove sessions entirely: + https://modelcontextprotocol.io/seps/2575-stateless-mcp). This means doing it ourselves + rather than using FastMCP's system for session-level enabling and disabling components, + but it doesn't come out much worse. + """ + async def on_initialize( self, context: MiddlewareContext[InitializeRequest], @@ -315,51 +391,47 @@ async def on_initialize( # the client are fetched from the InitializationOptions object tucked # away in the ServerSession object, so we need to modify that based # on whether we'll use mcp-apps with the client making the InitializeRequest. + # + # For consistency and simplicity of testing, we always set the + # instructions this way assert context.fastmcp_context session = context.fastmcp_context.session - if _use_mcp_app_for_client(context.message.params): - instructions = session._init_options.instructions - if instructions: - session._init_options.instructions = instructions.replace( - "run_script_with_confirmation", "run_script_interactive" - ) + instructions = _get_instructions() + + toolset = _current_toolset() + if "run_script" in toolset.tags and _use_mcp_app_for_client(context.message.params): + instructions = instructions.replace("run_script_with_confirmation", "run_script_interactive") + + session._init_options.instructions = instructions return await call_next(context) async def on_list_tools(self, context: MiddlewareContext, call_next): tools = await call_next(context) - # Eventually, the tagging of the tools via _meta.ui.visiblity as "app" or "model" will - # hide this tool but Goose doesn't support this yet. On the other hand, goose is happy - # if the app calls tools we don't list at all, so we just filter out the "app" tools - filtered_tools = [t for t in tools if "hidden_from_model" not in (t.tags)] + assert context.fastmcp_context + filter = ComponentFilter.get(context.fastmcp_context, is_list_tools=True) + return [t for t in tools if filter.includes(t)] - fastmcp_context = context.fastmcp_context - assert fastmcp_context is not None, ( - "FastMCP framework error: context.fastmcp_context should not be None inside on_list_tools" - ) + async def on_list_resources(self, context: MiddlewareContext, call_next): + resources = await call_next(context) - request_ctx = fastmcp_context.request_context - assert request_ctx is not None, ( - "FastMCP framework error: request context should not be None inside on_list_tools" - ) + assert context.fastmcp_context + filter = ComponentFilter.get(context.fastmcp_context) + return [r for r in resources if filter.includes(r)] - client_params = request_ctx.session.client_params - assert client_params is not None, ( - "FastMCP framework error: client_params should not be None inside on_list_tools" - ) + # on_call_tool: the filtering for this is handled in AuthorizationMiddleware + # (the two middlewares could be combined) + # + # on_read_resource: we consider it harmless if any app reads the static and + # public mcp-app HTML, so we don't provide a on_read_resource() handler. - if _use_mcp_app_for_client(client_params): - filtered_tools = [t for t in filtered_tools if "mcp_apps_exclude" not in t.tags] - else: - filtered_tools = [t for t in filtered_tools if "mcp_apps_only" not in t.tags] - return filtered_tools +mcp.add_middleware(AuthorizationMiddleware()) +mcp.add_middleware(DynamicDiscoveryMiddleware()) def main(): - mcp.add_middleware(AuthorizationMiddleware()) - mcp.add_middleware(DynamicDiscoveryMiddleware()) mcp.run(show_banner=False, transport=CONFIG.transport.value, **CONFIG.transport_kwargs) diff --git a/src/linux_mcp_server/tools/logs.py b/src/linux_mcp_server/tools/logs.py index 7d0afdb5..b80ee0c4 100644 --- a/src/linux_mcp_server/tools/logs.py +++ b/src/linux_mcp_server/tools/logs.py @@ -81,7 +81,7 @@ async def _get_journal_logs( @mcp.tool( title="Get journal logs", description="Get systemd journal logs.", - tags={"journal", "logs", "systemd", "troubleshooting"}, + tags={"fixed", "journal", "logs", "systemd", "troubleshooting"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -143,7 +143,7 @@ async def get_journal_logs( @mcp.tool( title="Read log file", description="Read a specific log file.", - tags={"files", "logs", "troubleshooting"}, + tags={"fixed", "files", "logs", "troubleshooting"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/tools/network.py b/src/linux_mcp_server/tools/network.py index fee3f424..4b21f19b 100644 --- a/src/linux_mcp_server/tools/network.py +++ b/src/linux_mcp_server/tools/network.py @@ -20,7 +20,7 @@ @mcp.tool( title="Get network interfaces", description="Get detailed information about network interfaces including address and traffic statistics.", - tags={"connectivity", "interfaces", "network"}, + tags={"fixed", "connectivity", "interfaces", "network"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -56,7 +56,7 @@ async def get_network_interfaces( @mcp.tool( title="Get network connections", description="Get detailed information about active network connections.", - tags={"connections", "connectivity", "network"}, + tags={"fixed", "connections", "connectivity", "network"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -82,7 +82,7 @@ async def get_network_connections( @mcp.tool( title="Get listening ports", description="Get details on listening port, protocols, and services.", - tags={"connectivity", "network", "ports"}, + tags={"fixed", "connectivity", "network", "ports"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/tools/processes.py b/src/linux_mcp_server/tools/processes.py index de40e58d..63b417a1 100644 --- a/src/linux_mcp_server/tools/processes.py +++ b/src/linux_mcp_server/tools/processes.py @@ -20,7 +20,7 @@ @mcp.tool( title="List processes", description="List running processes", - tags={"performance", "processes"}, + tags={"fixed", "performance", "processes"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -45,7 +45,7 @@ async def list_processes( @mcp.tool( title="Process details", description="Get information about a specific process.", - tags={"performance", "processes"}, + tags={"fixed", "performance", "processes"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/tools/run_script.py b/src/linux_mcp_server/tools/run_script.py index 9b6ccf01..4993cb50 100644 --- a/src/linux_mcp_server/tools/run_script.py +++ b/src/linux_mcp_server/tools/run_script.py @@ -233,7 +233,7 @@ class ExecuteScriptResult: @mcp.tool( - tags={"run_script", "hidden_from_model"}, + tags={"run_script", "mcp_apps_only", "hidden_from_model"}, description="Execute a script; this is only available to the our mcp-app", app=AppConfig(visibility=["app"]), ) @@ -270,7 +270,7 @@ async def execute_script( @mcp.tool( - tags={"run_script", "hidden_from_model"}, + tags={"run_script", "mcp_apps_only", "hidden_from_model"}, description="Reject a script; this is only available to the our mcp-app", app=AppConfig(visibility=["app"]), ) @@ -357,7 +357,7 @@ async def run_script_interactive( @mcp.tool( - tags={"run_script", "hidden_from_model"}, + tags={"run_script", "mcp_apps_only", "hidden_from_model"}, title="Get the execution state with request ID", description="Get the execution state with request ID", app=AppConfig(visibility=["app"]), diff --git a/src/linux_mcp_server/tools/services.py b/src/linux_mcp_server/tools/services.py index 50b4d11a..41af8657 100644 --- a/src/linux_mcp_server/tools/services.py +++ b/src/linux_mcp_server/tools/services.py @@ -20,7 +20,7 @@ @mcp.tool( title="List services", description="List all systemd services.", - tags={"services", "systemd"}, + tags={"fixed", "services", "systemd"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -53,7 +53,7 @@ async def list_services( @mcp.tool( title="Get service status", description="Get detailed status of a specific systemd service.", - tags={"services", "systemd"}, + tags={"fixed", "services", "systemd"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -93,7 +93,7 @@ async def get_service_status( @mcp.tool( title="Get service logs", description="Get recent logs for a specific systemd service.", - tags={"logs", "services", "systemd", "troubleshooting"}, + tags={"fixed", "logs", "services", "systemd", "troubleshooting"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/tools/storage.py b/src/linux_mcp_server/tools/storage.py index 658dad73..aecf8eaf 100644 --- a/src/linux_mcp_server/tools/storage.py +++ b/src/linux_mcp_server/tools/storage.py @@ -79,7 +79,7 @@ async def _list_resources( @mcp.tool( title="List block devices", description="List block devices on the system", - tags={"devices", "storage"}, + tags={"fixed", "devices", "storage"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -104,7 +104,7 @@ async def list_block_devices( @mcp.tool( title="List directories", description="List directories under a specified path with various sorting options.", - tags={"directories", "filesystem", "storage"}, + tags={"fixed", "directories", "filesystem", "storage"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -149,7 +149,7 @@ async def list_directories( @mcp.tool( title="List files", description="List files under a specified path with various sorting options.", - tags={"files", "filesystem", "storage"}, + tags={"fixed", "files", "filesystem", "storage"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -194,7 +194,7 @@ async def list_files( @mcp.tool( title="Read file", description="Read the contents of a text file up to a safe size limit.", - tags={"files", "filesystem", "storage"}, + tags={"fixed", "files", "filesystem", "storage"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/tools/system_info.py b/src/linux_mcp_server/tools/system_info.py index f6df99f9..be05fc85 100644 --- a/src/linux_mcp_server/tools/system_info.py +++ b/src/linux_mcp_server/tools/system_info.py @@ -24,7 +24,7 @@ @mcp.tool( title="Get system information", description="Get basic system information such as operating system, distribution, kernel version, uptime, and last boot time.", - tags={"hardware", "system"}, + tags={"fixed", "hardware", "system"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -53,7 +53,7 @@ async def get_system_information( @mcp.tool( title="Get CPU information", description="Get CPU information.", - tags={"cpu", "hardware", "performance", "system"}, + tags={"fixed", "cpu", "hardware", "performance", "system"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -81,7 +81,7 @@ async def get_cpu_information( @mcp.tool( title="Get memory information", description="Get detailed memory including physical and swap.", - tags={"hardware", "memory", "performance", "system"}, + tags={"fixed", "hardware", "memory", "performance", "system"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -110,7 +110,7 @@ async def get_memory_information( @mcp.tool( title="Get disk usage", description="Get detailed disk space information including size, mount points, and utilization.", - tags={"disk", "filesystem", "storage", "system"}, + tags={"fixed", "disk", "filesystem", "storage", "system"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call @@ -143,7 +143,7 @@ async def get_disk_usage( @mcp.tool( title="Get hardware information", description="Get hardware information such as CPU details, PCI devices, USB devices, and hardware information from DMI.", - tags={"hardware", "system"}, + tags={"fixed", "hardware", "system"}, annotations=ToolAnnotations(readOnlyHint=True), ) @log_tool_call diff --git a/src/linux_mcp_server/toolset.py b/src/linux_mcp_server/toolset.py index 5c0e8ff8..d730bd1e 100644 --- a/src/linux_mcp_server/toolset.py +++ b/src/linux_mcp_server/toolset.py @@ -7,37 +7,29 @@ @dataclass class Toolset: name: str - include_tags: set[str] # Tool must have all of these tags - exclude_tags: set[str] # Tool must have none of these tags + tags: set[str] # Tool must have one of these tags def includes_tool(self, tool_tags: set[str]) -> bool: - # Tool must have all required tags - if self.include_tags and not self.include_tags.issubset(tool_tags): - return False - - # Tool must not have any excluded tags - if self.exclude_tags and self.exclude_tags.intersection(tool_tags): - return False - - return True + # Tool must have one of the required tags + return not self.tags.isdisjoint(tool_tags) # Registry of available toolsets for policy matching _TOOLSETS = { "fixed": Toolset( name="fixed", - include_tags=set(), - exclude_tags={"run_script"}, + tags={"fixed"}, ), "run_script": Toolset( name="run_script", - include_tags={"run_script"}, - exclude_tags=set(), + tags={"run_script"}, ), "both": Toolset( name="both", - include_tags=set(), - exclude_tags=set(), + tags={ + "fixed", + "run_script", + }, ), } diff --git a/tests/conftest.py b/tests/conftest.py index 0ebccc65..78d00a53 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,9 @@ import os +from contextlib import AsyncExitStack +from contextlib import contextmanager +from unittest.mock import patch + # Register script tools on the in-process MCP instance used by ``mcp_client``. # Default CLI/config is FIXED-only; tests need validate_script / run_script / etc. @@ -9,8 +13,12 @@ import pytest from fastmcp.client import Client +from mcp.types import Implementation +from mcp.types import InitializeRequest from linux_mcp_server.audit import log_tool_call +from linux_mcp_server.config import CONFIG +from linux_mcp_server.config import Toolset from linux_mcp_server.server import mcp @@ -21,10 +29,79 @@ def clean_env(monkeypatch): monkeypatch.delenv(var, raising=False) +@contextmanager +def _with_mcp_app_client(): + """Context manager that runs the body so that the first created FastMCP client + in the body will have extended capabilities marking it as a mcp-app client. + + This requires a complicated dance to patch the InitializeRequest() constructor + but unpatch it before the in-process server side deserializes the request. + """ + OriginalInitializeRequest = InitializeRequest + patcher = patch("mcp.types.InitializeRequest") + + def intercept_init_request(*args, **kwargs): + # 3. SELF-DESTRUCT: Stop the patch immediately upon first call. + # This restores the real class globally before the Server ever sees the request. + patcher.stop() + + params = kwargs.get("params") + if params and hasattr(params, "capabilities"): + if not hasattr(params.capabilities, "extensions"): + params.capabilities.extensions = {} + params.capabilities.extensions["io.modelcontextprotocol/ui"] = {"mimeTypes": ["text/html;profile=mcp-app"]} + + return OriginalInitializeRequest(*args, **kwargs) + + mocked_class = patcher.start() + mocked_class.side_effect = intercept_init_request + + try: + yield + finally: + patcher.stop() + + +@pytest.fixture +async def setup_client(mocker): + # The AsyncExitStack allows us to flexibly clean things up when the + # test exits, even when we can't represent them as a with: stack + # (in particular, we want to add the cleanups when setup_fn is + # called, not when the fixture is called.) + async with AsyncExitStack() as stack: + is_setup = False + + async def setup_fn( + *, + toolset: Toolset = Toolset.FIXED, + mcp_apps: bool = False, + auto_initialize: bool = True, + client_info: Implementation | None = None, + ): + nonlocal is_setup + if is_setup: + raise RuntimeError("setup_client can only be called once per test.") + is_setup = True + + if mcp_apps: + # Patch things so the client advertises mcp-app support + stack.enter_context(_with_mcp_app_client()) + + # Patch the toolset + stack.enter_context(patch.object(CONFIG, "toolset", toolset)) + + client = Client(transport=mcp, auto_initialize=auto_initialize, client_info=client_info) + # Automatically connect to the client so that the caller + # doesn't need a with: block - the client will be disconnected + # via the AsyncExitStack when the test exits. + return await stack.enter_async_context(client) + + yield setup_fn + + @pytest.fixture -async def mcp_client(): - async with Client(transport=mcp) as mcp_client: - yield mcp_client +async def mcp_client(setup_client): + yield await setup_client(toolset=Toolset.FIXED, mcp_apps=False) @pytest.fixture diff --git a/tests/test_middleware.py b/tests/test_middleware.py deleted file mode 100644 index 9bc0d1e8..00000000 --- a/tests/test_middleware.py +++ /dev/null @@ -1,302 +0,0 @@ -import importlib - -import pytest - -from mcp import ServerSession -from mcp.types import InitializeRequestParams -from mcp.types import Tool - - -# Workaround: with python-3.10, mocker.patch("linux_mcp_server.server.X") -# doesn't work because it finds the imported server function rather than the module. -server_module = importlib.import_module("linux_mcp_server.server") - - -class TestUseMcpAppForClient: - @pytest.mark.parametrize("config_value,expected", [(True, True), (False, False)]) - def test_config_override(self, mocker, config_value, expected): - # Test CONFIG.use_mcp_apps override - mocker.patch.object(server_module, "CONFIG", use_mcp_apps=config_value) - result = server_module._use_mcp_app_for_client(mocker.Mock(spec=InitializeRequestParams)) - assert result is expected - - @pytest.mark.parametrize( - "extensions,expected", - [ - ({server_module.MCP_UI_EXTENSION: {"mimeTypes": [server_module.MCP_APP_MIME_TYPE]}}, True), - ({"other-extension": {}}, False), - ], - ) - def test_detection(self, mocker, extensions, expected): - # Test mcp-app detection from client capabilities - mocker.patch.object(server_module, "CONFIG", use_mcp_apps=None) - capabilities = mocker.Mock() - capabilities.extensions = extensions - params = mocker.Mock(spec=InitializeRequestParams) - params.capabilities = capabilities - result = server_module._use_mcp_app_for_client(params) - assert result is expected - - -class TestDynamicDiscoveryMiddlewareOnInitialize: - @pytest.fixture - def middleware(self): - return server_module.DynamicDiscoveryMiddleware() - - @pytest.fixture - def mock_context(self, mocker): - # Create a mock MiddlewareContext with message.params - context = mocker.Mock() - context.message = mocker.Mock() - context.message.params = mocker.Mock(spec=InitializeRequestParams) - context.fastmcp_context = mocker.Mock() - return context - - async def test_skips_modification_when_disabled(self, middleware, mock_context, mocker): - # Test that instructions are not modified when mcp-apps is disabled - mocker.patch.object(server_module, "_use_mcp_app_for_client", return_value=False) - - mock_session = mocker.Mock(spec=ServerSession) - mock_session._init_options = mocker.Mock() - mock_session._init_options.instructions = "Use run_script_with_confirmation for changes" - mock_context.fastmcp_context.session = mock_session - - await middleware.on_initialize(mock_context, mocker.AsyncMock(return_value=mocker.Mock())) - - assert mock_session._init_options.instructions == "Use run_script_with_confirmation for changes" - - async def test_modifies_instructions(self, middleware, mock_context, mocker): - mocker.patch.object(server_module, "_use_mcp_app_for_client", return_value=True) - - mock_session = mocker.Mock(spec=ServerSession) - mock_session._init_options = mocker.Mock() - mock_session._init_options.instructions = "Use run_script_with_confirmation for changes" - mock_context.fastmcp_context.session = mock_session - - await middleware.on_initialize(mock_context, mocker.AsyncMock(return_value=mocker.Mock())) - - assert mock_session._init_options.instructions == "Use run_script_interactive for changes" - - -class TestDynamicDiscoveryMiddlewareOnListTools: - @pytest.fixture - def middleware(self): - return server_module.DynamicDiscoveryMiddleware() - - @pytest.fixture - def mock_context(self, mocker): - # Create a mock context with full request_context chain - context = mocker.Mock() - fastmcp_context = mocker.Mock() - request_ctx = mocker.Mock() - session = mocker.Mock() - client_params = mocker.Mock(spec=InitializeRequestParams) - - session.client_params = client_params - request_ctx.session = session - fastmcp_context.request_context = request_ctx - context.fastmcp_context = fastmcp_context - - return context - - def _make_tool(self, mocker, name: str, tags: set[str] | None = None) -> Tool: - # Helper to create a Tool with tags - tool = mocker.Mock(spec=Tool) - tool.name = name - tool.tags = tags or set() # type: ignore - return tool - - async def test_filters_when_mcp_apps_enabled(self, middleware, mock_context, mocker): - # Test that hidden_from_model and mcp_apps_exclude are filtered when enabled - mocker.patch.object(server_module, "_use_mcp_app_for_client", return_value=True) - - tools = [ - self._make_tool(mocker, "regular", set()), - self._make_tool(mocker, "hidden", {"hidden_from_model"}), - self._make_tool(mocker, "excluded", {"mcp_apps_exclude"}), - self._make_tool(mocker, "app_only", {"mcp_apps_only"}), - ] - - result = await middleware.on_list_tools(mock_context, mocker.AsyncMock(return_value=tools)) - - assert len(result) == 2 - result_names = {tool.name for tool in result} - assert result_names == {"regular", "app_only"} - - async def test_filters_when_mcp_apps_disabled(self, middleware, mock_context, mocker): - # Test that hidden_from_model and mcp_apps_only are filtered when disabled - mocker.patch.object(server_module, "_use_mcp_app_for_client", return_value=False) - - tools = [ - self._make_tool(mocker, "regular", set()), - self._make_tool(mocker, "hidden", {"hidden_from_model"}), - self._make_tool(mocker, "excluded", {"mcp_apps_exclude"}), - self._make_tool(mocker, "app_only", {"mcp_apps_only"}), - ] - - result = await middleware.on_list_tools(mock_context, mocker.AsyncMock(return_value=tools)) - - assert len(result) == 2 - result_names = {tool.name for tool in result} - assert result_names == {"regular", "excluded"} - - -class TestAuthorizationMiddlewareOnCallTool: - @pytest.fixture - def middleware(self): - return server_module.AuthorizationMiddleware() - - @pytest.fixture - def mock_context(self, mocker): - # Create a mock MiddlewareContext for tool calls - context = mocker.Mock() - context.message = mocker.Mock() - context.message.name = "test_tool" - context.message.arguments = {} - - # Mock FastMCP context with get_tool - fastmcp_context = mocker.Mock() - tool = mocker.Mock() - tool.tags = set() - fastmcp_context.fastmcp.get_tool = mocker.AsyncMock(return_value=tool) - context.fastmcp_context = fastmcp_context - - # Mock user_info - context.user_info = None - - return context - - async def test_no_auth_stdio_transport_allows(self, middleware, mock_context, mocker): - # No auth configured, stdio transport with no policy, should allow - mocker.patch.object( - server_module, "CONFIG", auth=None, transport=server_module.Transport.stdio, policy_path=None - ) - - call_next = mocker.AsyncMock(return_value="result") - result = await middleware.on_call_tool(mock_context, call_next) - - assert result == "result" - call_next.assert_called_once() - - async def test_no_auth_http_with_all_users(self, middleware, mock_context, mocker): - # No auth, HTTP, but policy allows all_users - mocker.patch.object(server_module, "CONFIG", auth=None, transport=server_module.Transport.http) - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.LOCAL, None), - ) - - call_next = mocker.AsyncMock(return_value="result") - result = await middleware.on_call_tool(mock_context, call_next) - - assert result == "result" - - async def test_with_auth_policy_deny(self, middleware, mock_context, mocker): - # Auth configured, policy returns DENY - mock_context.user_info = {"email": "user@example.com"} - mocker.patch.object(server_module, "CONFIG", auth=mocker.Mock(), transport=server_module.Transport.http) - - # Mock get_access_token to return authenticated user - mock_access_token = mocker.Mock() - mock_access_token.claims = {"email": "user@example.com"} - mocker.patch("linux_mcp_server.server.get_access_token", return_value=mock_access_token) - - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.DENY, None), - ) - - call_next = mocker.AsyncMock() - - with pytest.raises(ValueError, match="Authorization denied"): - await middleware.on_call_tool(mock_context, call_next) - - async def test_with_auth_policy_local(self, middleware, mock_context, mocker): - # Auth configured, policy allows LOCAL action - mock_context.user_info = {"email": "user@example.com"} - mocker.patch.object(server_module, "CONFIG", auth=mocker.Mock(), transport=server_module.Transport.http) - - # Mock get_access_token to return authenticated user - mock_access_token = mocker.Mock() - mock_access_token.claims = {"email": "user@example.com"} - mocker.patch("linux_mcp_server.server.get_access_token", return_value=mock_access_token) - - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.LOCAL, None), - ) - - call_next = mocker.AsyncMock(return_value="result") - result = await middleware.on_call_tool(mock_context, call_next) - - assert result == "result" - - async def test_ssh_default_with_no_host_raises(self, middleware, mock_context, mocker): - # SSH_DEFAULT action but no host in arguments - mock_context.user_info = {"email": "user@example.com"} - mocker.patch.object(server_module, "CONFIG", auth=mocker.Mock(), transport=server_module.Transport.http) - - # Mock get_access_token to return authenticated user - mock_access_token = mocker.Mock() - mock_access_token.claims = {"email": "user@example.com"} - mocker.patch("linux_mcp_server.server.get_access_token", return_value=mock_access_token) - - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.SSH_DEFAULT, None), - ) - - call_next = mocker.AsyncMock() - - with pytest.raises(RuntimeError, match="Policy validation error: Cannot use SSH action"): - await middleware.on_call_tool(mock_context, call_next) - - async def test_ssh_key_with_config_allows(self, middleware, mock_context, mocker): - # SSH_KEY action with config, should allow and log - mock_context.message.arguments = {"host": "db-server"} - mock_context.user_info = {"email": "admin@example.com"} - - ssh_key_config = mocker.Mock() - ssh_key_config.path = "/keys/db-key" - ssh_key_config.user = "db-admin" - - mocker.patch.object(server_module, "CONFIG", auth=mocker.Mock(), transport=server_module.Transport.http) - - # Mock get_access_token to return authenticated user - mock_access_token = mocker.Mock() - mock_access_token.claims = {"email": "admin@example.com"} - mocker.patch("linux_mcp_server.server.get_access_token", return_value=mock_access_token) - - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.SSH_KEY, ssh_key_config), - ) - mock_logger = mocker.patch.object(server_module, "logger") - - call_next = mocker.AsyncMock(return_value="result") - result = await middleware.on_call_tool(mock_context, call_next) - - assert result == "result" - # Verify SSH key override was logged - mock_logger.debug.assert_called() - - async def test_local_action_with_host_raises(self, middleware, mock_context, mocker): - # LOCAL action but host is specified - mock_context.message.arguments = {"host": "remote-host"} - mock_context.user_info = {"email": "user@example.com"} - mocker.patch.object(server_module, "CONFIG", auth=mocker.Mock(), transport=server_module.Transport.http) - - # Mock get_access_token to return authenticated user - mock_access_token = mocker.Mock() - mock_access_token.claims = {"email": "user@example.com"} - mocker.patch("linux_mcp_server.server.get_access_token", return_value=mock_access_token) - - mocker.patch( - "linux_mcp_server.server.evaluate_policy", - return_value=(server_module.PolicyAction.LOCAL, None), - ) - - call_next = mocker.AsyncMock() - - with pytest.raises(RuntimeError, match="Policy validation error: Cannot use local action for remote host"): - await middleware.on_call_tool(mock_context, call_next) diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 00000000..197b1ada --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,336 @@ +"""tests for the server module, including middleware""" + +import logging + +from typing import Any +from typing import Literal +from typing import Protocol + +import pytest + +from fastmcp import Client +from fastmcp.client import FastMCPTransport +from fastmcp.exceptions import ToolError +from fastmcp.server.auth import AccessToken +from mcp.types import Implementation + +from linux_mcp_server.auth_policy import AuthPolicy +from linux_mcp_server.auth_policy import PolicyAction +from linux_mcp_server.auth_policy import PolicyRule +from linux_mcp_server.auth_policy import SSHKeyConfig +from linux_mcp_server.config import Toolset + + +FIXED_TOOLS = set( + [ + "get_cpu_information", + "get_disk_usage", + "get_hardware_information", + "get_journal_logs", + "get_listening_ports", + "get_memory_information", + "get_network_connections", + "get_network_interfaces", + "get_process_info", + "get_service_logs", + "get_service_status", + "get_system_information", + "list_block_devices", + "list_directories", + "list_files", + "list_processes", + "list_services", + "read_file", + "read_log_file", + ] +) + +RUN_SCRIPT_TOOLS = set(["run_script", "run_script_with_confirmation", "validate_script"]) +RUN_SCRIPT_APP_TOOLS = set( + [ + "execute_script", + "get_execution_state", + "reject_script", + "run_script_interactive", + "run_script", + "validate_script", + ] +) + + +@pytest.mark.parametrize( + "toolset,mcp_apps,expected", + [ + (Toolset.FIXED, False, FIXED_TOOLS), + (Toolset.FIXED, True, FIXED_TOOLS), + (Toolset.RUN_SCRIPT, False, RUN_SCRIPT_TOOLS), + (Toolset.RUN_SCRIPT, True, RUN_SCRIPT_APP_TOOLS), + (Toolset.BOTH, False, FIXED_TOOLS | RUN_SCRIPT_TOOLS), + (Toolset.BOTH, True, FIXED_TOOLS | RUN_SCRIPT_APP_TOOLS), + ], +) +async def test_list_tools(toolset: Toolset, mcp_apps: bool, expected: set[str], setup_client): + client = await setup_client(toolset=toolset, mcp_apps=mcp_apps) + tools = await client.list_tools() + tool_names = set(tool.name for tool in tools) + assert tool_names == expected + + +async def test_list_tools_with_use_mcp_apps_override(setup_client, mocker): + """Test that LINUX_MCP_USE_MCP_APPS turns off detected mcp-apps""" + mocker.patch("linux_mcp_server.server.CONFIG.use_mcp_apps", False) + client = await setup_client(toolset=Toolset.RUN_SCRIPT, mcp_apps=True) + tools = await client.list_tools() + tool_names = set(tool.name for tool in tools) + assert tool_names == RUN_SCRIPT_TOOLS + + +@pytest.mark.parametrize("version,list_app_only_tools", [("1.23.1", False), ("1.29", True), ("2beta", True)]) +async def test_list_tools_with_old_goose(setup_client, version: str, list_app_only_tools: bool): + """Test that visibility=["app"] tools are hidden from old versions of Goose""" + client_info = Implementation(name="goose-deskktop", version=version) + client = await setup_client(toolset=Toolset.RUN_SCRIPT, mcp_apps=True, client_info=client_info) + tools = await client.list_tools() + tool_names = set(tool.name for tool in tools) + if list_app_only_tools: + assert tool_names == RUN_SCRIPT_APP_TOOLS + else: + assert tool_names == {"run_script_interactive", "run_script", "validate_script"} + + +@pytest.mark.parametrize( + "toolset,mcp_apps,expected", + [ + (Toolset.FIXED, True, set()), + (Toolset.RUN_SCRIPT, False, set()), + (Toolset.RUN_SCRIPT, True, {"ui://run_script_readonly_with_mcp_app/run-script-app.html"}), + ], +) +async def test_list_resources(toolset: Toolset, mcp_apps: bool, expected: set[str], setup_client): + client: Client = await setup_client(toolset=toolset, mcp_apps=mcp_apps) + resources = await client.list_resources() + tool_names = set(str(resource.uri) for resource in resources) + assert tool_names == expected + + +@pytest.mark.parametrize( + "toolset,mcp_apps,expected_substrings", + [ + (Toolset.FIXED, False, ["You have access to predefined commands"]), + ( + Toolset.RUN_SCRIPT, + False, + [ + "You have access to tools that validate and execute", + "**run_script_with_confirmation:** Run a validated script", + ], + ), + ( + Toolset.RUN_SCRIPT, + True, + [ + "You have access to tools that validate and execute", + "**run_script_interactive:** Run a validated script", + ], + ), + ( + Toolset.BOTH, + False, + [ + "You have access to two kinds of tools", + "**run_script_with_confirmation:** Run a validated script", + ], + ), + ( + Toolset.BOTH, + True, + ["You have access to two kinds of tools", "**run_script_interactive:** Run a validated script"], + ), + ], +) +async def test_instructions(toolset: Toolset, mcp_apps: bool, expected_substrings: list[str], setup_client): + client = await setup_client(toolset=toolset, mcp_apps=mcp_apps, auto_initialize=False) + result = await client.initialize() + for substring in expected_substrings: + assert substring in result.instructions + + +async def test_read_mcp_app_resource(mcp_client): + """Test that we can read the mcp-app resource""" + result = await mcp_client.read_resource("ui://run_script_readonly_with_mcp_app/run-script-app.html") + + assert len(result) == 1 + assert "Run Script" in result[0].text + + +FREE_OUTPUT = """\ + total used free shared buff/cache available +Mem: 32490308 14925036 2497076 2138132 15082720 17565272 +Swap: 8388604 3555312 4833292 +""" + + +class TestAuthorizationMiddleware: + @pytest.fixture(autouse=True) + def setup(self, mocker, mcp_client: Client[FastMCPTransport]): + """A callable fixture to encapsulate the mocks we need for these tests""" + + # Mock out a single tool to use in the tests + mock_command = mocker.Mock() + mock_command.run = mocker.AsyncMock(return_value=(0, FREE_OUTPUT, "")) + mocker.patch("linux_mcp_server.tools.system_info.get_command", return_value=mock_command) + + def setup_fn( + transport: Literal["stdio", "http"] = "stdio", + policy: AuthPolicy | None = None, + claims: dict[str, Any] | None = None, + ): + mocker.patch("linux_mcp_server.server.CONFIG.transport", transport) + + if policy: + policy_path = "/etc/linux-mcp-server/auth_policy.json" + else: + policy_path = None + policy = AuthPolicy(rules=[]) + mocker.patch("linux_mcp_server.server.CONFIG.policy_path", policy_path) + mocker.patch("linux_mcp_server.auth_policy.get_policy", return_value=policy) + + if claims is not None: + access_token = mocker.Mock(spec=AccessToken) + access_token.claims = claims + mocker.patch("linux_mcp_server.server.get_access_token", return_value=access_token) + + return mcp_client + + yield setup_fn + + # pyright struggles to infer the type of the fixture, so add an explicit type + class SetupFn(Protocol): + def __call__( + self, + transport: Literal["stdio", "http"] = "stdio", + policy: AuthPolicy | None = None, + claims: dict[str, Any] | None = None, + ) -> Client[FastMCPTransport]: ... + + async def test_unknown_tool_raises(self, setup: SetupFn): + client = setup(transport="stdio") + with pytest.raises(ToolError, match=r"Unknown tool: 'not_a_tool'"): + await client.call_tool("not_a_tool") + + async def test_disabled_tool_raises(self, setup: SetupFn): + client = setup(transport="stdio") + with pytest.raises(ToolError, match=r"Unknown tool: 'validate_script'"): + await client.call_tool("validate_script") + + async def test_stdio_no_policy_allows(self, setup: SetupFn): + client = setup(transport="stdio") + result = await client.call_tool("get_memory_information") + assert result.structured_content and result.structured_content["ram"]["free"] == 2497076 + + async def test_http_no_policy_rejects(self, setup: SetupFn): + client = setup(transport="http") + with pytest.raises(ToolError, match=r"Authorization denied: tool 'get_memory_information'"): + await client.call_tool("get_memory_information") + + async def test_http_policy_all_users_allows(self, setup: SetupFn): + client = setup( + transport="http", + policy=AuthPolicy( + rules=[PolicyRule(host="localhost", tools=["@fixed"], all_users=True, action=PolicyAction.LOCAL)] + ), + ) + + result = await client.call_tool("get_memory_information") + assert result.structured_content and result.structured_content["ram"]["free"] == 2497076 + + @pytest.mark.parametrize("email,allowed", [("user1@example.com", True), ("user2@example.com", False)]) + async def test_http_policy_one_user(self, email: str, allowed: bool, setup: SetupFn): + client = setup( + transport="http", + policy=AuthPolicy( + rules=[ + PolicyRule( + host="localhost", + tools=["@fixed"], + claims={"email": "user1@example.com"}, + action=PolicyAction.LOCAL, + ) + ] + ), + claims={"email": email}, + ) + + if allowed: + result = await client.call_tool("get_memory_information") + assert result.structured_content and result.structured_content["ram"]["free"] == 2497076 + else: + with pytest.raises(ToolError, match=r"Authorization denied: tool 'get_memory_information'"): + await client.call_tool("get_memory_information") + + async def test_ssh_key_with_config_allows(self, setup: SetupFn, caplog): + """SSH_KEY action with config, should allow and log""" + caplog.set_level(logging.DEBUG) + + client = setup( + transport="http", + policy=AuthPolicy( + rules=[ + PolicyRule( + host="server1.example.com", + tools=["@fixed"], + all_users=True, + action=PolicyAction.SSH_KEY, + ssh_key=SSHKeyConfig(path="/keys/server1.key", user="serviceaccount"), + ) + ] + ), + ) + + result = await client.call_tool("get_memory_information", {"host": "server1.example.com"}) + assert result.structured_content and result.structured_content["ram"]["free"] == 2497076 + + assert "SSH key override: path=/keys/server1.key, user=serviceaccount" in caplog.text + + # Various scenarios for bad returns from evaluate_policy that should be prevented + # by policy validation. We mock the evaluate_policy() return directly + + async def test_local_action_with_host_raises(self, setup: SetupFn, mocker): + """LOCAL action but host is specified""" + client = setup(transport="http") + + mocker.patch( + "linux_mcp_server.server.evaluate_policy", + return_value=(PolicyAction.LOCAL, None), + ) + + with pytest.raises( + ToolError, match=r"Policy validation error: Cannot use local action for remote host 'server1.example.com'" + ): + await client.call_tool("get_memory_information", {"host": "server1.example.com"}) + + async def test_ssh_default_with_no_host_raises(self, setup: SetupFn, mocker): + """SSH_DEFAULT action but no host in arguments""" + client = setup(transport="http") + + mocker.patch( + "linux_mcp_server.server.evaluate_policy", + return_value=(PolicyAction.SSH_DEFAULT, None), + ) + + with pytest.raises( + ToolError, match=r"Policy validation error: Cannot use SSH action \('ssh_default'\) for local execution." + ): + await client.call_tool("get_memory_information") + + async def test_ssh_key_with_no_config_raises(self, setup: SetupFn, mocker): + """SSH_KEY action but no ssk_key configuration""" + client = setup(transport="http") + + mocker.patch( + "linux_mcp_server.server.evaluate_policy", + return_value=(PolicyAction.SSH_KEY, None), + ) + + with pytest.raises(ToolError, match=r"Policy validation error: SSH_KEY action requires ssh_key configuration."): + await client.call_tool("get_memory_information", {"host": "server1.example.com"}) diff --git a/tests/test_toolset.py b/tests/test_toolset.py index e598b5a7..d6b19879 100644 --- a/tests/test_toolset.py +++ b/tests/test_toolset.py @@ -7,8 +7,7 @@ def test_includes_tool_with_required_tags(self): # Toolset requires "run_script" tag toolset = Toolset( name="run_script", - include_tags={"run_script"}, - exclude_tags=set(), + tags={"run_script"}, ) # Tool has the required tag @@ -19,60 +18,32 @@ def test_includes_tool_missing_required_tags(self): # Toolset requires "run_script" tag toolset = Toolset( name="run_script", - include_tags={"run_script"}, - exclude_tags=set(), + tags={"run_script"}, ) # Tool doesn't have the required tag assert not toolset.includes_tool(set()) assert not toolset.includes_tool({"other_tag"}) - def test_includes_tool_with_excluded_tags(self): - # Toolset excludes "run_script" tag - toolset = Toolset( - name="fixed", - include_tags=set(), - exclude_tags={"run_script"}, - ) - - # Tool has the excluded tag - assert not toolset.includes_tool({"run_script"}) - assert not toolset.includes_tool({"run_script", "other_tag"}) - - def test_includes_tool_without_excluded_tags(self): - # Toolset excludes "run_script" tag - toolset = Toolset( - name="fixed", - include_tags=set(), - exclude_tags={"run_script"}, - ) - - # Tool doesn't have the excluded tag - assert toolset.includes_tool(set()) - assert toolset.includes_tool({"other_tag"}) - class TestGetToolset: def test_get_fixed_toolset(self): toolset = get_toolset("fixed") assert toolset is not None assert toolset.name == "fixed" - assert toolset.include_tags == set() - assert toolset.exclude_tags == {"run_script"} + assert toolset.tags == {"fixed"} def test_get_run_script_toolset(self): toolset = get_toolset("run_script") assert toolset is not None assert toolset.name == "run_script" - assert toolset.include_tags == {"run_script"} - assert toolset.exclude_tags == set() + assert toolset.tags == {"run_script"} def test_get_both_toolset(self): toolset = get_toolset("both") assert toolset is not None assert toolset.name == "both" - assert toolset.include_tags == set() - assert toolset.exclude_tags == set() + assert toolset.tags == {"fixed", "run_script"} def test_get_unknown_toolset(self): toolset = get_toolset("unknown") diff --git a/tests/tools/test_run_script.py b/tests/tools/test_run_script.py index 14a33a1d..89dc9cac 100644 --- a/tests/tools/test_run_script.py +++ b/tests/tools/test_run_script.py @@ -1,4 +1,4 @@ -"""Integration-style tests for ``run_script`` MCP tools via ``mcp_client``. +"""Integration-style tests for ``run_script`` MCP tools. Patches apply to ``linux_mcp_server.tools.run_script`` (the module object). The in-process server is built with ``LINUX_MCP_TOOLSET=both`` (see ``tests/conftest.py``). @@ -14,6 +14,7 @@ from fastmcp.exceptions import ToolError +from linux_mcp_server.config import Toolset from linux_mcp_server.connection.ssh import execute_command from linux_mcp_server.gatekeeper import GatekeeperResult from linux_mcp_server.gatekeeper import GatekeeperStatus @@ -33,6 +34,16 @@ def _tool_text(result: Any) -> str: return str(result.structured_content["result"]) +@pytest.fixture +async def client(setup_client): + yield await setup_client(toolset=Toolset.RUN_SCRIPT) + + +@pytest.fixture +async def app_client(setup_client): + yield await setup_client(toolset=Toolset.RUN_SCRIPT, mcp_apps=True) + + @pytest.fixture def script_store_fresh(monkeypatch: pytest.MonkeyPatch) -> ScriptStore: """Isolate ``script_store`` so tests do not share global script IDs.""" @@ -118,11 +129,11 @@ def test_bash_includes_strict_preamble_in_quoted_payload(self) -> None: class TestValidateScriptMCP: - """``validate_script`` through ``mcp_client.call_tool``.""" + """``validate_script`` through ``client.call_tool``.""" async def test_ok( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, monkeypatch: pytest.MonkeyPatch, @@ -130,7 +141,7 @@ async def test_ok( """Approval stores the script; structured content has ``token`` and ``needs_confirmation``.""" _stub_secrets_token(monkeypatch, "val-id") patch_check_run_script.return_value = _ok() - result = await mcp_client.call_tool( + result = await client.call_tool( "validate_script", { "description": "d", @@ -149,7 +160,7 @@ async def test_ok( async def test_gatekeeper_fail_raises_and_marks_rejected( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, monkeypatch: pytest.MonkeyPatch, @@ -158,7 +169,7 @@ async def test_gatekeeper_fail_raises_and_marks_rejected( _stub_secrets_token(monkeypatch, "val-id-2") patch_check_run_script.return_value = GatekeeperResult(status=GatekeeperStatus.POLICY, detail="bad") with pytest.raises(ToolError, match="Policy violation"): - await mcp_client.call_tool( + await client.call_tool( "validate_script", { "description": "d", @@ -171,7 +182,7 @@ async def test_gatekeeper_fail_raises_and_marks_rejected( async def test_bash_passes_strict_script_to_gatekeeper( self, - mcp_client: Any, + client: Any, patch_check_run_script: Any, patch_execute_command: Any, monkeypatch: pytest.MonkeyPatch, @@ -179,7 +190,7 @@ async def test_bash_passes_strict_script_to_gatekeeper( """Bash scripts are checked with the same strict preamble the runner will use.""" _stub_secrets_token(monkeypatch, "bash-val") patch_check_run_script.return_value = _ok() - await mcp_client.call_tool( + await client.call_tool( "validate_script", { "description": "d", @@ -198,11 +209,11 @@ async def test_bash_passes_strict_script_to_gatekeeper( class TestRunScriptMCP: - """``run_script`` (token only) via ``mcp_client``.""" + """``run_script`` (token only) via ``client``.""" async def test_success_string_stdout( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -216,14 +227,14 @@ async def test_success_string_stdout( readonly=True, ) patch_execute_command.return_value = (0, "output", "") - result = await mcp_client.call_tool("run_script", {"token": "tok"}) + result = await client.call_tool("run_script", {"token": "tok"}) assert _tool_text(result) == "output" patch_execute_command.assert_awaited_once() assert script_store_fresh.get_script_details("tok").state == "success" async def test_success_bytes_stdout( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -237,12 +248,12 @@ async def test_success_bytes_stdout( readonly=True, ) patch_execute_command.return_value = (0, "café".encode("utf-8"), "") - result = await mcp_client.call_tool("run_script", {"token": "tokb"}) + result = await client.call_tool("run_script", {"token": "tokb"}) assert _tool_text(result) == "café" async def test_nonzero_return( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -256,14 +267,14 @@ async def test_nonzero_return( readonly=True, ) patch_execute_command.return_value = (1, "", "err") - out = _tool_text(await mcp_client.call_tool("run_script", {"token": "toknz"})) + out = _tool_text(await client.call_tool("run_script", {"token": "toknz"})) assert "return code 1" in out assert "err" in out assert script_store_fresh.get_script_details("toknz").state == "failure" async def test_execute_exception_sets_failure( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -278,12 +289,12 @@ async def test_execute_exception_sets_failure( ) patch_execute_command.side_effect = ValueError("nope") with pytest.raises(ToolError, match="nope"): - await mcp_client.call_tool("run_script", {"token": "tokex"}) + await client.call_tool("run_script", {"token": "tokex"}) assert script_store_fresh.get_script_details("tokex").state == "failure" async def test_wrong_tool_when_confirmation_required( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -297,16 +308,16 @@ async def test_wrong_tool_when_confirmation_required( readonly=False, ) with pytest.raises(ToolError, match="run_script_with_confirmation"): - await mcp_client.call_tool("run_script", {"token": "mod"}) + await client.call_tool("run_script", {"token": "mod"}) patch_execute_command.assert_not_awaited() class TestRunScriptWithConfirmationMCP: - """``run_script_with_confirmation`` via ``mcp_client``.""" + """``run_script_with_confirmation`` via ``client``.""" async def test_success_matching_params( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -322,7 +333,7 @@ async def test_success_matching_params( ) patch_execute_command.return_value = (0, "ran", "") out = _tool_text( - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "same", @@ -339,7 +350,7 @@ async def test_success_matching_params( async def test_mismatch_revalidates_and_executes( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -356,7 +367,7 @@ async def test_mismatch_revalidates_and_executes( patch_check_run_script.return_value = _ok() patch_execute_command.return_value = (0, "out", "") out = _tool_text( - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "same", @@ -372,7 +383,7 @@ async def test_mismatch_revalidates_and_executes( async def test_mismatch_gatekeeper_fails( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -388,7 +399,7 @@ async def test_mismatch_gatekeeper_fails( ) patch_check_run_script.return_value = GatekeeperResult(status=GatekeeperStatus.MALICIOUS, detail="x") with pytest.raises(ToolError): - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "same", @@ -403,7 +414,7 @@ async def test_mismatch_gatekeeper_fails( async def test_nonzero_return( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -419,7 +430,7 @@ async def test_nonzero_return( ) patch_execute_command.return_value = (3, "", "stderr-here") out = _tool_text( - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "same", @@ -435,7 +446,7 @@ async def test_nonzero_return( async def test_wrong_tool_when_no_confirmation_needed( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -450,7 +461,7 @@ async def test_wrong_tool_when_no_confirmation_needed( readonly=True, ) with pytest.raises(ToolError, match="run_script instead"): - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "d", @@ -465,7 +476,7 @@ async def test_wrong_tool_when_no_confirmation_needed( async def test_bytes_stdout( self, - mcp_client: Any, + client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, patch_check_run_script: Any, @@ -482,7 +493,7 @@ async def test_bytes_stdout( raw = b"\xff\xfe" patch_execute_command.return_value = (0, raw, "") out = _tool_text( - await mcp_client.call_tool( + await client.call_tool( "run_script_with_confirmation", { "description": "same", @@ -497,11 +508,11 @@ async def test_bytes_stdout( class TestRunScriptInteractiveMCP: - """``run_script_interactive`` via ``mcp_client``.""" + """``run_script_interactive`` via ``app_client``.""" async def test_ok_matching_returns_same_token( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, ) -> None: @@ -514,7 +525,7 @@ async def test_ok_matching_returns_same_token( host=None, readonly=False, ) - result = await mcp_client.call_tool( + result = await app_client.call_tool( "run_script_interactive", { "description": "d", @@ -530,7 +541,7 @@ async def test_ok_matching_returns_same_token( async def test_mismatch_revalidates_new_id( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, monkeypatch: pytest.MonkeyPatch, @@ -546,7 +557,7 @@ async def test_mismatch_revalidates_new_id( ) patch_check_run_script.return_value = _ok() _stub_secrets_token(monkeypatch, "new-id") - result = await mcp_client.call_tool( + result = await app_client.call_tool( "run_script_interactive", { "description": "d", @@ -562,7 +573,7 @@ async def test_mismatch_revalidates_new_id( async def test_mismatch_gatekeeper_fails( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, ) -> None: @@ -577,7 +588,7 @@ async def test_mismatch_gatekeeper_fails( ) patch_check_run_script.return_value = GatekeeperResult(status=GatekeeperStatus.DANGEROUS, detail="no") with pytest.raises(ToolError, match="Dangerous"): - await mcp_client.call_tool( + await app_client.call_tool( "run_script_interactive", { "description": "d", @@ -591,7 +602,7 @@ async def test_mismatch_gatekeeper_fails( async def test_wrong_tool_when_no_confirmation_needed( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_check_run_script: Any, ) -> None: @@ -605,7 +616,7 @@ async def test_wrong_tool_when_no_confirmation_needed( readonly=True, ) with pytest.raises(ToolError, match="run_script instead"): - await mcp_client.call_tool( + await app_client.call_tool( "run_script_interactive", { "description": "d", @@ -619,11 +630,11 @@ async def test_wrong_tool_when_no_confirmation_needed( class TestExecuteScriptMCP: - """``execute_script`` via ``mcp_client``.""" + """``execute_script`` via ``app_client``.""" async def test_success( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -637,13 +648,13 @@ async def test_success( readonly=False, ) patch_execute_command.return_value = (0, "out", "") - result = await mcp_client.call_tool("execute_script", {"id": "tok"}) + result = await app_client.call_tool("execute_script", {"id": "tok"}) assert result.structured_content == {"state": "success", "output": "out"} assert script_store_fresh.get_script_details("tok").state == "success" async def test_failure_return_code( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -657,13 +668,13 @@ async def test_failure_return_code( readonly=False, ) patch_execute_command.return_value = (2, "", "stderr-msg") - result = await mcp_client.call_tool("execute_script", {"id": "tok2"}) + result = await app_client.call_tool("execute_script", {"id": "tok2"}) assert result.structured_content["state"] == "failure" assert "return code 2" in result.structured_content["output"] async def test_execute_exception_sets_failure( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, patch_execute_command: Any, ) -> None: @@ -678,16 +689,16 @@ async def test_execute_exception_sets_failure( ) patch_execute_command.side_effect = OSError("boom") with pytest.raises(ToolError, match="boom"): - await mcp_client.call_tool("execute_script", {"id": "tok3"}) + await app_client.call_tool("execute_script", {"id": "tok3"}) assert script_store_fresh.get_script_details("tok3").state == "failure" class TestRejectAndGetExecutionStateMCP: - """``reject_script`` and ``get_execution_state`` via ``mcp_client``.""" + """``reject_script`` and ``get_execution_state`` via ``app_client``.""" async def test_reject_script( self, - mcp_client: Any, + app_client: Any, script_store_fresh: ScriptStore, ) -> None: """``reject_script`` moves the stored entry to ``rejected-user``.""" @@ -699,10 +710,10 @@ async def test_reject_script( host=None, readonly=True, ) - await mcp_client.call_tool("reject_script", {"id": "r"}) + await app_client.call_tool("reject_script", {"id": "r"}) assert script_store_fresh.get_script_details("r").state == "rejected-user" - async def test_get_execution_state(self, mcp_client: Any, script_store_fresh: ScriptStore) -> None: + async def test_get_execution_state(self, app_client: Any, script_store_fresh: ScriptStore) -> None: """Expose the current lifecycle state string for UI polling.""" script_store_fresh._scripts["g"] = ScriptDetails( state="executing", @@ -712,5 +723,5 @@ async def test_get_execution_state(self, mcp_client: Any, script_store_fresh: Sc host="host", readonly=True, ) - result = await mcp_client.call_tool("get_execution_state", {"id": "g"}) + result = await app_client.call_tool("get_execution_state", {"id": "g"}) assert result.structured_content == {"state": "executing"}