matlab-mcp-python 1.0.0__py3-none-any.whl

matlab_mcp/pool/engine.py

@@ -0,0 +1,216 @@
+"""MATLAB engine wrapper with lifecycle management.
+
+Wraps a single matlab.engine instance with state tracking, health checks,
+and workspace reset capabilities. Uses lazy import so tests can mock the
+matlab.engine module.
+"""
+from __future__ import annotations
+
+import importlib
+import logging
+import time
+from enum import Enum, auto
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class EngineState(Enum):
+    STOPPED = auto()
+    STARTING = auto()
+    IDLE = auto()
+    BUSY = auto()
+
+
+class MatlabEngineWrapper:
+    """Wraps a single MATLAB engine instance.
+
+    Parameters
+    ----------
+    engine_id:
+        Unique identifier for this engine (e.g., ``"engine-0"``).
+    pool_config:
+        ``PoolConfig`` instance (used for matlab_root if needed).
+    workspace_config:
+        ``WorkspaceConfig`` instance with default_paths and startup_commands.
+    """
+
+    def __init__(self, engine_id: str, pool_config: Any, workspace_config: Any) -> None:
+        self.engine_id = engine_id
+        self._pool_config = pool_config
+        self._workspace_config = workspace_config
+
+        self._engine: Any = None
+        self._state: EngineState = EngineState.STOPPED
+        self._idle_since: float = time.monotonic()
+
+    # ------------------------------------------------------------------
+    # State properties
+    # ------------------------------------------------------------------
+
+    @property
+    def state(self) -> EngineState:
+        return self._state
+
+    @property
+    def idle_seconds(self) -> float:
+        """Seconds since the engine last became idle (0 if not idle)."""
+        if self._state == EngineState.IDLE:
+            return time.monotonic() - self._idle_since
+        return 0.0
+
+    @property
+    def is_alive(self) -> bool:
+        """True if the underlying engine object exists and reports alive."""
+        if self._engine is None:
+            return False
+        alive_attr = getattr(self._engine, "is_alive", None)
+        if alive_attr is None:
+            # Real matlab.engine doesn't have is_alive — assume alive
+            return True
+        if callable(alive_attr):
+            return bool(alive_attr())
+        return bool(alive_attr)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def _get_matlab_engine_module(self) -> Any:
+        """Lazily import matlab.engine so tests can mock it."""
+        return importlib.import_module("matlab.engine")
+
+    def start(self) -> None:
+        """Start the MATLAB engine and apply default paths and startup commands."""
+        logger.info("[%s] Starting MATLAB engine", self.engine_id)
+        self._state = EngineState.STARTING
+
+        matlab_engine = self._get_matlab_engine_module()
+        self._engine = matlab_engine.start_matlab()
+
+        # Apply default workspace paths
+        for path in self._workspace_config.default_paths:
+            try:
+                self._engine.addpath(path)
+            except Exception:
+                logger.warning("[%s] Failed to addpath: %s", self.engine_id, path)
+
+        # Run startup commands
+        for cmd in self._workspace_config.startup_commands:
+            try:
+                self._engine.eval(cmd, nargout=0)
+            except Exception:
+                logger.warning("[%s] Startup command failed: %s", self.engine_id, cmd)
+
+        self._state = EngineState.IDLE
+        self._idle_since = time.monotonic()
+        logger.info("[%s] MATLAB engine started", self.engine_id)
+
+    def stop(self) -> None:
+        """Quit the MATLAB engine."""
+        if self._engine is not None:
+            try:
+                self._engine.quit()
+            except Exception:
+                logger.warning("[%s] Exception during engine quit", self.engine_id)
+            finally:
+                self._engine = None
+        self._state = EngineState.STOPPED
+        logger.info("[%s] MATLAB engine stopped", self.engine_id)
+
+    # ------------------------------------------------------------------
+    # Operations
+    # ------------------------------------------------------------------
+
+    def health_check(self) -> bool:
+        """Run a trivial eval to confirm the engine is responsive."""
+        if self._engine is None:
+            return False
+        try:
+            self._engine.eval("1", nargout=0)
+            return True
+        except Exception:
+            return False
+
+    def execute(self, code: str, nargout: int = 0, background: bool = False,
+                stdout: Any = None, stderr: Any = None) -> Any:
+        """Run MATLAB code on the engine.
+
+        Parameters
+        ----------
+        code:       MATLAB code to evaluate.
+        nargout:    Number of output arguments expected.
+        background: If True, return a future-like object immediately.
+        stdout:     Stream to capture standard output (e.g. ``io.StringIO``).
+        stderr:     Stream to capture standard error (e.g. ``io.StringIO``).
+        """
+        if self._engine is None:
+            raise RuntimeError(f"[{self.engine_id}] Engine is not started")
+        kwargs: dict[str, Any] = {"nargout": nargout, "background": background}
+        if stdout is not None:
+            kwargs["stdout"] = stdout
+        if stderr is not None:
+            kwargs["stderr"] = stderr
+        return self._engine.eval(code, **kwargs)
+
+    def reset_workspace(self) -> None:
+        """Reset the MATLAB workspace to a clean state.
+
+        Sequence: clear all, clear global, clear functions, fclose all,
+        restoredefaultpath, re-add configured paths, re-run startup commands.
+        """
+        if self._engine is None:
+            raise RuntimeError(f"[{self.engine_id}] Engine is not started")
+
+        cleanup_commands = [
+            "clear all",
+            "clear global",
+            "clear functions",
+            "fclose all",
+        ]
+        for cmd in cleanup_commands:
+            try:
+                self._engine.eval(cmd, nargout=0)
+            except Exception:
+                logger.warning("[%s] Reset command failed: %s", self.engine_id, cmd)
+
+        try:
+            self._engine.restoredefaultpath()
+        except Exception:
+            logger.warning("[%s] restoredefaultpath failed", self.engine_id)
+
+        # Re-apply configured paths
+        for path in self._workspace_config.default_paths:
+            try:
+                self._engine.addpath(path)
+            except Exception:
+                logger.warning("[%s] Re-addpath failed: %s", self.engine_id, path)
+
+        # Re-run startup commands
+        for cmd in self._workspace_config.startup_commands:
+            try:
+                self._engine.eval(cmd, nargout=0)
+            except Exception:
+                logger.warning("[%s] Startup command failed on reset: %s", self.engine_id, cmd)
+
+        logger.debug("[%s] Workspace reset complete", self.engine_id)
+
+    # ------------------------------------------------------------------
+    # State transitions
+    # ------------------------------------------------------------------
+
+    def mark_busy(self) -> None:
+        """Transition engine to BUSY state."""
+        self._state = EngineState.BUSY
+
+    def mark_idle(self) -> None:
+        """Transition engine to IDLE state and record timestamp."""
+        self._state = EngineState.IDLE
+        self._idle_since = time.monotonic()
+
+    # ------------------------------------------------------------------
+    # Dunder
+    # ------------------------------------------------------------------
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"MatlabEngineWrapper(id={self.engine_id!r}, state={self._state.name})"

matlab_mcp/pool/manager.py

@@ -0,0 +1,227 @@
+"""MATLAB engine pool manager.
+
+Manages a pool of ``MatlabEngineWrapper`` instances, handling:
+- Startup of a minimum number of engines in parallel
+- Acquire/release with queueing when all engines are busy
+- Scale-up on demand up to max_engines
+- Scale-down of idle engines beyond min_engines
+- Periodic health checks that replace dead engines
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Dict, List
+
+from matlab_mcp.pool.engine import EngineState, MatlabEngineWrapper
+
+logger = logging.getLogger(__name__)
+
+
+class EnginePoolManager:
+    """Manages a pool of MATLAB engine wrappers.
+
+    Parameters
+    ----------
+    config:
+        ``AppConfig`` instance.  Uses ``config.pool`` and ``config.workspace``.
+    """
+
+    def __init__(self, config: Any, collector: Any = None) -> None:
+        self._config = config
+        self._pool_config = config.pool
+        self._workspace_config = config.workspace
+        self._collector = collector
+
+        # All engines (available + busy)
+        self._all_engines: List[MatlabEngineWrapper] = []
+        # Queue of engines available for acquisition
+        self._available: asyncio.Queue = asyncio.Queue()
+        # Lock to guard scale-up logic
+        self._scale_lock: asyncio.Lock = asyncio.Lock()
+        self._next_id: int = 0
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _make_engine(self) -> MatlabEngineWrapper:
+        engine_id = f"engine-{self._next_id}"
+        self._next_id += 1
+        return MatlabEngineWrapper(engine_id, self._pool_config, self._workspace_config)
+
+    async def _start_engine_async(self) -> MatlabEngineWrapper:
+        """Start a single engine in a thread executor and return it."""
+        loop = asyncio.get_running_loop()
+        engine = self._make_engine()
+        self._all_engines.append(engine)
+        await loop.run_in_executor(None, engine.start)
+        return engine
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    async def start(self) -> None:
+        """Start min_engines engines in parallel."""
+        min_engines = self._pool_config.min_engines
+        logger.info("Starting pool with %d engine(s)", min_engines)
+
+        tasks = [self._start_engine_async() for _ in range(min_engines)]
+        engines = await asyncio.gather(*tasks)
+
+        for engine in engines:
+            await self._available.put(engine)
+
+        logger.info("Pool started: %d engine(s) ready", min_engines)
+
+    async def acquire(self) -> MatlabEngineWrapper:
+        """Return an available engine, scaling up if needed.
+
+        Blocks if the pool is at max capacity and all engines are busy.
+        """
+        # Try to get an immediately available engine
+        try:
+            engine = self._available.get_nowait()
+            engine.mark_busy()
+            logger.info("Acquired engine %s (available=%d, busy=%d)",
+                        engine.engine_id, self._available.qsize(),
+                        len(self._all_engines) - self._available.qsize())
+            return engine
+        except asyncio.QueueEmpty:
+            pass
+
+        # Attempt scale-up under a lock to avoid races
+        async with self._scale_lock:
+            total = len(self._all_engines)
+            if total < self._pool_config.max_engines:
+                logger.info("Scaling up pool: starting new engine (%d/%d)",
+                            total + 1, self._pool_config.max_engines)
+                engine = await self._start_engine_async()
+                engine.mark_busy()
+                if self._collector:
+                    self._collector.record_event("engine_scale_up", {"engine_id": engine.engine_id, "total_after": len(self._all_engines)})
+                return engine
+
+        # At max capacity — wait for one to become available
+        logger.warning("Pool at max capacity (%d/%d busy) — waiting for available engine",
+                        len(self._all_engines), self._pool_config.max_engines)
+        engine = await self._available.get()
+        engine.mark_busy()
+        logger.info("Acquired engine %s after wait (pool was full)", engine.engine_id)
+        return engine
+
+    async def release(self, engine: MatlabEngineWrapper) -> None:
+        """Return an engine to the pool."""
+        logger.info("Releasing engine %s — resetting workspace", engine.engine_id)
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, engine.reset_workspace)
+        engine.mark_idle()
+        await self._available.put(engine)
+        logger.info("Engine %s returned to pool (available=%d)",
+                     engine.engine_id, self._available.qsize())
+
+    async def stop(self) -> None:
+        """Stop all engines in the pool."""
+        logger.info("Stopping pool (%d engines)", len(self._all_engines))
+        loop = asyncio.get_running_loop()
+        stop_tasks = [
+            loop.run_in_executor(None, engine.stop)
+            for engine in self._all_engines
+        ]
+        await asyncio.gather(*stop_tasks, return_exceptions=True)
+        self._all_engines.clear()
+        # Drain the queue
+        while not self._available.empty():
+            try:
+                self._available.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        logger.info("Pool stopped")
+
+    async def run_health_checks(self) -> None:
+        """Check idle engines, replace dead ones, scale down excess idle engines.
+
+        - Engines that fail health check are stopped and replaced.
+        - Idle engines beyond min_engines are stopped and removed.
+        """
+        loop = asyncio.get_running_loop()
+        min_engines = self._pool_config.min_engines
+        idle_timeout = self._pool_config.scale_down_idle_timeout
+
+        # Collect currently available (idle) engines by draining the queue
+        idle_engines: List[MatlabEngineWrapper] = []
+        while not self._available.empty():
+            try:
+                idle_engines.append(self._available.get_nowait())
+            except asyncio.QueueEmpty:
+                break
+
+        to_replace: List[MatlabEngineWrapper] = []
+        to_remove: List[MatlabEngineWrapper] = []
+        to_keep: List[MatlabEngineWrapper] = []
+
+        # How many busy engines do we have?
+        busy_count = sum(1 for e in self._all_engines if e.state == EngineState.BUSY)
+        idle_count = len(idle_engines)
+        total = busy_count + idle_count
+
+        for engine in idle_engines:
+            # Check health
+            healthy = await loop.run_in_executor(None, engine.health_check)
+            if not healthy:
+                logger.warning("[%s] Health check failed; replacing", engine.engine_id)
+                if self._collector:
+                    self._collector.record_event("health_check_fail", {"engine_id": engine.engine_id, "error": "health check failed"})
+                to_replace.append(engine)
+                continue
+
+            # Scale down if idle beyond timeout and above min
+            if (engine.idle_seconds > idle_timeout
+                    and total > min_engines
+                    and len(to_keep) + busy_count >= min_engines):
+                logger.info("[%s] Scaling down idle engine (idle %.0fs)",
+                            engine.engine_id, engine.idle_seconds)
+                if self._collector:
+                    self._collector.record_event("engine_scale_down", {"engine_id": engine.engine_id, "total_after": total - 1})
+                to_remove.append(engine)
+                total -= 1
+            else:
+                to_keep.append(engine)
+
+        # Stop engines to remove / replace
+        engines_to_stop = to_replace + to_remove
+        stop_tasks = [loop.run_in_executor(None, e.stop) for e in engines_to_stop]
+        await asyncio.gather(*stop_tasks, return_exceptions=True)
+        for e in engines_to_stop:
+            try:
+                self._all_engines.remove(e)
+            except ValueError:
+                pass
+
+        # Start replacement engines for failed ones
+        replacement_tasks = [self._start_engine_async() for _ in to_replace]
+        new_engines = await asyncio.gather(*replacement_tasks, return_exceptions=True)
+
+        # Return healthy engines + replacements to queue
+        for engine in to_keep:
+            await self._available.put(engine)
+        for old_engine, new_engine in zip(to_replace, new_engines):
+            if isinstance(new_engine, Exception):
+                logger.error("Replacement engine failed to start: %s", new_engine)
+            else:
+                if self._collector:
+                    self._collector.record_event("engine_replaced", {"old_id": old_engine.engine_id, "new_id": new_engine.engine_id})
+                await self._available.put(new_engine)
+
+    def get_status(self) -> Dict[str, int]:
+        """Return pool status summary."""
+        total = len(self._all_engines)
+        available = self._available.qsize()
+        busy = total - available
+        return {
+            "total": total,
+            "available": available,
+            "busy": max(busy, 0),
+            "max": self._pool_config.max_engines,
+        }

matlab_mcp/security/validator.py

@@ -0,0 +1,154 @@
+"""Security validator for MATLAB code and filenames.
+
+Provides:
+- ``BlockedFunctionError`` — raised when blocked MATLAB code is detected.
+- ``SecurityValidator``   — checks code and sanitizes filenames.
+"""
+from __future__ import annotations
+
+import logging
+import re
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class BlockedFunctionError(Exception):
+    """Raised when MATLAB code contains a blocked function or construct."""
+
+
+# Precompiled patterns for _strip_string_literals
+_DOUBLE_QUOTED_RE = re.compile(r'"[^"\n]*"')
+_SINGLE_QUOTED_RE = re.compile(r"(?<![a-zA-Z0-9_\)\]])'[^'\n]*'")
+_COMMENT_RE = re.compile(r'%')
+
+
+class SecurityValidator:
+    """Validates MATLAB code and filenames against a security policy.
+
+    Parameters
+    ----------
+    security_config:
+        ``SecurityConfig`` instance with ``blocked_functions_enabled`` and
+        ``blocked_functions`` attributes.
+    """
+
+    def __init__(self, security_config: Any, collector: Any = None) -> None:
+        self._config = security_config
+        self._collector = collector
+
+        # Precompile blocked-function patterns for check_code hot path
+        self._call_patterns: dict[str, re.Pattern] = {}
+        self._cmd_patterns: dict[str, re.Pattern] = {}
+        for func in security_config.blocked_functions:
+            if func != "!":
+                self._call_patterns[func] = re.compile(rf"\b{re.escape(func)}\s*\(")
+                self._cmd_patterns[func] = re.compile(rf"(?:^|;)\s*{re.escape(func)}\s+\S", re.MULTILINE)
+
+    # ------------------------------------------------------------------
+    # Code checking
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _strip_string_literals(code: str) -> str:
+        """Remove MATLAB string literals and comments to avoid false positives.
+
+        Processing order per line:
+        1. Remove double-quoted strings "..."
+        2. Remove single-quoted strings '...' (MATLAB char arrays)
+           - A quote preceded by [a-zA-Z0-9_)] is a transpose operator, not a string.
+        3. Remove MATLAB comments (% to end of line)
+
+        Note: This is a best-effort heuristic; it handles the common cases
+        tested without a full MATLAB parser.
+        """
+        processed_lines = []
+        for line in code.splitlines():
+            line = _DOUBLE_QUOTED_RE.sub('""', line)
+            line = _SINGLE_QUOTED_RE.sub("''", line)
+            comment_match = _COMMENT_RE.search(line)
+            if comment_match:
+                line = line[:comment_match.start()]
+            processed_lines.append(line)
+        return "\n".join(processed_lines)
+
+    def check_code(self, code: str) -> None:
+        """Scan *code* for blocked functions/constructs.
+
+        Parameters
+        ----------
+        code:
+            MATLAB source code to check.
+
+        Raises
+        ------
+        BlockedFunctionError
+            If a blocked construct is found (and blocking is enabled).
+        """
+        if not self._config.blocked_functions_enabled:
+            logger.debug("Security check skipped (blocked_functions disabled)")
+            return
+
+        # Strip string literals to avoid false positives
+        sanitized = self._strip_string_literals(code)
+
+        for func in self._config.blocked_functions:
+            if func == "!":
+                # Shell escape: lines starting with ! (after optional whitespace)
+                for line in sanitized.splitlines():
+                    if line.lstrip().startswith("!"):
+                        logger.warning("BLOCKED: shell escape '!' in code: %s", repr(code[:120]))
+                        if self._collector:
+                            self._collector.record_event("blocked_function", {"function": "!"})
+                        raise BlockedFunctionError(
+                            "Shell escape '!' is not allowed"
+                        )
+            else:
+                call_re = self._call_patterns[func]
+                cmd_re = self._cmd_patterns[func]
+                if call_re.search(sanitized) or cmd_re.search(sanitized):
+                    logger.warning("BLOCKED: function '%s' in code: %s", func, repr(code[:120]))
+                    if self._collector:
+                        self._collector.record_event("blocked_function", {"function": func})
+                    raise BlockedFunctionError(
+                        f"Function '{func}' is not allowed"
+                    )
+
+    # ------------------------------------------------------------------
+    # Filename sanitization
+    # ------------------------------------------------------------------
+
+    _SAFE_FILENAME_RE = re.compile(r'^[a-zA-Z0-9._-]+$')
+
+    def sanitize_filename(self, filename: str) -> str:
+        """Validate and return *filename* if safe.
+
+        Parameters
+        ----------
+        filename:
+            Proposed filename (basename only, no directory separators).
+
+        Returns
+        -------
+        str
+            The filename unchanged if it passes validation.
+
+        Raises
+        ------
+        ValueError
+            If the filename is empty, contains path traversal (``..``),
+            or contains characters outside ``[a-zA-Z0-9._-]``.
+        """
+        if not filename:
+            raise ValueError("Filename must not be empty")
+
+        if ".." in filename:
+            raise ValueError(f"Path traversal not allowed in filename: {filename!r}")
+
+        if not self._SAFE_FILENAME_RE.match(filename):
+            raise ValueError(
+                f"Filename contains invalid characters: {filename!r}. "
+                "Only [a-zA-Z0-9._-] are allowed."
+            )
+
+        return filename