baltamatica-mcp 0.2.0__py3-none-any.whl

baltamatica_mcp/__init__.py

@@ -0,0 +1,3 @@
+"""baltamatica-mcp: MCP server for Baltamatica scientific computing engine."""
+
+__version__ = "0.1.0"

baltamatica_mcp/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for running baltamatica-mcp as a module: python -m baltamatica_mcp"""
+
+from baltamatica_mcp.server import main
+
+if __name__ == "__main__":
+    main()

baltamatica_mcp/backend_bex.py

@@ -0,0 +1,316 @@
+"""BEX backend client for the Baltamatica JSON-over-TCP bridge."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import json
+from dataclasses import replace
+from typing import Any
+
+from baltamatica_mcp.engine import (
+    Artifact,
+    EngineError,
+    EngineUnavailableError,
+    ExecutionResult,
+    VariableInfo,
+    VariableListResult,
+)
+from baltamatica_mcp.serializer import (
+    encode_for_set,
+    present_binary_value,
+    present_structured,
+)
+
+DEFAULT_BEX_HOST = "127.0.0.1"
+DEFAULT_BEX_PORT = 31415
+DEFAULT_ARTIFACT_TYPE = "application/octet-stream"
+# Max bytes for a single response line (large binary get_variable payloads).
+_STREAM_LIMIT = 64 * 1024 * 1024
+
+
+class BexProtocolError(EngineError):
+    """Raised when the BEX bridge returns malformed protocol data."""
+
+
+class _BexConnectionLost(EngineUnavailableError):
+    """Internal signal used to retry a dropped transport once."""
+
+
+class BexEngine:
+    """Client for a Baltamatica BEX plugin speaking newline-delimited JSON."""
+
+    backend = "bex"
+
+    def __init__(
+        self,
+        host: str = DEFAULT_BEX_HOST,
+        port: int = DEFAULT_BEX_PORT,
+        timeout: float = 30.0,
+        reconnect: bool = True,
+    ) -> None:
+        if port <= 0 or port > 65535:
+            raise ValueError(f"Invalid BEX port: {port}")
+        if timeout <= 0:
+            raise ValueError(f"Invalid BEX timeout: {timeout}")
+
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+        self.reconnect = reconnect
+        self._reader: asyncio.StreamReader | None = None
+        self._writer: asyncio.StreamWriter | None = None
+        self._request_counter = 0
+
+    async def execute_code(self, code: str) -> ExecutionResult:
+        response = await self._request("execute_code", {"code": code})
+        return _with_output_artifacts(_execution_result_from_response(response))
+
+    async def run_script(self, file_path: str) -> ExecutionResult:
+        response = await self._request("run_script", {"file_path": file_path})
+        return _with_output_artifacts(_execution_result_from_response(response))
+
+    async def clear_workspace(self) -> ExecutionResult:
+        response = await self._request("clear_workspace", {})
+        return _execution_result_from_response(response)
+
+    async def set_variable(self, name: str, data: Any, dtype: str | None = None) -> ExecutionResult:
+        resolved_dtype, dims, raw = encode_for_set(data, dtype)
+        response = await self._request(
+            "set_variable",
+            {
+                "name": name,
+                "dtype": resolved_dtype,
+                "dims": dims,
+                "data_b64": base64.b64encode(raw).decode("ascii"),
+            },
+        )
+        return _execution_result_from_response(response)
+
+    async def list_variables(self) -> VariableListResult:
+        response = await self._request("list_variables", {})
+        return VariableListResult(
+            success=bool(response.get("success", False)),
+            variables=_variables_from_response(response),
+            output=_string_field(response, "output"),
+            error=_error_from_response(response),
+        )
+
+    async def get_variable(self, name: str) -> ExecutionResult:
+        response = await self._request("get_variable", {"name": name})
+        result = _execution_result_from_response(response)
+        value = result.value
+        if isinstance(value, dict) and "data_b64" in value:
+            presented, artifacts = present_binary_value(value, name=name)
+            result = replace(
+                result,
+                value=presented,
+                artifacts=(result.artifacts or []) + artifacts,
+            )
+        elif isinstance(value, dict) and value.get("type") in {
+            "numeric_array",
+            "logical_array",
+            "struct",
+            "cell",
+        }:
+            result = replace(result, value=present_structured(value))
+        return result
+
+    async def close(self) -> None:
+        """Close the current TCP connection, if one is open."""
+
+        writer = self._writer
+        self._reader = None
+        self._writer = None
+        if writer is None:
+            return
+        writer.close()
+        await writer.wait_closed()
+
+    async def _request(self, method: str, params: dict[str, object]) -> dict[str, Any]:
+        payload = {
+            "id": self._next_request_id(),
+            "method": method,
+            "params": params,
+        }
+
+        try:
+            return await self._send_once(payload)
+        except _BexConnectionLost:
+            if not self.reconnect:
+                raise
+            return await self._send_once(payload)
+
+    async def _send_once(self, payload: dict[str, object]) -> dict[str, Any]:
+        reader, writer = await self._ensure_connection()
+        data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + b"\n"
+
+        try:
+            writer.write(data)
+            await asyncio.wait_for(writer.drain(), timeout=self.timeout)
+            line = await asyncio.wait_for(reader.readline(), timeout=self.timeout)
+        except asyncio.TimeoutError as exc:
+            await self.close()
+            raise TimeoutError(
+                f"BEX request '{payload['method']}' timed out after {self.timeout:g} seconds."
+            ) from exc
+        except (ConnectionError, OSError) as exc:
+            await self.close()
+            raise _BexConnectionLost(f"BEX connection lost: {exc}") from exc
+
+        if not line:
+            await self.close()
+            raise _BexConnectionLost("BEX connection closed before a response was received.")
+
+        try:
+            response = json.loads(line.decode("utf-8"))
+        except json.JSONDecodeError as exc:
+            raise BexProtocolError(f"BEX response is not valid JSON: {exc}") from exc
+
+        if not isinstance(response, dict):
+            raise BexProtocolError("BEX response must be a JSON object.")
+        if response.get("id") != payload["id"]:
+            raise BexProtocolError(
+                f"BEX response id mismatch: expected {payload['id']!r}, "
+                f"got {response.get('id')!r}."
+            )
+        return response
+
+    async def _ensure_connection(self) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
+        if self._connection_is_usable():
+            return self._reader, self._writer  # type: ignore[return-value]
+
+        await self.close()
+        try:
+            self._reader, self._writer = await asyncio.wait_for(
+                # A large stream limit so big get_variable responses (base64 of a
+                # large matrix) fit in a single readline() instead of raising
+                # LimitOverrunError at the 64 KB default.
+                asyncio.open_connection(self.host, self.port, limit=_STREAM_LIMIT),
+                timeout=self.timeout,
+            )
+        except asyncio.TimeoutError as exc:
+            raise EngineUnavailableError(
+                f"Timed out connecting to BEX bridge at {self.host}:{self.port}."
+            ) from exc
+        except OSError as exc:
+            raise EngineUnavailableError(
+                f"Failed to connect to BEX bridge at {self.host}:{self.port}: {exc}"
+            ) from exc
+        return self._reader, self._writer
+
+    def _connection_is_usable(self) -> bool:
+        return (
+            self._reader is not None
+            and self._writer is not None
+            and not self._writer.is_closing()
+            and not self._reader.at_eof()
+        )
+
+    def _next_request_id(self) -> str:
+        self._request_counter += 1
+        return str(self._request_counter)
+
+
+def _with_output_artifacts(result: ExecutionResult) -> ExecutionResult:
+    """Merge BALTAMATICA_ARTIFACT markers found in captured output into artifacts.
+
+    Baltamatica cannot export figures to files, but data-export functions
+    (writematrix, writetable, save, ...) can, and a script/command can announce
+    the resulting file with a BALTAMATICA_ARTIFACT line. Since execute_code and
+    run_script now capture console output, those markers are reported here too.
+    """
+
+    from baltamatica_mcp.backend_cli import parse_artifacts
+
+    if not result.output:
+        return result
+    marker_artifacts = parse_artifacts(result.output)
+    if not marker_artifacts:
+        return result
+    return replace(result, artifacts=(result.artifacts or []) + marker_artifacts)
+
+
+def _execution_result_from_response(response: dict[str, Any]) -> ExecutionResult:
+    return ExecutionResult(
+        success=bool(response.get("success", False)),
+        output=_string_field(response, "output"),
+        error=_error_from_response(response),
+        artifacts=_artifacts_from_response(response),
+        value=_value_from_response(response),
+    )
+
+
+def _error_from_response(response: dict[str, Any]) -> str | None:
+    error = response.get("error")
+    if error is None:
+        return None if response.get("success", False) else "BEX request failed."
+    if isinstance(error, dict):
+        code = str(error.get("code") or "ERROR")
+        message = str(error.get("message") or "")
+        return f"{code}: {message}" if message else code
+    return str(error)
+
+
+def _artifacts_from_response(response: dict[str, Any]) -> list[Artifact]:
+    raw_artifacts = response.get("artifacts") or []
+    if not isinstance(raw_artifacts, list):
+        return []
+
+    artifacts: list[Artifact] = []
+    for item in raw_artifacts:
+        if not isinstance(item, dict):
+            continue
+        path = str(item.get("path") or "")
+        if not path:
+            continue
+        artifacts.append(
+            Artifact(
+                path=path,
+                type=str(item.get("type") or DEFAULT_ARTIFACT_TYPE),
+                exists=bool(item.get("exists", False)),
+                size=_int_field(item.get("size")),
+            )
+        )
+    return artifacts
+
+
+def _value_from_response(response: dict[str, Any]) -> dict[str, Any] | None:
+    value = response.get("value")
+    return value if isinstance(value, dict) else None
+
+
+def _variables_from_response(response: dict[str, Any]) -> list[VariableInfo]:
+    raw_variables = response.get("variables") or []
+    if not isinstance(raw_variables, list):
+        return []
+
+    variables: list[VariableInfo] = []
+    for item in raw_variables:
+        if not isinstance(item, dict):
+            continue
+        name = str(item.get("name") or "")
+        if not name:
+            continue
+        variables.append(
+            VariableInfo(
+                name=name,
+                size=str(item.get("size") or ""),
+                bytes=_int_field(item.get("bytes")),
+                class_name=str(item.get("class_name") or ""),
+                attributes=str(item.get("attributes") or ""),
+            )
+        )
+    return variables
+
+
+def _string_field(data: dict[str, Any], key: str) -> str:
+    value = data.get(key)
+    return "" if value is None else str(value)
+
+
+def _int_field(value: object) -> int:
+    try:
+        return int(value) if value is not None else 0
+    except (TypeError, ValueError):
+        return 0

baltamatica_mcp/backend_cli.py

@@ -0,0 +1,323 @@
+"""CLI backend for running Baltamatica commands through baltamaticaC.sh."""
+
+from __future__ import annotations
+
+import asyncio
+import locale
+import os
+import re
+import shutil
+import subprocess
+import tempfile
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+from baltamatica_mcp.engine import (
+    Artifact,
+    ExecutionResult,
+    EngineUnavailableError,
+    VariableInfo,
+    VariableListResult,
+)
+
+DEFAULT_EXECUTABLE = "baltamaticaC.sh"
+ENV_EXECUTABLE = "BALTAMATICA_CLI"
+ARTIFACT_PREFIX = "BALTAMATICA_ARTIFACT="
+VAR_NAME_PATTERN = re.compile(r"^[A-Za-z]\w*$")
+ANSI_PATTERN = re.compile(r"\x1b\[[0-9;]*m")
+NOISE_OUTPUT_PATTERNS = (
+    "Failed to delete old log:",
+)
+ERROR_OUTPUT_PATTERNS = (
+    "错误使用函数",
+    "未定义的变量或函数",
+    "未定义的函数",
+    "是未定义的变量或函数",
+    "位于输入的第",
+    "位于文件",
+)
+WHOS_ROW_PATTERN = re.compile(
+    r"^\s*(?P<name>[A-Za-z]\w*)\s+"
+    r"(?P<size>\S+)\s+"
+    r"(?P<bytes>\d+)\s+"
+    r"(?P<class_name>\S+)"
+    r"(?:\s+(?P<attributes>.*?))?\s*$"
+)
+
+
+@dataclass(frozen=True)
+class CliProcessResult:
+    """Raw process result from the Baltamatica CLI."""
+
+    returncode: int
+    stdout: str
+    stderr: str
+
+
+class CliEngine:
+    """Execute Baltamatica code by spawning the command-line runtime."""
+
+    backend = "cli"
+
+    def __init__(
+        self,
+        executable: str | None = None,
+        timeout: float = 30.0,
+        state_file: str | Path | None = None,
+    ) -> None:
+        self._configured_executable = executable
+        self.timeout = timeout
+        self.state_file = Path(state_file) if state_file else _default_state_file()
+
+    async def execute_code(self, code: str) -> ExecutionResult:
+        return await self._execute_command(self._wrap_stateful_code(code))
+
+    async def run_script(self, file_path: str) -> ExecutionResult:
+        script_command = f"run('{_escape_baltamatica_string(Path(file_path).as_posix())}')"
+        return await self.execute_code(script_command)
+
+    async def clear_workspace(self) -> ExecutionResult:
+        if self.state_file.exists():
+            self.state_file.unlink()
+        return await self._execute_command("clear")
+
+    async def set_variable(self, name: str, data: object, dtype: str | None = None) -> ExecutionResult:
+        from baltamatica_mcp.serializer import to_baltamatica_literal
+
+        return await self.execute_code(f"{name} = {to_baltamatica_literal(data, dtype)};")
+
+    async def list_variables(self) -> VariableListResult:
+        result = await self._execute_command(self._wrap_readonly_code("whos"))
+        if not result.success:
+            return VariableListResult(
+                success=False,
+                variables=[],
+                output=result.output,
+                error=result.error,
+            )
+        return VariableListResult(
+            success=True,
+            variables=parse_whos_output(result.output),
+            output=result.output,
+        )
+
+    async def get_variable(self, name: str) -> ExecutionResult:
+        if not VAR_NAME_PATTERN.match(name):
+            raise ValueError(f"Invalid variable name: {name}")
+        return await self._execute_command(self._wrap_readonly_code(f"disp({name})"))
+
+    def _resolve_executable(self) -> str:
+        configured = self._configured_executable or os.environ.get(ENV_EXECUTABLE)
+        if configured:
+            path = Path(configured).expanduser()
+            if path.is_absolute() or len(path.parts) > 1:
+                if not path.exists():
+                    raise EngineUnavailableError(
+                        f"Baltamatica CLI executable does not exist: {path}"
+                    )
+                if not path.is_file():
+                    raise EngineUnavailableError(f"Baltamatica CLI path is not a file: {path}")
+                return str(path)
+            resolved = shutil.which(configured)
+            if resolved:
+                return resolved
+            raise EngineUnavailableError(
+                f"Baltamatica CLI executable not found on PATH: {configured}"
+            )
+
+        resolved = shutil.which(DEFAULT_EXECUTABLE)
+        if resolved:
+            return resolved
+        raise EngineUnavailableError(
+            f"Baltamatica CLI executable not found. Set {ENV_EXECUTABLE} or pass --cli-executable."
+        )
+
+    async def _run_cli(self, argv: Sequence[str]) -> CliProcessResult:
+        return await asyncio.to_thread(self._run_cli_sync, argv)
+
+    async def _execute_command(self, code: str) -> ExecutionResult:
+        process = await self._run_cli([self._resolve_executable(), "-nodesktop", "-s", code])
+        output = _combine_output(process.stdout, process.stderr)
+        artifacts = parse_artifacts(output)
+        if process.returncode != 0:
+            return ExecutionResult(
+                success=False,
+                output=output,
+                error=_process_error(process),
+                artifacts=artifacts,
+            )
+        output_error = detect_baltamatica_error(output)
+        if output_error:
+            return ExecutionResult(
+                success=False,
+                output=output,
+                error=output_error,
+                artifacts=artifacts,
+            )
+        return ExecutionResult(success=True, output=output, artifacts=artifacts)
+
+    def _wrap_stateful_code(self, code: str) -> str:
+        state_path = _escape_baltamatica_string(self.state_file.as_posix())
+        return (
+            f"if exist('{state_path}','file'), load('{state_path}'); end; "
+            f"{code}; "
+            f"save('{state_path}');"
+        )
+
+    def _wrap_readonly_code(self, code: str) -> str:
+        state_path = _escape_baltamatica_string(self.state_file.as_posix())
+        return f"if exist('{state_path}','file'), load('{state_path}'); end; {code};"
+
+    def _run_cli_sync(self, argv: Sequence[str]) -> CliProcessResult:
+        creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
+        try:
+            process = subprocess.run(
+                argv,
+                capture_output=True,
+                check=False,
+                stdin=subprocess.DEVNULL,
+                creationflags=creationflags,
+                timeout=self.timeout,
+            )
+        except subprocess.TimeoutExpired as exc:
+            raise TimeoutError(
+                f"Baltamatica CLI timed out after {self.timeout:g} seconds."
+            ) from exc
+        except OSError as exc:
+            raise EngineUnavailableError(f"Failed to start Baltamatica CLI: {exc}") from exc
+
+        return CliProcessResult(
+            returncode=process.returncode,
+            stdout=_decode_process_output(process.stdout),
+            stderr=_decode_process_output(process.stderr),
+        )
+
+
+def _combine_output(stdout: str | None, stderr: str | None) -> str:
+    output = _clean_process_output(stdout or "")
+    error_output = _clean_process_output(stderr or "")
+    if output and error_output:
+        return f"{output}\n{error_output}"
+    return output or error_output
+
+
+def _decode_process_output(output: bytes | None) -> str:
+    if not output:
+        return ""
+    for encoding in ("utf-8", locale.getpreferredencoding(False), "gbk"):
+        try:
+            return output.decode(encoding)
+        except UnicodeDecodeError:
+            continue
+    return output.decode("utf-8", errors="replace")
+
+
+def _clean_process_output(output: str) -> str:
+    lines = []
+    for line in _strip_ansi(output).splitlines():
+        stripped = line.strip()
+        if any(stripped.startswith(pattern) for pattern in NOISE_OUTPUT_PATTERNS):
+            continue
+        lines.append(line.rstrip())
+    return "\n".join(lines).strip()
+
+
+def _process_error(process: CliProcessResult) -> str:
+    message = process.stderr.strip() or process.stdout.strip()
+    if message:
+        return f"Baltamatica CLI exited with code {process.returncode}: {message}"
+    return f"Baltamatica CLI exited with code {process.returncode}."
+
+
+def detect_baltamatica_error(output: str) -> str | None:
+    """Detect Baltamatica error text when the CLI exits with status 0."""
+
+    plain_output = _strip_ansi(output)
+    for pattern in ERROR_OUTPUT_PATTERNS:
+        if pattern in plain_output:
+            first_line = next(
+                (line.strip() for line in plain_output.splitlines() if line.strip()),
+                plain_output.strip(),
+            )
+            return f"Baltamatica reported an error: {first_line}"
+    return None
+
+
+def _strip_ansi(value: str) -> str:
+    return ANSI_PATTERN.sub("", value)
+
+
+def _escape_baltamatica_string(value: str) -> str:
+    return value.replace("'", "''")
+
+
+def _default_state_file() -> Path:
+    return Path(
+        tempfile.NamedTemporaryFile(prefix="baltamatica_mcp_", suffix=".mat", delete=True).name
+    )
+
+
+def parse_whos_output(output: str) -> list[VariableInfo]:
+    variables: list[VariableInfo] = []
+    for line in output.splitlines():
+        match = WHOS_ROW_PATTERN.match(line)
+        if not match:
+            continue
+        variables.append(
+            VariableInfo(
+                name=match.group("name"),
+                size=match.group("size"),
+                bytes=int(match.group("bytes")),
+                class_name=match.group("class_name"),
+                attributes=(match.group("attributes") or "").strip(),
+            )
+        )
+    return variables
+
+
+def parse_artifacts(output: str) -> list[Artifact]:
+    artifacts: list[Artifact] = []
+    for line in output.splitlines():
+        stripped = line.strip()
+        if not stripped.startswith(ARTIFACT_PREFIX):
+            continue
+        spec = stripped.removeprefix(ARTIFACT_PREFIX).strip()
+        if not spec:
+            continue
+        artifact_type, artifact_path = _parse_artifact_spec(spec)
+        path = Path(artifact_path).expanduser().resolve()
+        artifacts.append(
+            Artifact(
+                path=str(path),
+                type=artifact_type,
+                exists=path.exists() and path.is_file(),
+                size=path.stat().st_size if path.exists() and path.is_file() else 0,
+            )
+        )
+    return artifacts
+
+
+def _parse_artifact_spec(spec: str) -> tuple[str, str]:
+    if ":" not in spec:
+        return _guess_artifact_type(spec), spec
+
+    maybe_type, maybe_path = spec.split(":", 1)
+    if "/" not in maybe_type:
+        return _guess_artifact_type(spec), spec
+    return maybe_type.strip(), maybe_path.strip()
+
+
+def _guess_artifact_type(path: str) -> str:
+    suffix = Path(path).suffix.lower()
+    if suffix == ".png":
+        return "image/png"
+    if suffix in {".jpg", ".jpeg"}:
+        return "image/jpeg"
+    if suffix == ".svg":
+        return "image/svg+xml"
+    if suffix == ".pdf":
+        return "application/pdf"
+    if suffix == ".csv":
+        return "text/csv"
+    return "application/octet-stream"