diff --git a/revit_mcp/code_execution.py b/revit_mcp/code_execution.py index faae198..4a3ae56 100644 --- a/revit_mcp/code_execution.py +++ b/revit_mcp/code_execution.py @@ -15,7 +15,11 @@ def register_code_execution_routes(api): - """Register code execution routes with the API.""" + """Register code execution routes with the API. + + ** WARNING ** : this can run 'ANY' code you give it, this can be a major security risk, only send it code that you trust. + + """ @api.route("/execute_code/", methods=["POST"]) def execute_code(doc, uidoc, request): @@ -26,10 +30,10 @@ def execute_code(doc, uidoc, request): { "code": "python code as string", "description": "optional description of what the code does", - "use_transaction": true # set false for UI ops like switching the active view } """ try: + # Parse the request data data = ( json.loads(request.data) @@ -84,7 +88,6 @@ def execute_code(doc, uidoc, request): except Exception as exec_error: sys.stdout = old_stdout partial_output = captured_output.getvalue() - captured_output.close() error_traceback = traceback.format_exc() error_type = type(exec_error).__name__ @@ -136,5 +139,8 @@ def execute_code(doc, uidoc, request): except Exception as e: logger.error("Execute code request failed: {}".format(str(e))) return routes.make_response(data={"error": str(e)}, status=500) + finally: + sys.stdout = old_stdout + captured_output.close() logger.info("Code execution routes registered successfully.") diff --git a/revit_mcp/commands.py b/revit_mcp/commands.py new file mode 100644 index 0000000..b655aae --- /dev/null +++ b/revit_mcp/commands.py @@ -0,0 +1,163 @@ +# -*- coding: UTF-8 -*- +""" +Commands Module for Revit MCP +Handles Command execution +""" +from pyrevit import HOST_APP +from pyrevit.coreutils.logger import get_logger +from pyrevit.loader import sessioninfo + +import glob +import os + +from pyrevit import routes +from pyrevit.routes.server import serverinfo + +from pyrevit import routes +import logging + +logger = logging.getLogger(__name__) + + +# Helper methods +def _find_active_log(): + # type: ()->(str, int) + """Return (log_path, current_byte_offset) for the most recent PyRevit log. + returned offset can be used to capture a range of text inside said log. + + Not 100% reliable when multiple revits are in play, and file logging must be enabled for this to work. + """ + appdata = os.environ.get('APPDATA', '') + if not appdata: + return None, 0 + pattern = os.path.join(appdata, 'pyRevit', '*', 'pyRevit_*_runtime.log') + matches = glob.glob(pattern) + if not matches: + return None, 0 + log_path = max(matches, key=os.path.getmtime) + try: + offset = os.path.getsize(log_path) + except OSError: + offset = 0 + return log_path, offset + + +def _read_log_since(log_path, offset): + # type: (str, int)->(str|None) + """reads log data from a file FROM an offset position""" + try: + with open(log_path, 'rb') as f: + f.seek(offset) + raw = f.read() + return raw.decode('utf-8', errors='replace') + except Exception: + return None + +def register_commands_routes(api): + """Register all commands-related routes with the API""" + + @api.route('/commands_list', methods=['GET']) + def get_commands(uiapp): + """List all loaded pyRevit commands with their control IDs.""" + from pyrevit.loader import sessionmgr + commands = sessionmgr.find_all_commands(cache=True) + return [ + { + "name": cmd.name, + "control_id": cmd.control_id, + "bundle": cmd.bundle, + "extension": cmd.extension, + "unique_id": cmd.unique_id, + } + for cmd in commands + ] + + @api.route('/commands_run', methods=['POST']) + def run_command(request, uiapp): + """Run a pyRevit command by control ID. + + Has 2 modes depending on provided body: + wait = true => runs execution immediatly and returns script logs (where possible) + Wait = false => posts command to revit ui thread, to be executed after this api call. + + Request body (JSON): { + "control_id": "CustomCtrl_%CustomCtrl_%...", + "wait": true (optional, default true — set false for fire-and-forget) + } + """ + from pyrevit.loader import sessionmgr + from datetime import datetime + data = request.data or {} + control_id = data.get('control_id', None) + wait = data.get('wait', True) + + if not control_id: + return {"error": "control_id is required in request body"} + + # fire-and-forget via PostCommand + if not wait: + command_id = UI.RevitCommandId.LookupCommandId(control_id) + if command_id is None: + return {"error": "Command not found: {}".format(control_id)} + uiapp.PostCommand(command_id) + return {"result": "posted", "control_id": control_id} + + + cmd = next((c for c in sessionmgr.find_all_commands() + if c.control_id == control_id), None) + if cmd is None: + return {"error": "Command not found: {}".format(control_id)} + + + # PyRevit reload destroys the IronPython engine mid-execution, so the HTTP + # response can never be sent. Waiting on it will always time out so we don't allow this + # command to be run with 'wait' + if cmd.unique_id == 'pyrevitcore_pyrevit_pyrevit_tools_reload': + return { + "error": "Cannot await PyRevit reload: the reload script destroys " + "the runtime engine before a response can be sent. " + "Use wait=false to fire-and-forget instead." + } + + # Snapshot log position before any execution so we can return only the + # lines produced by this command. + log_path, log_offset = _find_active_log() + + now = datetime.now() + envvars.set_pyrevit_env_var('PYREVIT_HEADLESS', '1') + result = None + except_info = None + try: + mlogger.debug('[HEADLESS:START] command=%s controlId=%s', cmd.unique_id, control_id) + result = sessionmgr.execute_command_cls(cmd.extcmd_type) + + except Exception as e: + except_info = except_info + logger.exception(e) + finally: + mlogger.debug('[HEADLESS:END] command=%s controlId=%s result=%s', cmd.unique_id, control_id, result) + envvars.set_pyrevit_env_var('PYREVIT_HEADLESS', '') + + # Read log lines produced between [HEADLESS:START] and [HEADLESS:END]. + execution_log = _read_log_since(log_path, log_offset) if log_path else None + + response = { + "result": str(result) if result else 'error', + "execution_time": str(datetime.now() - now), + "command": { + "name": cmd.name, + "control_id": cmd.control_id, + "bundle": cmd.bundle, + "extension": cmd.extension, + "unique_id": cmd.unique_id, + }, + } + if except_info is not None: + response["error"] = except_info + + if execution_log is not None: + response["output"] = execution_log + else: + response["output"] = "Logs are disabled or failed to collect. You can turn them on in the PyRevit Settings" + + return response diff --git a/startup.py b/startup.py index 477e61d..7a44e74 100644 --- a/startup.py +++ b/startup.py @@ -41,6 +41,10 @@ def register_routes(): register_code_execution_routes(api) + from revit_mcp.commands import register_commands_routes + + register_commands_routes(api) + logger.info("All MCP routes registered successfully") except Exception as e: diff --git a/tools/__init__.py b/tools/__init__.py index b787a86..402ddff 100644 --- a/tools/__init__.py +++ b/tools/__init__.py @@ -11,6 +11,7 @@ def register_tools(mcp_server, revit_get_func, revit_post_func, revit_image_func from .model_tools import register_model_tools from .colors_tools import register_colors_tools from .code_execution_tools import register_code_execution_tools + from .commands_tools import register_commands_tools # Register tools from each module register_status_tools(mcp_server, revit_get_func) @@ -21,3 +22,4 @@ def register_tools(mcp_server, revit_get_func, revit_post_func, revit_image_func register_code_execution_tools( mcp_server, revit_get_func, revit_post_func, revit_image_func ) + register_commands_tools(mcp_server, revit_get_func, revit_post_func) diff --git a/tools/commands_tools.py b/tools/commands_tools.py new file mode 100644 index 0000000..b12ec35 --- /dev/null +++ b/tools/commands_tools.py @@ -0,0 +1,37 @@ +# -*- coding: utf-8 -*- +"""Commands tools for the MCP server.""" +from mcp.server.fastmcp import Context +from .utils import format_response + + +def register_commands_tools(mcp, revit_get, revit_post, revit_image=None): + """Register your tools with the MCP server.""" + + # ---- Tool for the GET request ---- + @mcp.tool() + async def list_commands( + ctx: Context = None + ) -> str: + + """ + Return a list of all pyrevit commands, including control & uniqueid + """ + response = await revit_get("/commands_run", ctx) + return format_response(response) + + + # ---- Tool for the POST request ---- + @mcp.tool() + async def run_command_by_control_id(control_id: str, config:bool = False, wait:bool=False ,ctx: Context) -> str: + """ + Runs a command using its control_id + + Args: + control_id: The ID of the command, as found in the revit journal or list_commands tool + config: Run tool in config (shift+click) mode. + wait: Wait for tool to finish and return response + + """ + payload = {"control_id": control_id, "config": config, "wait" = wait} + response = await revit_post("/commands_run", payload, ctx) + return format_response(response)