diff --git a/README.md b/README.md index 2b438d05..419cfcb9 100644 --- a/README.md +++ b/README.md @@ -221,6 +221,8 @@ The following table lists the available MCP functions for use: These are the list of HTTP endpoints that can be called: +- `/health`: Report server thread health, active request count, binary availability, and the last server error. + - `/allStrings`: All strings in one response. - `/formatValue?address=&text=&size=`: Convert and set a comment at an address. - `/getXrefsTo?address=`: Xrefs to address (code+data). diff --git a/plugin/__init__.py b/plugin/__init__.py index 2fd7b33f..bfdcd4ba 100644 --- a/plugin/__init__.py +++ b/plugin/__init__.py @@ -1,3 +1,5 @@ +import atexit + import binaryninja as bn from binaryninja import Settings @@ -47,7 +49,7 @@ def start_server(self, bv): _show_no_bv_popup() return # Avoid duplicate starts - if self.server and self.server.server: + if self.server and self.server.is_running: bn.log_info("MCP Max server already running; skip new start") # Ensure BV is set if not already if self.server.binary_ops.current_view is None: @@ -83,7 +85,7 @@ def start_server(self, bv): def stop_server(self, bv): try: # If not running, inform the user - if not (self.server and self.server.server): + if not (self.server and self.server.is_running): bn.log_info("MCP Max server stop requested but server is not running") _show_popup("MCP Server", "Server is not running.") return @@ -181,6 +183,34 @@ def _show_no_bv_popup(): _status_container = None _indicator_timer = None _bv_monitor_timer = None +_runtime_shutdown = False + + +def _shutdown_runtime(): + """Stop plugin-owned runtime resources before Binary Ninja tears down.""" + global _runtime_shutdown + if _runtime_shutdown: + return + _runtime_shutdown = True + + for timer in (_indicator_timer, _bv_monitor_timer): + try: + if timer is not None and hasattr(timer, "stop"): + timer.stop() + except Exception: + pass + + try: + if plugin.server and plugin.server.is_running: + plugin.server.stop() + except Exception as exc: + try: + bn.log_debug(f"MCP shutdown cleanup failed: {exc}") + except Exception: + pass + + +atexit.register(_shutdown_runtime) def _sidebar_icon_margin_default() -> int: @@ -266,7 +296,7 @@ def _create(): # Set initial visible state so the indicator shows up immediately try: - running_now = bool(plugin.server and plugin.server.server) + running_now = bool(plugin.server and plugin.server.is_running) except Exception: running_now = False if running_now: @@ -283,7 +313,7 @@ def _create(): # Click handler to toggle server state def _on_click(): try: - running = bool(plugin.server and plugin.server.server) + running = bool(plugin.server and plugin.server.is_running) if running: plugin.stop_server(None) else: @@ -304,7 +334,7 @@ def _on_click(): return plugin.start_server(bv) finally: - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) _status_button.clicked.connect(_on_click) @@ -376,7 +406,7 @@ def _start_indicator_watcher(): def _tick(): try: _ensure_status_indicator() - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) if _status_button is not None and hasattr(_indicator_timer, "stop"): _indicator_timer.stop() except Exception: @@ -410,7 +440,7 @@ def _schedule_status_init(): def _init_once(): try: _ensure_status_indicator() - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) except Exception: pass @@ -603,7 +633,7 @@ class _MCPMaxUINotification(ui.UIContextNotification): def __init__(self): super().__init__() ui.UIContext.registerNotification(self) - + def _get_active_bv(self): try: ctx = ui.UIContext.activeContext() @@ -620,7 +650,7 @@ def OnViewChange(self, *args): # signature varies across versions bv = self._get_active_bv() # Ensure the status indicator exists and reflects current state _ensure_status_indicator() - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) _start_indicator_watcher() _start_bv_monitor() if bv: @@ -639,7 +669,7 @@ def OnAfterOpenFile(self, *args): # type: ignore[override] bv = self._get_active_bv() # Ensure the status indicator is present as soon as a file opens _ensure_status_indicator() - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) _start_indicator_watcher() _start_bv_monitor() if bv: @@ -677,11 +707,19 @@ def OnAfterCloseFile(self, *args): # type: ignore[override] def OnContextOpen(self, *args): # type: ignore[override] try: _ensure_status_indicator() - _set_status_indicator(bool(plugin.server and plugin.server.server)) + _set_status_indicator(bool(plugin.server and plugin.server.is_running)) _start_indicator_watcher() except Exception: pass + def OnContextClose(self, *args): # type: ignore[override] + try: + all_contexts = getattr(ui.UIContext, "allContexts", None) + if callable(all_contexts) and len(list(all_contexts())) <= 1: + _shutdown_runtime() + except Exception: + pass + notification = _MCPMaxUINotification() bn.log_info("MCP Max UI notifications installed") # Ensure status control is present at startup with retries @@ -732,7 +770,7 @@ def _kick_autostart(): def _is_server_running() -> bool: try: - return bool(plugin.server and plugin.server.server) + return bool(plugin.server and plugin.server.is_running) except Exception: return False diff --git a/plugin/core/binary_operations.py b/plugin/core/binary_operations.py index 1ac91e18..14876b28 100644 --- a/plugin/core/binary_operations.py +++ b/plugin/core/binary_operations.py @@ -220,7 +220,7 @@ def list_open_binaries(self) -> list[dict[str, str]]: vb_canon = vb entries.append((canonical_id, fn, bool(vb_canon is self._current_view))) # Sort by filename for stable ordering - entries.sort(key=lambda t: (t[1] or "")) + entries.sort(key=lambda t: t[1] or "") for cid, fn, active in entries: items.append({"id": cid, "filename": fn, "active": active}) return items diff --git a/plugin/server/http_server.py b/plugin/server/http_server.py index 8a2df226..0625c5b0 100644 --- a/plugin/server/http_server.py +++ b/plugin/server/http_server.py @@ -1,7 +1,8 @@ import json import threading +import time import urllib.parse -from http.server import BaseHTTPRequestHandler, HTTPServer +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any import binaryninja as bn @@ -14,12 +15,75 @@ from ..utils.string_utils import parse_int_or_default +class MCPThreadingHTTPServer(ThreadingHTTPServer): + daemon_threads = True + allow_reuse_address = True + + class MCPRequestHandler(BaseHTTPRequestHandler): binary_ops = None # Will be set by the server + mcp_server = None + slow_request_seconds = 10.0 def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + def parse_request(self): + parsed = super().parse_request() + if parsed: + self._request_started_at = time.monotonic() + self._request_finished = False + request_path = urllib.parse.urlparse(self.path).path + if self.mcp_server: + self.mcp_server.request_started(id(self), self.command, request_path) + self._slow_request_timer = threading.Timer( + self.slow_request_seconds, + self._log_slow_request, + ) + self._slow_request_timer.daemon = True + self._slow_request_timer.start() + bn.log_debug(f"MCP request started: {self.command} {request_path}") + return parsed + + def handle_one_request(self): + try: + super().handle_one_request() + except Exception as exc: + method = getattr(self, "command", "UNKNOWN") + path = urllib.parse.urlparse(getattr(self, "path", "")).path + bn.log_error(f"MCP request failed: {method} {path}: {exc!r}") + raise + finally: + self._finish_request_tracking() + + def _log_slow_request(self): + if getattr(self, "_request_finished", True): + return + method = getattr(self, "command", "UNKNOWN") + path = urllib.parse.urlparse(getattr(self, "path", "")).path + bn.log_warn( + f"MCP request still running after {self.slow_request_seconds:g}s: {method} {path}" + ) + + def _finish_request_tracking(self): + started_at = getattr(self, "_request_started_at", None) + if started_at is None or getattr(self, "_request_finished", False): + return + self._request_finished = True + timer = getattr(self, "_slow_request_timer", None) + if timer: + timer.cancel() + elapsed = time.monotonic() - started_at + if self.mcp_server: + self.mcp_server.request_finished(id(self)) + method = getattr(self, "command", "UNKNOWN") + path = urllib.parse.urlparse(getattr(self, "path", "")).path + message = f"MCP request finished in {elapsed:.3f}s: {method} {path}" + if elapsed >= self.slow_request_seconds: + bn.log_warn(message) + else: + bn.log_debug(message) + @property def endpoints(self): # Create endpoints on demand to ensure binary_ops is set @@ -238,10 +302,11 @@ def _check_binary_loaded(self): def do_GET(self): try: - # For all endpoints except /status, /convertNumber, /platforms, /binaries, /views, /selectBinary, check loaded + # For all endpoints except server-level utilities, check that a binary is loaded. if ( not ( self.path.startswith("/status") + or self.path.startswith("/health") or self.path.startswith("/convertNumber") or self.path.startswith("/platforms") or self.path.startswith("/binaries") @@ -261,7 +326,11 @@ def do_GET(self): else: limit = parse_int_or_default(params.get("limit"), 100) - if path == "/status": + if path == "/health": + health = self.mcp_server.health() if self.mcp_server else {"status": "unknown"} + self._send_json_response(health) + + elif path == "/status": status = { "loaded": self.binary_ops and self.binary_ops.current_view is not None, "filename": self.binary_ops.current_view.file.filename @@ -2355,32 +2424,89 @@ def __init__(self, config: Config): self.server = None self.thread = None self.binary_ops = BinaryOperations(config.binary_ninja) + self._ready = threading.Event() + self._state_lock = threading.Lock() + self._active_requests: dict[int, dict[str, Any]] = {} + self._serve_error: str | None = None + + @property + def is_running(self) -> bool: + return bool(self.server and self.thread and self.thread.is_alive()) + + def request_started(self, request_id: int, method: str, path: str): + with self._state_lock: + self._active_requests[request_id] = { + "method": method, + "path": path, + "started_at": time.monotonic(), + } + + def request_finished(self, request_id: int): + with self._state_lock: + self._active_requests.pop(request_id, None) + + def health(self) -> dict[str, Any]: + with self._state_lock: + active_requests = len(self._active_requests) + return { + "status": "ok" if self.is_running else "error", + "server_thread_alive": bool(self.thread and self.thread.is_alive()), + "active_requests": active_requests, + "binary_loaded": self.binary_ops.current_view is not None, + "last_server_error": self._serve_error, + } + + def _serve(self): + self._ready.set() + try: + self.server.serve_forever() + except Exception as exc: + self._serve_error = repr(exc) + bn.log_error(f"MCP server thread stopped unexpectedly: {exc!r}") def start(self): """Start the HTTP server in a background thread.""" + if self.is_running: + return + if self.server: + self.server.server_close() + self.server = None + self.thread = None server_address = (self.config.server.host, self.config.server.port) # Create handler with access to binary operations handler_class = type( "MCPRequestHandlerWithOps", (MCPRequestHandler,), - {"binary_ops": self.binary_ops}, + {"binary_ops": self.binary_ops, "mcp_server": self}, ) - self.server = HTTPServer(server_address, handler_class) - self.thread = threading.Thread(target=self.server.serve_forever) + self._ready.clear() + self._serve_error = None + self.server = MCPThreadingHTTPServer(server_address, handler_class) + self.thread = threading.Thread(target=self._serve, name="BinaryNinjaMCPServer") self.thread.daemon = True self.thread.start() + if not self._ready.wait(timeout=2.0) or not self.thread.is_alive(): + self.server.server_close() + self.server = None + self.thread = None + raise RuntimeError("MCP server thread failed to become ready") bn.log_info(f"Server started on {self.config.server.host}:{self.config.server.port}") def stop(self): """Stop the HTTP server and clean up resources.""" if self.server: - self.server.shutdown() + if self.thread and self.thread.is_alive(): + self.server.shutdown() self.server.server_close() if self.thread: - self.thread.join() + self.thread.join(timeout=5.0) + if self.thread.is_alive(): + bn.log_warn("MCP server thread did not stop within 5 seconds") # Clear references so callers can reliably detect stopped state self.thread = None self.server = None + with self._state_lock: + self._active_requests.clear() bn.log_info("Server stopped") diff --git a/tests/test_http_server.py b/tests/test_http_server.py new file mode 100644 index 00000000..a3a98714 --- /dev/null +++ b/tests/test_http_server.py @@ -0,0 +1,133 @@ +import importlib.util +import json +import sys +import threading +import time +import types +import unittest +import urllib.request +from pathlib import Path + + +def _load_http_server_module(): + log_messages = {"debug": [], "info": [], "warn": [], "error": []} + binaryninja = types.ModuleType("binaryninja") + for level in log_messages: + setattr(binaryninja, f"log_{level}", log_messages[level].append) + settings = types.ModuleType("binaryninja.settings") + settings.Settings = object + sys.modules["binaryninja"] = binaryninja + sys.modules["binaryninja.settings"] = settings + + for package in ("plugin", "plugin.server", "plugin.api", "plugin.core", "plugin.utils"): + module = types.ModuleType(package) + module.__path__ = [] + sys.modules[package] = module + + endpoints = types.ModuleType("plugin.api.endpoints") + endpoints.BinaryNinjaEndpoints = object + sys.modules[endpoints.__name__] = endpoints + + binary_operations = types.ModuleType("plugin.core.binary_operations") + + class BinaryOperations: + def __init__(self, _config): + self.current_view = None + + binary_operations.BinaryOperations = BinaryOperations + sys.modules[binary_operations.__name__] = binary_operations + + config = types.ModuleType("plugin.core.config") + config.Config = object + sys.modules[config.__name__] = config + + number_utils = types.ModuleType("plugin.utils.number_utils") + number_utils.convert_number = lambda value, size: (value, size) + sys.modules[number_utils.__name__] = number_utils + + string_utils = types.ModuleType("plugin.utils.string_utils") + string_utils.parse_int_or_default = lambda value, default: ( + int(value) if value is not None else default + ) + sys.modules[string_utils.__name__] = string_utils + + path = Path(__file__).parents[1] / "plugin" / "server" / "http_server.py" + spec = importlib.util.spec_from_file_location("plugin.server.http_server", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module, log_messages + + +class ServerConfig: + server = types.SimpleNamespace(host="127.0.0.1", port=0) + binary_ninja = object() + + +class HTTPServerTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.module, cls.logs = _load_http_server_module() + + def setUp(self): + self.server = self.module.MCPServer(ServerConfig()) + + def tearDown(self): + self.server.stop() + + def url(self, path): + port = self.server.server.server_address[1] + return f"http://127.0.0.1:{port}{path}" + + def test_health_reports_live_server(self): + self.server.start() + with urllib.request.urlopen(self.url("/health"), timeout=1) as response: + health = json.load(response) + + self.assertEqual(health["status"], "ok") + self.assertTrue(health["server_thread_alive"]) + self.assertGreaterEqual(health["active_requests"], 1) + + def test_slow_request_does_not_block_other_requests(self): + slow_started = threading.Event() + release_slow = threading.Event() + original_do_get = self.module.MCPRequestHandler.do_GET + original_threshold = self.module.MCPRequestHandler.slow_request_seconds + + def do_get(handler): + if handler.path == "/slow": + slow_started.set() + release_slow.wait(timeout=2) + handler._send_json_response({"path": handler.path}) + + self.module.MCPRequestHandler.do_GET = do_get + self.module.MCPRequestHandler.slow_request_seconds = 0.05 + try: + self.server.start() + slow_request = threading.Thread( + target=lambda: urllib.request.urlopen(self.url("/slow"), timeout=2).read() + ) + slow_request.start() + self.assertTrue(slow_started.wait(timeout=1)) + + with urllib.request.urlopen(self.url("/fast"), timeout=1) as response: + self.assertEqual(json.load(response), {"path": "/fast"}) + + time.sleep(0.1) + self.assertTrue(any("GET /slow" in message for message in self.logs["warn"])) + + stop_started = time.monotonic() + self.server.stop() + self.assertLess(time.monotonic() - stop_started, 1.0) + + release_slow.set() + slow_request.join(timeout=1) + self.assertFalse(slow_request.is_alive()) + finally: + release_slow.set() + self.module.MCPRequestHandler.do_GET = original_do_get + self.module.MCPRequestHandler.slow_request_seconds = original_threshold + + +if __name__ == "__main__": + unittest.main()