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

matlab_mcp/session/manager.py

@@ -0,0 +1,210 @@
+"""Session manager for MATLAB MCP Server.
+
+Manages the lifecycle of user sessions, each of which owns a temporary
+directory and a set of associated jobs.
+"""
+from __future__ import annotations
+
+import logging
+import shutil
+import threading
+import time
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_SESSION_ID = "default"
+
+
+@dataclass
+class Session:
+    """Represents a single user session.
+
+    Parameters
+    ----------
+    session_id:
+        Unique identifier for this session.
+    temp_dir:
+        Path to the session's temporary working directory.
+    """
+
+    session_id: str
+    temp_dir: str
+    created_at: float = field(default_factory=time.time)
+    last_active: float = field(default_factory=time.time)
+
+    # ------------------------------------------------------------------
+    # Methods
+    # ------------------------------------------------------------------
+
+    def touch(self) -> None:
+        """Update last_active to the current time."""
+        self.last_active = time.time()
+
+    @property
+    def idle_seconds(self) -> float:
+        """Seconds since the session was last active."""
+        return time.time() - self.last_active
+
+
+class SessionManager:
+    """Manages the lifecycle of user sessions.
+
+    Parameters
+    ----------
+    config:
+        The full :class:`~matlab_mcp.config.AppConfig` instance.
+    """
+
+    def __init__(self, config: Any = None, collector: Any = None) -> None:
+        self._config = config
+        self._collector = collector
+        self._sessions: Dict[str, Session] = {}
+
+        # Derive limits from config or use sensible defaults
+        if config is not None:
+            self._max_sessions: int = config.sessions.max_sessions
+            self._session_timeout: int = config.sessions.session_timeout
+            base_temp: str = config.execution.temp_dir
+        else:
+            self._max_sessions = 50
+            self._session_timeout = 3600
+            base_temp = "/tmp/matlab_mcp"
+
+        self._base_temp = Path(base_temp)
+        self._lock = threading.Lock()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def create_session(self, *, session_id: Optional[str] = None) -> Session:
+        """Create a new session with a temporary directory.
+
+        Parameters
+        ----------
+        session_id
+            Optional explicit ID for the session.  When *None* (the default),
+            a random UUID is generated.
+
+        Raises
+        ------
+        RuntimeError
+            If the maximum number of sessions has been reached.
+        """
+        with self._lock:
+            if len(self._sessions) >= self._max_sessions:
+                raise RuntimeError(
+                    f"Maximum number of sessions reached ({self._max_sessions})"
+                )
+
+            session_id = session_id or str(uuid.uuid4())
+            temp_dir = self._base_temp / session_id
+            temp_dir.mkdir(parents=True, exist_ok=True)
+
+            session = Session(session_id=session_id, temp_dir=str(temp_dir))
+            self._sessions[session_id] = session
+            logger.info("Session created: %s (temp_dir=%s, total=%d/%d)",
+                         session_id[:8], temp_dir, len(self._sessions), self._max_sessions)
+            if self._collector:
+                self._collector.record_event("session_created", {"session_id_short": session.session_id[-8:]})
+            return session
+
+    def get_session(self, session_id: str) -> Optional[Session]:
+        """Return the session with the given ID, or None if not found."""
+        with self._lock:
+            return self._sessions.get(session_id)
+
+    def get_or_create_default(self) -> Session:
+        """Return (or create) the default session for single-user stdio mode."""
+        with self._lock:
+            session = self._sessions.get(_DEFAULT_SESSION_ID)
+            if session is not None:
+                return session
+        return self.create_session(session_id=_DEFAULT_SESSION_ID)
+
+    def destroy_session(self, session_id: str) -> bool:
+        """Destroy a session and remove its temporary directory.
+
+        Returns True if the session existed and was destroyed, False otherwise.
+        """
+        with self._lock:
+            session = self._sessions.pop(session_id, None)
+            remaining = len(self._sessions)
+        if session is None:
+            return False
+
+        idle_s = session.idle_seconds
+        logger.info("Destroying session %s (idle=%.0fs, remaining=%d)",
+                     session_id[:8], idle_s, remaining)
+
+        # Clean up temp directory
+        temp_path = Path(session.temp_dir)
+        if temp_path.exists():
+            try:
+                shutil.rmtree(temp_path)
+                logger.info("Removed temp dir %s for session %s", temp_path, session_id[:8])
+            except Exception:
+                logger.warning(
+                    "Failed to remove temp dir %s for session %s",
+                    temp_path,
+                    session_id[:8],
+                )
+        return True
+
+    def cleanup_expired(
+        self,
+        has_active_jobs_fn: Optional[Callable[[str], bool]] = None,
+    ) -> int:
+        """Remove sessions that have been idle beyond the session timeout.
+
+        Sessions with active jobs are skipped.
+
+        Parameters
+        ----------
+        has_active_jobs_fn:
+            A callable that takes a session_id and returns True if the session
+            has active (PENDING or RUNNING) jobs.  If None, all idle sessions
+            are eligible for removal.
+
+        Returns
+        -------
+        int
+            The number of sessions removed.
+        """
+        # Collect idle candidates under lock, then check external callback outside
+        # to avoid holding self._lock while calling into JobTracker.
+        with self._lock:
+            candidates = [
+                sid for sid, s in self._sessions.items()
+                if s.idle_seconds >= self._session_timeout
+            ]
+
+        to_destroy = []
+        for session_id in candidates:
+            if has_active_jobs_fn is not None and has_active_jobs_fn(session_id):
+                logger.debug(
+                    "Skipping cleanup of session %s — has active jobs", session_id
+                )
+                continue
+            to_destroy.append(session_id)
+
+        for session_id in to_destroy:
+            self.destroy_session(session_id)
+
+        if to_destroy:
+            logger.info("Cleaned up %d expired sessions", len(to_destroy))
+        return len(to_destroy)
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def session_count(self) -> int:
+        """Current number of active sessions."""
+        with self._lock:
+            return len(self._sessions)

matlab_mcp/tools/admin.py

@@ -0,0 +1,28 @@
+"""Admin MCP tool implementations.
+
+Provides:
+- get_pool_status_impl — delegate to pool.get_status()
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+async def get_pool_status_impl(pool: Any) -> dict:
+    """Return the current status of the engine pool.
+
+    Parameters
+    ----------
+    pool:
+        An :class:`~matlab_mcp.pool.manager.EnginePoolManager` instance
+        (or any object with a ``get_status()`` method).
+
+    Returns
+    -------
+    dict
+        Status summary with ``total``, ``available``, ``busy``, and ``max`` keys.
+    """
+    return pool.get_status()

matlab_mcp/tools/core.py

@@ -0,0 +1,144 @@
+"""Core MCP tool implementations for MATLAB MCP Server.
+
+Provides the primary execution tools:
+- execute_code_impl  — run MATLAB code (with security check)
+- check_code_impl    — lint MATLAB code via checkcode/mlint
+- get_workspace_impl — retrieve workspace variables via 'whos'
+"""
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any, Optional
+
+from matlab_mcp.security.validator import BlockedFunctionError
+
+logger = logging.getLogger(__name__)
+
+
+async def execute_code_impl(
+    code: str,
+    session_id: str,
+    executor: Any,
+    security: Any,
+    temp_dir: Optional[str] = None,
+) -> dict:
+    """Execute MATLAB code with security check.
+
+    Parameters
+    ----------
+    code:
+        MATLAB source code to execute.
+    session_id:
+        ID of the owning session.
+    executor:
+        A :class:`~matlab_mcp.jobs.executor.JobExecutor` instance.
+    security:
+        A :class:`~matlab_mcp.security.validator.SecurityValidator` instance.
+
+    Returns
+    -------
+    dict
+        Result dict with at minimum ``status`` and ``job_id`` keys.
+        On security violation returns ``{"status": "failed", "error": {...}}``.
+    """
+    # Check security blocklist first
+    try:
+        security.check_code(code)
+    except BlockedFunctionError as exc:
+        return {
+            "status": "failed",
+            "error": {
+                "type": "ValidationError",
+                "message": f"Blocked: {exc}",
+                "matlab_id": None,
+                "stack_trace": None,
+            },
+        }
+
+    # Delegate to executor
+    return await executor.execute(session_id=session_id, code=code, temp_dir=temp_dir)
+
+
+async def check_code_impl(
+    code: str,
+    session_id: str,
+    executor: Any,
+    temp_dir: str,
+) -> dict:
+    """Run checkcode/mlint on MATLAB code.
+
+    Writes *code* to a temporary ``.m`` file, calls ``mcp_checkcode()`` via
+    the executor, parses the JSON result, and cleans up the temp file.
+
+    Parameters
+    ----------
+    code:
+        MATLAB source code to check.
+    session_id:
+        ID of the owning session.
+    executor:
+        A :class:`~matlab_mcp.jobs.executor.JobExecutor` instance.
+    temp_dir:
+        Temporary directory path for writing the ``.m`` file.
+
+    Returns
+    -------
+    dict
+        Parsed result from ``mcp_checkcode`` or an error dict.
+    """
+    # Write code to a temp .m file
+    td = Path(temp_dir)
+    td.mkdir(parents=True, exist_ok=True)
+
+    tmp_file = td / f"_check_{session_id}.m"
+    try:
+        tmp_file.write_text(code, encoding="utf-8")
+
+        # Build the MATLAB call
+        escaped_path = str(tmp_file).replace("\\", "\\\\").replace("'", "''")
+        matlab_cmd = f"mcp_checkcode('{escaped_path}')"
+
+        result = await executor.execute(session_id=session_id, code=matlab_cmd)
+
+        # Try to parse JSON from the output text
+        if result.get("status") == "completed":
+            raw_text = result.get("text", "")
+            try:
+                parsed = json.loads(raw_text)
+                return {"status": "completed", "issues": parsed}
+            except (json.JSONDecodeError, ValueError):
+                # Return raw text if not valid JSON
+                return {"status": "completed", "issues": [], "raw": raw_text}
+
+        return result
+    finally:
+        # Clean up temp file
+        try:
+            if tmp_file.exists():
+                tmp_file.unlink()
+        except Exception:
+            logger.debug("Failed to remove temp file %s", tmp_file)
+
+
+async def get_workspace_impl(
+    session_id: str,
+    executor: Any,
+) -> dict:
+    """Get workspace variables via 'whos'.
+
+    Parameters
+    ----------
+    session_id:
+        ID of the owning session.
+    executor:
+        A :class:`~matlab_mcp.jobs.executor.JobExecutor` instance.
+
+    Returns
+    -------
+    dict
+        Result dict with variables list or raw whos output.
+    """
+    result = await executor.execute(session_id=session_id, code="whos")
+    return result

matlab_mcp/tools/custom.py

@@ -0,0 +1,241 @@
+"""Custom tool support for MATLAB MCP Server.
+
+Provides:
+- CustomToolParam  — pydantic model for a tool parameter definition
+- CustomToolDef    — pydantic model for a complete custom tool definition
+- load_custom_tools  — load custom tools from a YAML config file
+- make_custom_tool_handler  — factory that creates a typed async handler function
+  with a proper ``inspect.Signature`` so FastMCP can introspect it correctly.
+"""
+from __future__ import annotations
+
+import inspect
+import logging
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+import yaml
+from pydantic import BaseModel, Field
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Pydantic models
+# ---------------------------------------------------------------------------
+
+_TYPE_MAP: Dict[str, type] = {
+    "str": str,
+    "string": str,
+    "int": int,
+    "integer": int,
+    "float": float,
+    "number": float,
+    "bool": bool,
+    "boolean": bool,
+    "list": list,
+    "dict": dict,
+    "any": Any,
+}
+
+
+class CustomToolParam(BaseModel):
+    """Definition of a single parameter for a custom tool.
+
+    Parameters
+    ----------
+    name:
+        Parameter name (valid Python identifier).
+    type:
+        Parameter type as a string (e.g. ``"str"``, ``"int"``, ``"float"``).
+    required:
+        Whether the parameter is required.  Defaults to True.
+    default:
+        Default value when ``required`` is False.
+    """
+
+    name: str
+    type: str = "str"
+    required: bool = True
+    default: Optional[Any] = None
+
+
+class CustomToolDef(BaseModel):
+    """Definition of a custom MATLAB-backed tool.
+
+    Parameters
+    ----------
+    name:
+        Tool name (used as the MCP tool name).
+    matlab_function:
+        MATLAB function to call when this tool is invoked.
+    description:
+        Human-readable description shown in MCP tool listings.
+    parameters:
+        List of :class:`CustomToolParam` definitions.
+    returns:
+        Description of the return value.
+    """
+
+    name: str
+    matlab_function: str
+    description: str = ""
+    parameters: List[CustomToolParam] = Field(default_factory=list)
+    returns: str = ""
+
+
+# ---------------------------------------------------------------------------
+# Loader
+# ---------------------------------------------------------------------------
+
+def load_custom_tools(config_path: str) -> List[CustomToolDef]:
+    """Load custom tool definitions from a YAML file.
+
+    Parameters
+    ----------
+    config_path:
+        Path to the YAML configuration file.
+
+    Returns
+    -------
+    list of CustomToolDef
+        Parsed tool definitions, or an empty list if the file does not exist
+        or contains no ``tools`` section.
+    """
+    path = Path(config_path)
+    if not path.exists():
+        logger.debug("Custom tools config not found: %s", path)
+        return []
+
+    try:
+        with open(path, "r", encoding="utf-8") as fh:
+            data = yaml.safe_load(fh) or {}
+    except Exception as exc:
+        logger.error("Failed to load custom tools from %s: %s", path, exc)
+        return []
+
+    raw_tools = data.get("tools", []) or []
+    tools: List[CustomToolDef] = []
+    for raw in raw_tools:
+        try:
+            tools.append(CustomToolDef.model_validate(raw))
+        except Exception as exc:
+            logger.warning("Invalid custom tool definition %r: %s", raw, exc)
+
+    return tools
+
+
+# ---------------------------------------------------------------------------
+# Handler factory
+# ---------------------------------------------------------------------------
+
+def make_custom_tool_handler(
+    tool_def: CustomToolDef,
+    server_state: Any,
+) -> Callable:
+    """Create an async handler function for a custom tool.
+
+    FastMCP introspects the function signature to determine parameter names and
+    types. This factory builds a proper ``inspect.Signature`` with:
+    - ``ctx: Context`` as the first parameter
+    - One parameter per entry in ``tool_def.parameters``
+
+    The returned function delegates to the executor stored in
+    ``server_state.executor``, building a MATLAB function call from
+    ``tool_def.matlab_function`` and the supplied arguments.
+
+    Parameters
+    ----------
+    tool_def:
+        The :class:`CustomToolDef` that describes this tool.
+    server_state:
+        An object with at minimum:
+        - ``executor`` — a :class:`~matlab_mcp.jobs.executor.JobExecutor`
+        - ``session_id`` — default session ID string
+
+    Returns
+    -------
+    callable
+        An async function with ``__name__``, ``__doc__``, and ``__signature__``
+        properly set.
+    """
+    # Try to import Context from fastmcp; fall back to Any for tests without fastmcp
+    try:
+        from fastmcp import Context as _Context
+    except ImportError:  # pragma: no cover
+        _Context = Any  # type: ignore[misc,assignment]
+
+    # Build the inspect.Parameter list
+    params: List[inspect.Parameter] = [
+        inspect.Parameter(
+            "ctx",
+            kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=_Context,
+        )
+    ]
+
+    for p in tool_def.parameters:
+        py_type = _TYPE_MAP.get(p.type.lower(), str)
+        if p.required:
+            params.append(
+                inspect.Parameter(
+                    p.name,
+                    kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                    annotation=py_type,
+                )
+            )
+        else:
+            params.append(
+                inspect.Parameter(
+                    p.name,
+                    kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                    annotation=py_type,
+                    default=p.default,
+                )
+            )
+
+    sig = inspect.Signature(params, return_annotation=dict)
+
+    # Build the actual handler
+    # Capture tool_def and server_state in closure
+    _tool_def = tool_def
+    _server_state = server_state
+
+    async def _handler(*args, **kwargs):
+        # Bind arguments to the signature (skip ctx which is first positional)
+        # args[0] is ctx
+        # remaining args / kwargs are tool params
+        bound = sig.bind(*args, **kwargs)
+        bound.apply_defaults()
+        arguments = dict(bound.arguments)
+        # Remove ctx
+        arguments.pop("ctx", None)
+
+        # Build MATLAB function call: func(arg1, arg2, ...)
+        matlab_args = []
+        for param in _tool_def.parameters:
+            val = arguments.get(param.name)
+            py_type = _TYPE_MAP.get(param.type.lower(), str)
+            if py_type in (str,):
+                # Escape single quotes to prevent MATLAB injection
+                safe_val = str(val).replace("'", "''")
+                matlab_args.append(f"'{safe_val}'")
+            elif py_type in (int, float):
+                matlab_args.append(str(val))
+            elif py_type is bool:
+                matlab_args.append("true" if val else "false")
+            else:
+                matlab_args.append(str(val))
+
+        matlab_call = f"{_tool_def.matlab_function}({', '.join(matlab_args)})"
+
+        session_id = getattr(_server_state, "session_id", "default")
+        executor = _server_state.executor
+
+        return await executor.execute(session_id=session_id, code=matlab_call)
+
+    # Set metadata so FastMCP can use it
+    _handler.__name__ = tool_def.name
+    _handler.__doc__ = tool_def.description or f"Custom tool: {tool_def.name}"
+    _handler.__signature__ = sig
+
+    return _handler