obsidian-wiki/scripts/manage_mcp_service.py

#!/usr/bin/env python3
"""Manage the user-level HTTP MCP service for obsidian-wiki source files."""

from __future__ import annotations

import argparse
import os
import secrets
import shlex
import shutil
import socket
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence


SERVICE_NAME = "obsidian-wiki-mcp.service"
ENV_RELATIVE_PATH = Path(".runtime") / "obsidian-wiki-mcp.env"
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8765
DEFAULT_PATH = "/mcp"
DEFAULT_TOKEN_BYTES = 32
DEFAULT_UV_CACHE_DIR = "/tmp/uv-cache"


@dataclass(frozen=True)
class ServicePaths:
    project_root: Path
    unit_path: Path
    env_path: Path


def generate_token(num_bytes: int = DEFAULT_TOKEN_BYTES) -> str:
    return secrets.token_urlsafe(num_bytes)


def systemd_quote(value: str | Path) -> str:
    return shlex.quote(str(value))


def find_uv_path() -> Path:
    configured = os.environ.get("UV_PATH", "").strip()
    discovered = configured or shutil.which("uv")
    if not discovered:
        raise RuntimeError("Could not find uv on PATH. Install uv or set UV_PATH to its absolute path.")
    return Path(discovered).expanduser().resolve()


def resolve_service_paths(project_root: Path, config_home: Path | None = None) -> ServicePaths:
    root = project_root.expanduser().resolve()
    base_config = config_home or Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"))
    unit_path = base_config.expanduser() / "systemd" / "user" / SERVICE_NAME
    return ServicePaths(project_root=root, unit_path=unit_path, env_path=root / ENV_RELATIVE_PATH)


def render_service_unit(
    *,
    project_root: Path,
    env_file: Path,
    host: str = DEFAULT_HOST,
    port: int = DEFAULT_PORT,
    path: str = DEFAULT_PATH,
    uv_cache_dir: str = DEFAULT_UV_CACHE_DIR,
    uv_path: Path | None = None,
) -> str:
    resolved_uv_path = uv_path or find_uv_path()
    command = " ".join(
        [
            systemd_quote(resolved_uv_path),
            "run",
            "--no-sync",
            "obsidian-wiki-mcp",
            "--vault",
            systemd_quote(project_root),
            "--transport",
            "streamable-http",
            "--host",
            host,
            "--port",
            str(port),
            "--path",
            path,
        ]
    )
    return (
        "[Unit]\n"
        "Description=obsidian-wiki readonly HTTP MCP server\n"
        "After=network-online.target\n"
        "Wants=network-online.target\n"
        "\n"
        "[Service]\n"
        "Type=simple\n"
        f"WorkingDirectory={systemd_quote(project_root)}\n"
        f"EnvironmentFile={systemd_quote(env_file)}\n"
        f"Environment=UV_CACHE_DIR={systemd_quote(uv_cache_dir)}\n"
        f"ExecStart={command}\n"
        "Restart=on-failure\n"
        "RestartSec=5\n"
        "\n"
        "[Install]\n"
        "WantedBy=default.target\n"
    )


def write_token_env_file(
    env_path: Path,
    token: str,
    *,
    host: str = DEFAULT_HOST,
    port: int = DEFAULT_PORT,
    path: str = DEFAULT_PATH,
) -> None:
    env_path.parent.mkdir(parents=True, exist_ok=True)
    env_path.parent.chmod(0o700)
    fd = os.open(env_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as env_file:
        env_file.write(f"OBSIDIAN_MCP_TOKEN={token}\n")
        env_file.write(f"OBSIDIAN_MCP_HOST={host}\n")
        env_file.write(f"OBSIDIAN_MCP_PORT={port}\n")
        env_file.write(f"OBSIDIAN_MCP_PATH={path}\n")
    env_path.chmod(0o600)


def read_service_env_file(env_path: Path) -> dict[str, str]:
    if not env_path.exists():
        raise RuntimeError(f"Missing MCP service env file: {env_path}. Run scripts/new_mcp_service_token.sh first.")

    values: dict[str, str] = {}
    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        key, separator, value = line.partition("=")
        if separator:
            values[key.strip()] = value.strip()
    return values


def saved_service_config(project_root: Path) -> tuple[str, str, int, str]:
    paths = resolve_service_paths(project_root)
    values = read_service_env_file(paths.env_path)
    token = values.get("OBSIDIAN_MCP_TOKEN", "").strip()
    if not token:
        raise RuntimeError(f"Missing OBSIDIAN_MCP_TOKEN in {paths.env_path}")
    host = values.get("OBSIDIAN_MCP_HOST") or DEFAULT_HOST
    port = int(values.get("OBSIDIAN_MCP_PORT") or DEFAULT_PORT)
    path = values.get("OBSIDIAN_MCP_PATH") or DEFAULT_PATH
    return token, host, port, path


def write_service_unit(
    paths: ServicePaths,
    *,
    host: str = DEFAULT_HOST,
    port: int = DEFAULT_PORT,
    path: str = DEFAULT_PATH,
) -> None:
    paths.unit_path.parent.mkdir(parents=True, exist_ok=True)
    unit = render_service_unit(
        project_root=paths.project_root,
        env_file=paths.env_path,
        host=host,
        port=port,
        path=path,
    )
    paths.unit_path.write_text(unit, encoding="utf-8")


def detect_remote_hosts() -> list[str]:
    hosts: list[str] = []
    try:
        result = subprocess.run(["hostname", "-I"], capture_output=True, text=True, check=False)
    except OSError:
        result = None
    if result and result.returncode == 0:
        hosts.extend(result.stdout.split())

    try:
        for info in socket.getaddrinfo(socket.gethostname(), None):
            hosts.append(info[4][0])
    except OSError:
        pass

    cleaned: list[str] = []
    for host in hosts:
        normalized = host.strip()
        if not normalized or normalized.startswith("127.") or normalized == "::1" or normalized.startswith("fe80:"):
            continue
        if normalized not in cleaned:
            cleaned.append(normalized)
    return cleaned


def render_remote_config(
    *,
    hosts: Sequence[str],
    token: str,
    port: int = DEFAULT_PORT,
    path: str = DEFAULT_PATH,
) -> str:
    host = hosts[0] if hosts else "<server-ip>"
    url = f"http://{host}:{port}{path}"
    alternatives = "\n".join(f"- http://{candidate}:{port}{path}" for candidate in hosts[1:])
    alternatives_section = f"\nOther detected addresses:\n{alternatives}\n" if alternatives else ""
    return (
        "MCP service is running.\n"
        "\n"
        "Use this URL from another machine:\n"
        f"{url}\n"
        f"{alternatives_section}"
        "\n"
        "Use this header:\n"
        f"Authorization: Bearer {token}\n"
        "\n"
        "Example MCP client config:\n"
        "{\n"
        '  "mcpServers": {\n'
        '    "obsidian-wiki": {\n'
        '      "type": "streamable-http",\n'
        f'      "url": "{url}",\n'
        '      "headers": {\n'
        f'        "Authorization": "Bearer {token}"\n'
        "      }\n"
        "    }\n"
        "  }\n"
        "}\n"
    )


def show_config(project_root: Path) -> str:
    token, _host, port, path = saved_service_config(project_root)
    return render_remote_config(hosts=detect_remote_hosts(), token=token, port=port, path=path)


def run_systemctl(args: Sequence[str], *, check: bool = True) -> subprocess.CompletedProcess[str]:
    return subprocess.run(["systemctl", "--user", *args], text=True, check=check)


def restart_service_if_installed(paths: ServicePaths) -> bool:
    if not paths.unit_path.exists():
        return False
    run_systemctl(["restart", SERVICE_NAME])
    return True


def new_token(
    *,
    project_root: Path,
    host: str = DEFAULT_HOST,
    port: int = DEFAULT_PORT,
    path: str = DEFAULT_PATH,
    token_bytes: int = DEFAULT_TOKEN_BYTES,
) -> str:
    paths = resolve_service_paths(project_root)
    token = generate_token(token_bytes)
    write_token_env_file(paths.env_path, token, host=host, port=port, path=path)
    restart_service_if_installed(paths)
    return render_remote_config(hosts=detect_remote_hosts(), token=token, port=port, path=path)


def start_service(
    *,
    project_root: Path,
) -> str:
    paths = resolve_service_paths(project_root)
    token, host, port, path = saved_service_config(project_root)
    write_service_unit(paths, host=host, port=port, path=path)

    run_systemctl(["daemon-reload"])
    run_systemctl(["enable", SERVICE_NAME])
    run_systemctl(["restart", SERVICE_NAME])
    return render_remote_config(hosts=detect_remote_hosts(), token=token, port=port, path=path)


def stop_service() -> None:
    run_systemctl(["stop", SERVICE_NAME])


def uninstall_service(project_root: Path) -> None:
    paths = resolve_service_paths(project_root)
    run_systemctl(["disable", "--now", SERVICE_NAME], check=False)
    if paths.unit_path.exists():
        paths.unit_path.unlink()
    if paths.env_path.exists():
        paths.env_path.unlink()
    try:
        paths.env_path.parent.rmdir()
    except OSError:
        pass
    run_systemctl(["daemon-reload"], check=False)
    run_systemctl(["reset-failed", SERVICE_NAME], check=False)


def service_status() -> None:
    run_systemctl(["status", SERVICE_NAME], check=False)


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Manage the readonly HTTP obsidian-wiki MCP user service.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    new_token_parser = subparsers.add_parser("new-token", help="Generate a token and print remote MCP client configuration.")
    new_token_parser.add_argument("--project-root", default=Path(__file__).resolve().parent.parent, type=Path)
    new_token_parser.add_argument("--host", default=DEFAULT_HOST)
    new_token_parser.add_argument("--port", default=DEFAULT_PORT, type=int)
    new_token_parser.add_argument("--path", default=DEFAULT_PATH)
    new_token_parser.add_argument("--token-bytes", default=DEFAULT_TOKEN_BYTES, type=int)

    start_parser = subparsers.add_parser("start", help="Install or update the user service using the saved token.")
    start_parser.add_argument("--project-root", default=Path(__file__).resolve().parent.parent, type=Path)

    subparsers.add_parser("stop", help="Stop the user service without deleting service files.")

    uninstall_parser = subparsers.add_parser("uninstall", help="Stop and remove the managed user service and token file.")
    uninstall_parser.add_argument("--project-root", default=Path(__file__).resolve().parent.parent, type=Path)

    show_config_parser = subparsers.add_parser("show-config", help="Print the saved token and remote MCP client configuration.")
    show_config_parser.add_argument("--project-root", default=Path(__file__).resolve().parent.parent, type=Path)

    subparsers.add_parser("status", help="Show systemd status for the managed service.")

    args = parser.parse_args(argv)

    if args.command == "new-token":
        print(
            new_token(
                project_root=args.project_root,
                host=args.host,
                port=args.port,
                path=args.path,
                token_bytes=args.token_bytes,
            ),
            end="",
        )
        return 0
    if args.command == "start":
        print(start_service(project_root=args.project_root), end="")
        return 0
    if args.command == "stop":
        stop_service()
        print(f"Stopped {SERVICE_NAME}.")
        return 0
    if args.command == "uninstall":
        uninstall_service(args.project_root)
        print(f"Uninstalled {SERVICE_NAME}.")
        return 0
    if args.command == "show-config":
        print(show_config(args.project_root), end="")
        return 0
    if args.command == "status":
        service_status()
        return 0
    raise AssertionError(f"unhandled command: {args.command}")


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