From d3e90649560b1c2de77d1ca609bfb4b2e70ec326 Mon Sep 17 00:00:00 2001 From: Maicon Berlofa Date: Thu, 23 Apr 2026 18:02:17 -0300 Subject: [PATCH 1/2] fix(server): harden dynamic tool discovery --- README.md | 24 ++++- src/server_builder.py | 186 ++++++++++++++++++++++++++--------- tests/test_server_builder.py | 59 +++++++++++ 3 files changed, 221 insertions(+), 48 deletions(-) diff --git a/README.md b/README.md index 44f8884..fb0ed58 100644 --- a/README.md +++ b/README.md @@ -127,7 +127,7 @@ For private repos, set `SOURCE_GIT_TOKEN` with a personal access token. ## Writing Tools -Create a `.py` file in `tools/`. Every public function becomes an MCP tool: +Create a `.py` file in `tools/`. By default, every public function defined in that module becomes an MCP tool: ```python def get_weather(city: str) -> str: @@ -141,6 +141,28 @@ def roll_dice(sides: int = 6) -> int: return random.randint(1, sides) ``` +Utility modules are supported too. Files named like `*_helpers.py` are skipped by default. For any other utility module, keep helpers private (`_helper`) or disable auto-registration explicitly: + +```python +__mcp_auto_register__ = False + +def evidence_true(data, *keys): + return True +``` + +If you want exact control over exports, declare a `TOOLS` allowlist: + +```python +TOOLS = ["deploy"] + +def deploy(service: str, version: str) -> str: + """Deploy a service.""" + return f"Deployed {service}@{version}" + +def helper() -> str: + return "not registered" +``` + ### Tool Metadata Add optional module-level variables to control tool registration: diff --git a/src/server_builder.py b/src/server_builder.py index c9e569a..6c1d198 100644 --- a/src/server_builder.py +++ b/src/server_builder.py @@ -11,6 +11,7 @@ """ import importlib.util +import inspect import logging import os import sys @@ -20,6 +21,101 @@ from fastmcp import FastMCP logger = logging.getLogger("fastmcp-server.builder") +HELPER_MODULE_SUFFIXES = ("_helpers",) + + +def _get_explicit_tool_names(module: types.ModuleType) -> list[str] | None: + """Return explicit tool export names when the module declares them.""" + explicit = getattr(module, "TOOLS", None) + if explicit is None: + return None + + if isinstance(explicit, (list, tuple, set)): + return [str(name) for name in explicit] + + logger.warning( + "Ignoring invalid TOOLS declaration in %s; expected list/tuple/set, got %s", + module.__name__, + type(explicit).__name__, + ) + return None + + +def _is_supported_tool_signature(func: types.FunctionType) -> tuple[bool, str]: + """Return whether a function signature is safe to register as an MCP tool.""" + try: + signature = inspect.signature(func) + except (TypeError, ValueError) as exc: + return False, f"signature inspection failed: {exc}" + + for parameter in signature.parameters.values(): + if parameter.kind == inspect.Parameter.VAR_POSITIONAL: + return False, "functions with *args are not supported" + if parameter.kind == inspect.Parameter.VAR_KEYWORD: + return False, "functions with **kwargs are not supported" + + return True, "" + + +def _iter_tool_candidates(module: types.ModuleType, py_file: Path): + """Yield tool candidates from a module using explicit exports or safe heuristics.""" + auto_register = getattr(module, "__mcp_auto_register__", True) + explicit_names = _get_explicit_tool_names(module) + module_stem = py_file.stem.lower() + + if not auto_register and explicit_names is None: + logger.info( + "Skipping tool auto-registration for %s (__mcp_auto_register__=False)", + py_file.name, + ) + return + + if explicit_names is None and ( + module_stem == "helpers" or module_stem.endswith(HELPER_MODULE_SUFFIXES) + ): + logger.info( + "Skipping helper module %s during tool auto-registration", + py_file.name, + ) + return + + names = explicit_names if explicit_names is not None else dir(module) + seen: set[str] = set() + + for name in names: + if name in seen: + continue + seen.add(name) + + if name.startswith("_"): + continue + + obj = getattr(module, name, None) + if not (callable(obj) and isinstance(obj, types.FunctionType)): + continue + + is_explicit = explicit_names is not None + + if not is_explicit and obj.__module__ != module.__name__: + logger.debug( + "Skipping imported function %s from %s (defined in %s)", + name, + py_file.name, + obj.__module__, + ) + continue + + supported, reason = _is_supported_tool_signature(obj) + if not supported: + logger.warning( + "Skipping function %s from %s: %s", + name, + py_file.name, + reason, + ) + continue + + yield name, obj def build_server(workspace: str) -> tuple[FastMCP, dict]: @@ -238,53 +334,49 @@ def _load_tools(mcp: FastMCP, tools_dir: Path, strict: bool = False) -> int: mod_cache_ttl = getattr(module, "__cache_ttl__", 0) registered = False - for name in dir(module): - if name.startswith("_"): - continue - obj = getattr(module, name) - if callable(obj) and isinstance(obj, types.FunctionType): - # v1.0.0: Apply sandboxing (memory + output limits) - obj = sandbox_tool(obj, name, mod_max_memory_mb, mod_max_output_size_kb) - - # v1.0.0: Apply rate limiting - obj = rate_limit_tool(obj, name, mod_rate_limit) - - # v1.0.0: Apply caching - if mod_cache_ttl > 0: - obj = cache_tool(obj, name, float(mod_cache_ttl)) - - # Instrument with metrics if enabled - if metrics_enabled(): - obj = instrument_tool(obj, name) - - # Build kwargs for mcp.tool() - tool_kwargs: dict = {} - if mod_tags is not None: - tool_kwargs["tags"] = set(mod_tags) - if mod_timeout is not None: - tool_kwargs["timeout"] = float(mod_timeout) - if mod_annotations is not None: - tool_kwargs["annotations"] = dict(mod_annotations) - - # Store required scopes in annotations (v0.7.0) - if mod_scopes is not None: - annotations = tool_kwargs.get("annotations", {}) - annotations["requiredScopes"] = list(mod_scopes) - tool_kwargs["annotations"] = annotations - - # Mark idempotent hint if cached (v1.0.0) - if mod_cache_ttl > 0: - annotations = tool_kwargs.get("annotations", {}) - annotations["idempotentHint"] = True - tool_kwargs["annotations"] = annotations - - if tool_kwargs: - mcp.tool(obj, **tool_kwargs) - else: - mcp.tool(obj) - logger.info("Registered tool: %s (from %s)", name, py_file.name) - count += 1 - registered = True + for name, obj in _iter_tool_candidates(module, py_file): + # v1.0.0: Apply sandboxing (memory + output limits) + obj = sandbox_tool(obj, name, mod_max_memory_mb, mod_max_output_size_kb) + + # v1.0.0: Apply rate limiting + obj = rate_limit_tool(obj, name, mod_rate_limit) + + # v1.0.0: Apply caching + if mod_cache_ttl > 0: + obj = cache_tool(obj, name, float(mod_cache_ttl)) + + # Instrument with metrics if enabled + if metrics_enabled(): + obj = instrument_tool(obj, name) + + # Build kwargs for mcp.tool() + tool_kwargs: dict = {} + if mod_tags is not None: + tool_kwargs["tags"] = set(mod_tags) + if mod_timeout is not None: + tool_kwargs["timeout"] = float(mod_timeout) + if mod_annotations is not None: + tool_kwargs["annotations"] = dict(mod_annotations) + + # Store required scopes in annotations (v0.7.0) + if mod_scopes is not None: + annotations = tool_kwargs.get("annotations", {}) + annotations["requiredScopes"] = list(mod_scopes) + tool_kwargs["annotations"] = annotations + + # Mark idempotent hint if cached (v1.0.0) + if mod_cache_ttl > 0: + annotations = tool_kwargs.get("annotations", {}) + annotations["idempotentHint"] = True + tool_kwargs["annotations"] = annotations + + if tool_kwargs: + mcp.tool(obj, **tool_kwargs) + else: + mcp.tool(obj) + logger.info("Registered tool: %s (from %s)", name, py_file.name) + count += 1 + registered = True if not registered: msg = f"No tools found in {py_file.name}" diff --git a/tests/test_server_builder.py b/tests/test_server_builder.py index 8aa5cd3..a10266a 100644 --- a/tests/test_server_builder.py +++ b/tests/test_server_builder.py @@ -23,6 +23,65 @@ def test_build_with_multiple_tools(workspace, monkeypatch): assert counts["tool_count"] == 2 +def test_helper_module_without_docstrings_is_not_registered(workspace, monkeypatch): + """Public helper functions without docstrings should not become tools.""" + monkeypatch.setenv("MCP_SERVER_NAME", "test-server") + (workspace / "tools" / "evidence_helpers.py").write_text( + "def evidence_true(data, *keys):\n return True\n\n" + "def notes_reviewed(data):\n return True\n" + ) + (workspace / "tools" / "real_tool.py").write_text( + 'def check_release_readiness() -> dict:\n """Check release readiness."""\n return {"status": "PASS"}\n' + ) + + mcp, counts = build_server(str(workspace)) + + assert counts["tool_count"] == 1 + + +def test_tool_manifest_registers_selected_functions(workspace, monkeypatch): + """TOOLS manifest should allow explicit registration from mixed modules.""" + monkeypatch.setenv("MCP_SERVER_NAME", "test-server") + (workspace / "tools" / "manifested.py").write_text( + 'TOOLS = ["hello"]\n\n' + 'def hello() -> str:\n return "hi"\n\n' + 'def helper() -> str:\n return "helper"\n' + ) + + mcp, counts = build_server(str(workspace)) + + assert counts["tool_count"] == 1 + + +def test_varargs_tool_is_skipped_instead_of_crashing(workspace, monkeypatch): + """Functions with *args or **kwargs should be skipped with no crash.""" + monkeypatch.setenv("MCP_SERVER_NAME", "test-server") + (workspace / "tools" / "bad_signature.py").write_text( + 'TOOLS = ["bad", "good"]\n\n' + 'def bad(*args) -> str:\n """Bad signature."""\n return "bad"\n\n' + 'def good(name: str) -> str:\n """Good signature."""\n return name\n' + ) + + mcp, counts = build_server(str(workspace)) + + assert counts["tool_count"] == 1 + + +def test_imported_function_is_not_auto_registered(workspace, monkeypatch): + """Imported helper functions should not become tools via auto-discovery.""" + monkeypatch.setenv("MCP_SERVER_NAME", "test-server") + (workspace / "tools" / "shared.py").write_text( + 'def shared_tool() -> str:\n """Shared tool."""\n return "shared"\n' + ) + (workspace / "tools" / "consumer.py").write_text( + "from shared import shared_tool\n" + ) + + mcp, counts = build_server(str(workspace)) + + assert counts["tool_count"] == 1 + + def test_build_with_resources(workspace, sample_resource, monkeypatch): """Resources are registered correctly.""" monkeypatch.setenv("MCP_SERVER_NAME", "test-server") From e0dd2525244af3f01491bf0ee3bbeb8b99fc98a7 Mon Sep 17 00:00:00 2001 From: Maicon Berlofa Date: Thu, 23 Apr 2026 18:03:55 -0300 Subject: [PATCH 2/2] style(tests): format server builder regression tests --- tests/test_server_builder.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_server_builder.py b/tests/test_server_builder.py index a10266a..e3d83b1 100644 --- a/tests/test_server_builder.py +++ b/tests/test_server_builder.py @@ -73,9 +73,7 @@ def test_imported_function_is_not_auto_registered(workspace, monkeypatch): (workspace / "tools" / "shared.py").write_text( 'def shared_tool() -> str:\n """Shared tool."""\n return "shared"\n' ) - (workspace / "tools" / "consumer.py").write_text( - "from shared import shared_tool\n" - ) + (workspace / "tools" / "consumer.py").write_text("from shared import shared_tool\n") mcp, counts = build_server(str(workspace))