obsidian-wiki/obsidian_wiki/mcp_server.py

"""stdio MCP server exposing safe wiki source-file resources."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Any

from obsidian_wiki.config import resolve_vault_path
from obsidian_wiki.file_assets import clean_config_value, db_path_for_vault, registry_path_for_vault
from obsidian_wiki.mcp_assets import (
    DEFAULT_MAX_FILE_BYTES,
    DEFAULT_PROFILE_NAME,
    AssetAccessError,
    AssetNotFoundError,
    AssetTooLargeError,
    get_file_asset_metadata as core_get_file_asset_metadata,
    load_registry_by_file_id,
    mcp_runtime_options,
    read_file_asset_resource,
    search_file_assets as core_search_file_assets,
)


def create_server(
    *,
    vault_path: str | Path,
    db_path: str | Path | None = None,
    profile_name: str = DEFAULT_PROFILE_NAME,
    max_file_bytes: int = DEFAULT_MAX_FILE_BYTES,
    expose_paths: bool = False,
):
    if sys.version_info < (3, 10):
        raise RuntimeError("obsidian-wiki MCP server requires Python 3.10 or newer")
    try:
        from mcp.server.fastmcp import FastMCP
    except ImportError as exc:
        raise RuntimeError('MCP support is not installed. Run: pip install "obsidian-wiki[mcp]"') from exc

    vault = Path(clean_config_value(str(vault_path))).expanduser()
    db = _resolve_db_path(vault, db_path)
    server = FastMCP("obsidian-wiki")

    @server.tool()
    def file_index_status() -> dict[str, Any]:
        """Show the configured vault, file index, registry, and MCP limits."""
        registry = registry_path_for_vault(vault)
        return {
            "vault_path": str(vault),
            "db_path": str(db),
            "db_exists": db.is_file(),
            "registry_path": str(registry),
            "registry_exists": registry.is_file(),
            "source_registry_records": len(load_registry_by_file_id(vault)),
            "profile_name": profile_name,
            "max_file_bytes": max_file_bytes,
            "expose_paths": expose_paths,
        }

    @server.tool()
    def search_file_assets(
        query: str,
        doc_type: str | None = None,
        source_type: str | None = None,
        product_line: str | None = None,
        customer: str | None = None,
        competitor: str | None = None,
        limit: int = 10,
    ) -> list[dict[str, Any]]:
        """Search safe engineering source files such as drawings, 3D models, specs, and test reports."""
        return _as_mcp_results(
            core_search_file_assets(
                vault,
                query,
                db_path=db,
                limit=limit,
                doc_type=doc_type,
                source_type=source_type,
                product_line=product_line,
                customer=customer,
                competitor=competitor,
                profile_name=profile_name,
                expose_paths=expose_paths,
            )
        )

    @server.tool()
    def get_file_asset_metadata(file_id: str) -> dict[str, Any]:
        """Return safe metadata and the resource URI for one indexed source file."""
        return _as_mcp_result(
            core_get_file_asset_metadata(
                vault,
                file_id,
                db_path=db,
                profile_name=profile_name,
                expose_paths=expose_paths,
            )
        )

    @server.resource("asset://source/sha256/{digest}")
    def read_file_asset(digest: str) -> bytes:
        """Read one allowed source file as a binary MCP resource."""
        resource = read_file_asset_resource(
            vault,
            f"asset://source/sha256/{digest}",
            db_path=db,
            profile_name=profile_name,
            max_file_bytes=max_file_bytes,
        )
        return resource["bytes"]

    return server


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="obsidian-wiki-mcp", description="Run the obsidian-wiki MCP server")
    parser.add_argument("--vault", metavar="PATH", help="absolute path to the Obsidian vault")
    parser.add_argument("--db", metavar="PATH", help="SQLite file asset index path")
    parser.add_argument("--profile", default=None, help=f"bot profile name, default: {DEFAULT_PROFILE_NAME}")
    parser.add_argument("--max-file-bytes", type=int, default=None, help="maximum resource size, 0 disables the limit")
    parser.add_argument("--expose-paths", action="store_true", help="include archive/absolute paths in tool output")
    args = parser.parse_args(argv)

    vault = resolve_vault_path(args.vault)
    if not vault:
        print("error: vault path is not configured; run `obsidian-wiki setup --vault /path/to/vault`", file=sys.stderr)
        return 1
    options = mcp_runtime_options(
        profile_name=args.profile,
        max_file_bytes=args.max_file_bytes,
        expose_paths=args.expose_paths,
    )
    try:
        server = create_server(
            vault_path=vault,
            db_path=args.db,
            profile_name=options["profile_name"],
            max_file_bytes=options["max_file_bytes"],
            expose_paths=options["expose_paths"],
        )
        try:
            server.run(transport="stdio")
        except TypeError:
            server.run()
    except (RuntimeError, FileNotFoundError, AssetAccessError, AssetNotFoundError, AssetTooLargeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


def _resolve_db_path(vault: Path, raw_path: str | Path | None) -> Path:
    if raw_path is None:
        return db_path_for_vault(vault)
    db_path = Path(clean_config_value(str(raw_path))).expanduser()
    if not db_path.is_absolute():
        db_path = vault / db_path
    return db_path


def _as_mcp_results(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [_as_mcp_result(row) for row in rows]


def _as_mcp_result(row: dict[str, Any]) -> dict[str, Any]:
    resource_uri = str(row.get("resource_uri") or "")
    return {
        **row,
        "resource_link": {
            "uri": resource_uri,
            "name": row.get("original_name") or row.get("title") or resource_uri,
            "mime_type": row.get("mime_type") or "application/octet-stream",
        },
    }


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