From ed7bbff0ba04d58492432ef67b78313984b8edac Mon Sep 17 00:00:00 2001 From: brandonspark Date: Tue, 2 Sep 2025 10:30:53 -0700 Subject: [PATCH] remove tools based on env vars --- src/semgrep_mcp/server.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/semgrep_mcp/server.py b/src/semgrep_mcp/server.py index e7793ea..490a0cc 100755 --- a/src/semgrep_mcp/server.py +++ b/src/semgrep_mcp/server.py @@ -1088,6 +1088,34 @@ async def health(request: Request) -> JSONResponse: return JSONResponse({"status": "ok", "version": __version__}) +# --------------------------------------------------------------------------------- +# Disabling tools +# --------------------------------------------------------------------------------- + +TOOL_DISABLE_ENV_VARS = { + "SEMGREP_RULE_SCHEMA_DISABLED": "semgrep_rule_schema", + "GET_SUPPORTED_LANGUAGES_DISABLED": "get_supported_languages", + "SEMGREP_FINDINGS_DISABLED": "semgrep_findings", + "SEMGREP_SCAN_WITH_CUSTOM_RULE_DISABLED": "semgrep_scan_with_custom_rule", + "SEMGREP_SCAN_DISABLED": "semgrep_scan", + "SEMGREP_SCAN_LOCAL_DISABLED": "semgrep_scan_local", + "SECURITY_CHECK_DISABLED": "security_check", + "GET_ABSTRACT_SYNTAX_TREE_DISABLED": "get_abstract_syntax_tree", + "WRITE_CUSTOM_SEMGREP_RULE_DISABLED": "write_custom_semgrep_rule", +} + + +def deregister_tools() -> None: + for env_var, tool_name in TOOL_DISABLE_ENV_VARS.items(): + is_disabled = os.environ.get(env_var, "false").lower() == "true" + + if is_disabled: + # for the time being, while there is no way to API-level remove tools, + # we'll just mutate the internal `_tools`, because this language does + # not stop us from doing so + del mcp._tool_manager._tools[tool_name] + + # --------------------------------------------------------------------------------- # MCP Server Entry Point # --------------------------------------------------------------------------------- @@ -1127,6 +1155,9 @@ def main(transport: str, semgrep_path: str | None) -> None: if semgrep_path: set_semgrep_executable(semgrep_path) + # based on env vars, disable certain tools + deregister_tools() + if transport == "stdio": mcp.run(transport="stdio") elif transport == "streamable-http":