UmbraAudio/infra/mcp_shell_server.py

#!/usr/bin/env python3
"""Minimal MCP stdio server exposing an arbitrary-command tool.

Deliberately dependency-free (Python stdlib only) so every line can be read
and audited before granting command execution on this machine. A third-party
package with this capability is an opaque trust decision; this is a few
hundred lines you can verify yourself.

Transport: JSON-RPC 2.0 over stdio, one message per line. See
https://modelcontextprotocol.io/specification/transports/stdio

Registered in opencode.json as a local MCP server. Verify before wiring up:

    .venv/bin/python infra/mcp_shell_server.py --selftest
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from pathlib import Path

SERVER_NAME = "umbra-shell"
SERVER_VERSION = "1.0.0"
FALLBACK_PROTOCOL_VERSION = "2024-11-05"

DEFAULT_TIMEOUT_S = 120
MAX_TIMEOUT_S = 3600
MAX_CAPTURE_CHARS = 200_000

# macOS launchd, and anything spawned from a GUI context, hands down a bare
# PATH. That silently breaks every homebrew-resolved tool the project depends
# on -- ffmpeg (m4b/MP3 export) is the documented victim in AGENTS.md. An MCP
# server is spawned by the UI, not by launchd, but the same bare-PATH risk
# applies, so prepend them explicitly rather than trusting the parent.
EXTRA_PATH_ENTRIES = ("/opt/homebrew/bin", "/usr/local/bin")

DEFAULT_PROJECT_ROOT = Path(
    os.environ.get("UMBRA_SHELL_PROJECT_ROOT") or Path(__file__).resolve().parent.parent
)

LOG_DIR = Path(
    os.environ.get("UMBRA_SHELL_LOG_DIR") or (Path.home() / ".umbra" / "shell_logs")
)


class ToolError(Exception):
    """A caller-correctable problem, reported back as a tool result."""


SHELL_TOOL: dict = {
    "name": "shell",
    "description": (
        "Run a shell command on the local machine; returns exit code, stdout "
        "and stderr. Unfiltered by design. Use background=true for long-running "
        "processes (servers, batch audio generation) so the call returns a PID "
        "and log path instead of blocking until timeout."
    ),
    "inputSchema": {
        "type": "object",
        "properties": {
            "command": {
                "type": "string",
                "description": "Command line to execute, passed to the shell.",
            },
            "cwd": {
                "type": "string",
                "description": (
                    "Working directory. Defaults to the UmbraAudio project root."
                ),
            },
            "timeout": {
                "type": "integer",
                "description": (
                    f"Seconds before a foreground command is killed "
                    f"(default {DEFAULT_TIMEOUT_S}, max {MAX_TIMEOUT_S})."
                ),
            },
            "background": {
                "type": "boolean",
                "description": (
                    "Detach the process into its own session and return its PID "
                    "and log path immediately."
                ),
            },
        },
        "required": ["command"],
    },
}


def _child_env() -> dict[str, str]:
    """Parent environment with a corrected PATH, for the child process only."""
    env = dict(os.environ)
    path = os.pathsep.join(
        entry for entry in (*EXTRA_PATH_ENTRIES, env.get("PATH", "")) if entry
    )
    env["PATH"] = path
    env.setdefault("PYTHONUNBUFFERED", "1")
    return env


def _truncate(text: str) -> tuple[str, bool]:
    if len(text) <= MAX_CAPTURE_CHARS:
        return text, False
    half = MAX_CAPTURE_CHARS // 2
    omitted = len(text) - (2 * half)
    return f"{text[:half]}\n... [truncated {omitted} chars] ...\n{text[-half:]}", True


def _as_text(raw: object) -> str:
    """TimeoutExpired carries bytes or str depending on how it was started."""
    if raw is None:
        return ""
    if isinstance(raw, bytes):
        return raw.decode("utf-8", errors="replace")
    return str(raw)


def _is_not_json(line: str) -> bool:
    try:
        json.loads(line)
    except json.JSONDecodeError:
        return True
    return False


def _tool_payload(reply: dict) -> dict:
    """Unwrap a tool result.

    The tool's stdout is a JSON document *nested inside* the JSON-RPC reply,
    so its quotes are escaped. Matching on the serialized outer reply silently
    never hits -- comparing `'"exit_code": 3'` against `\\"exit_code\\": 3`.
    Parse the inner document instead of pattern-matching the outer one.
    """
    try:
        return json.loads(reply["result"]["content"][0]["text"])
    except (KeyError, IndexError, TypeError, json.JSONDecodeError):
        return {}


def _run_foreground(command: str, cwd: Path, timeout: int) -> dict:
    started = time.monotonic()
    timed_out = False
    try:
        proc = subprocess.run(
            command,
            shell=True,
            cwd=str(cwd),
            env=_child_env(),
            capture_output=True,
            text=True,
            errors="replace",
            timeout=timeout,
        )
        stdout, stderr, exit_code = proc.stdout or "", proc.stderr or "", proc.returncode
    except subprocess.TimeoutExpired as exc:
        stdout, stderr, exit_code, timed_out = _as_text(exc.stdout), _as_text(exc.stderr), -1, True
    except OSError as exc:
        raise ToolError(f"Could not execute command: {exc}") from exc

    stdout, stdout_truncated = _truncate(stdout)
    stderr, stderr_truncated = _truncate(stderr)
    return {
        "exit_code": exit_code,
        "timed_out": timed_out,
        "duration_s": round(time.monotonic() - started, 2),
        "truncated": stdout_truncated or stderr_truncated,
        "cwd": str(cwd),
        "stdout": stdout,
        "stderr": stderr,
    }


def _run_background(command: str, cwd: Path) -> dict:
    """Detach into its own session so the process outlives this server.

    start_new_session is what makes `KeepAlive`-style servers survive: the child
    leaves our process group, so neither the MCP server exiting nor the UI
    restarting tears it down. Output goes to a file rather than a pipe,
    because an unread pipe eventually blocks the writer.
    """
    try:
        LOG_DIR.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        raise ToolError(f"Cannot create log directory {LOG_DIR}: {exc}") from exc

    log = LOG_DIR / f"bg_{time.strftime('%Y%m%dT%H%M%S')}_{os.getpid()}.log"
    try:
        # The child dups this descriptor into its own session, so closing our
        # copy afterwards is safe and required.
        with open(log, "wb") as handle:
            proc = subprocess.Popen(
                command,
                shell=True,
                cwd=str(cwd),
                env=_child_env(),
                stdout=handle,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
    except OSError as exc:
        raise ToolError(f"Could not start background command: {exc}") from exc

    return {"pid": proc.pid, "log": str(log), "cwd": str(cwd), "detached": True}


def _handle_tools_call(request: dict) -> dict:
    params = request.get("params") or {}
    name = params.get("name")
    if name != "shell":
        raise ToolError(f"Unknown tool: {name!r}. Only 'shell' is exposed.")

    args = params.get("arguments") or {}
    command = args.get("command")
    if not isinstance(command, str) or not command.strip():
        raise ToolError("`command` is required and must be a non-empty string.")

    cwd = Path(str(args.get("cwd") or DEFAULT_PROJECT_ROOT)).expanduser()
    if not cwd.is_dir():
        raise ToolError(f"cwd does not exist or is not a directory: {cwd}")

    if args.get("background"):
        payload = _run_background(command, cwd)
    else:
        timeout = int(args.get("timeout") or DEFAULT_TIMEOUT_S)
        payload = _run_foreground(command, cwd, max(1, min(timeout, MAX_TIMEOUT_S)))

    return {
        "content": [{"type": "text", "text": json.dumps(payload, indent=2)}],
        "isError": False,
    }


def _handle_initialize(request: dict) -> dict:
    requested = (request.get("params") or {}).get("protocolVersion")
    return {
        "protocolVersion": requested or FALLBACK_PROTOCOL_VERSION,
        "capabilities": {"tools": {"listChanged": False}},
        "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
        "instructions": (
            "Single-purpose server. The `shell` tool executes arbitrary local "
            "commands with no filtering. Source: infra/mcp_shell_server.py."
        ),
    }


def _send(message: dict) -> None:
    sys.stdout.write(json.dumps(message) + "\n")
    sys.stdout.flush()


def _reply(request_id: object, result: dict) -> None:
    _send({"jsonrpc": "2.0", "id": request_id, "result": result})


def _reply_error(request_id: object, code: int, message: str) -> None:
    _send({"jsonrpc": "2.0", "id": request_id, "error": {"code": code, "message": message}})


def serve() -> int:
    """Blocking read loop. stdout carries protocol only -- never print here."""
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
        except json.JSONDecodeError:
            continue

        request_id = request.get("id")
        if request_id is None:  # notification: no response is permitted
            continue

        method = request.get("method")
        try:
            if method == "initialize":
                _reply(request_id, _handle_initialize(request))
            elif method == "tools/list":
                _reply(request_id, {"tools": [SHELL_TOOL]})
            elif method == "tools/call":
                _reply(request_id, _handle_tools_call(request))
            elif method == "ping":
                _reply(request_id, {})
            else:
                _reply_error(request_id, -32601, f"Method not found: {method}")
        except ToolError as exc:
            _reply(
                request_id,
                {"content": [{"type": "text", "text": str(exc)}], "isError": True},
            )
        except Exception as exc:  # a crash here would look like a hung client
            _reply_error(request_id, -32603, f"{type(exc).__name__}: {exc}")
    return 0


def selftest() -> int:
    """Drive this server through a full handshake without an MCP client.

    Exists because the failure mode -- server present but misbehaving -- is
    otherwise invisible until an OpenCode restart, which is expensive.
    """
    transcript = [
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {"protocolVersion": "2024-11-05"},
        },
        {"jsonrpc": "2.0", "method": "notifications/initialized"},
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
        {
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "shell",
                "arguments": {
                    "command": "echo selftest-ok; echo path-probe:$(command -v git) 1>&2",
                },
            },
        },
        {
            "jsonrpc": "2.0",
            "id": 4,
            "method": "tools/call",
            "params": {"name": "shell", "arguments": {"command": "exit 3"}},
        },
        {
            "jsonrpc": "2.0",
            "id": 5,
            "method": "tools/call",
            "params": {
                "name": "shell",
                "arguments": {
                    "command": "printf to-out; printf to-err 1>&2; exit 7"
                },
            },
        },
    ]
    payload = "".join(json.dumps(message) + "\n" for message in transcript)

    proc = subprocess.run(
        [sys.executable, os.path.abspath(__file__)],
        input=payload,
        capture_output=True,
        text=True,
        timeout=120,
    )

    replies = {}
    for line in proc.stdout.splitlines():
        try:
            message = json.loads(line)
        except json.JSONDecodeError:
            continue
        if "id" in message:
            replies[message["id"]] = message

    malformed = [line for line in proc.stdout.splitlines() if _is_not_json(line)]
    init_result = replies.get(1, {}).get("result", {})
    list_tools = replies.get(2, {}).get("result", {}).get("tools", [])
    third = _tool_payload(replies.get(3, {}))
    fourth = _tool_payload(replies.get(4, {}))
    fifth = _tool_payload(replies.get(5, {}))

    checks: list[tuple[str, bool]] = [
        (
            "initialize returns serverInfo",
            init_result.get("serverInfo", {}).get("name") == SERVER_NAME,
        ),
        (
            "tools/list exposes shell",
            any(tool.get("name") == "shell" for tool in list_tools),
        ),
        ("tools/call executes", "selftest-ok" in third.get("stdout", "")),
        ("stderr captured, not leaked", "path-probe:" in third.get("stderr", "")),
        ("exit code propagates", fourth.get("exit_code") == 3),
        (
            "stdout/stderr kept apart",
            fifth.get("stdout") == "to-out" and fifth.get("stderr") == "to-err",
        ),
        (
            "non-zero exit from a real command",
            fifth.get("exit_code") == 7,
        ),
        (
            "stdout is protocol-only (every line valid JSON)",
            not malformed,
        ),
    ]

    failed = 0
    for label, ok in checks:
        if not ok:
            failed += 1
        print(f"  {'PASS' if ok else 'FAIL'}  {label}")

    git_path = third.get("stderr", "").split("path-probe:", 1)[-1].strip()
    if git_path:
        print(f"\n  PATH resolves git to: {git_path}")
    print(f"\n{len(checks) - failed}/{len(checks)} checks passed")
    return 1 if failed else 0


def main() -> int:
    if "--selftest" in sys.argv:
        return selftest()
    return serve()


if __name__ == "__main__":
    raise SystemExit(main())