obsidian-wiki/obsidian_wiki/mcp_server.py

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

from __future__ import annotations

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

from obsidian_wiki.config import config_value, 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,
)

DEFAULT_HTTP_HOST = "127.0.0.1"
DEFAULT_HTTP_PORT = 8765
DEFAULT_HTTP_PATH = "/mcp"
DEFAULT_TOKEN_ENV = "OBSIDIAN_MCP_TOKEN"
MCP_SERVER_NAME = "wiki-file-assets"


def bearer_token(authorization: str) -> str:
    scheme, _, token = authorization.partition(" ")
    if scheme.lower() != "bearer" or not token:
        return ""
    return token.strip()


class BearerTokenAuthMiddleware:
    def __init__(self, app: Any, token: str) -> None:
        if not token:
            raise ValueError("HTTP MCP token must not be empty")
        self.app = app
        self.token = token

    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return
        headers = {key.lower(): value for key, value in scope.get("headers", [])}
        authorization = headers.get(b"authorization", b"").decode("latin1")
        token = bearer_token(authorization)
        if not hmac.compare_digest(token, self.token):
            await send(
                {
                    "type": "http.response.start",
                    "status": 401,
                    "headers": [
                        (b"content-type", b"application/json"),
                        (b"www-authenticate", b"Bearer"),
                    ],
                }
            )
            await send({"type": "http.response.body", "body": b'{"error":"unauthorized"}'})
            return
        await self.app(scope, receive, send)


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,
    host: str = DEFAULT_HTTP_HOST,
    port: int = DEFAULT_HTTP_PORT,
    streamable_http_path: str = DEFAULT_HTTP_PATH,
    allowed_hosts: list[str] | None = None,
    allowed_origins: list[str] | None = None,
):
    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
        try:
            from mcp.server.fastmcp.server import TransportSecuritySettings
        except ImportError:
            TransportSecuritySettings = None
    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)
    transport_security = None
    if TransportSecuritySettings is not None and (allowed_hosts is not None or allowed_origins is not None):
        transport_security = TransportSecuritySettings(
            allowed_hosts=allowed_hosts or [],
            allowed_origins=allowed_origins or [],
        )
    server = FastMCP(
        MCP_SERVER_NAME,
        host=host,
        port=port,
        streamable_http_path=streamable_http_path,
        transport_security=transport_security,
    )

    @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 list_file_assets(
        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]]:
        """List safe source files matching optional filters. Use search_file_assets when the user gives keywords such as product codes, sizes, or drawing names."""
        return _as_mcp_results(
            core_search_file_assets(
                vault,
                "",
                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}", mime_type="application/octet-stream")
    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"]

    @server.prompt(
        "how_to_use_wiki_file_assets",
        description="Explain how to search and retrieve original source files from the wiki-file-assets MCP server.",
    )
    def how_to_use_wiki_file_assets() -> str:
        return (
            "wiki-file-assets serves original source files archived from an Obsidian wiki vault. "
            "Use search_file_assets for keyword searches such as product codes, drawing names, sizes, "
            "or document types. Use list_file_assets to browse allowed engineering files with optional "
            "filters. Use get_file_asset_metadata with a file_id from search/list results to get the "
            "asset://source/sha256/... URI, then read that resource to fetch the original file bytes. "
            "This MCP is intentionally limited to safe engineering files such as drawings, CAD/3D models, "
            "specifications, product parameters, and test reports; it blocks raw orders, prices, email "
            "correspondence, complaint evidence, claims, and personal data."
        )

    return server


def read_token_from_env(env_name: str = DEFAULT_TOKEN_ENV) -> str:
    token = os.environ.get(env_name, "").strip()
    if not token:
        raise RuntimeError(f"{env_name} must be set for HTTP MCP transport")
    return token


def runtime_config_value(key: str, default: str) -> str:
    raw = os.environ.get(key)
    if raw is None:
        raw = config_value(key, default)
    value = clean_config_value(raw).strip()
    return value or default


def runtime_config_int(key: str, default: int) -> int:
    raw = runtime_config_value(key, str(default))
    try:
        return int(raw)
    except ValueError:
        return default


def create_http_app(server: Any, token: str) -> Any:
    return BearerTokenAuthMiddleware(server.streamable_http_app(), token)


def run_http_server(
    *,
    vault_path: str | Path,
    db_path: str | Path | None,
    profile_name: str,
    max_file_bytes: int,
    expose_paths: bool,
    host: str,
    port: int,
    token_env: str,
    streamable_http_path: str,
    allowed_hosts: list[str] | None,
    allowed_origins: list[str] | None,
) -> None:
    token = read_token_from_env(token_env)
    try:
        import uvicorn
    except ImportError as exc:
        raise RuntimeError('HTTP MCP support requires uvicorn. Run: pip install "obsidian-wiki[mcp]"') from exc

    server = create_server(
        vault_path=vault_path,
        db_path=db_path,
        profile_name=profile_name,
        max_file_bytes=max_file_bytes,
        expose_paths=expose_paths,
        host=host,
        port=port,
        streamable_http_path=streamable_http_path,
        allowed_hosts=allowed_hosts,
        allowed_origins=allowed_origins,
    )
    app = create_http_app(server, token)
    uvicorn.run(app, host=host, port=port, log_level="info")


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")
    parser.add_argument(
        "--transport",
        choices=["stdio", "streamable-http"],
        default=None,
        help="MCP transport. HTTP mode requires a bearer token.",
    )
    parser.add_argument("--host", default=None, help="HTTP bind host. Used only with --transport streamable-http.")
    parser.add_argument("--port", default=None, type=int, help="HTTP port. Used only with --transport streamable-http.")
    parser.add_argument("--path", default=None, help="Streamable HTTP MCP path.")
    parser.add_argument("--token-env", default=None, help="Environment variable containing the HTTP bearer token.")
    parser.add_argument(
        "--allowed-host",
        action="append",
        dest="allowed_hosts",
        help="Allowed HTTP Host pattern for DNS rebinding protection. Can be repeated.",
    )
    parser.add_argument(
        "--allowed-origin",
        action="append",
        dest="allowed_origins",
        help="Allowed Origin pattern for DNS rebinding protection. Can be repeated.",
    )
    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,
    )
    transport = args.transport or runtime_config_value("OBSIDIAN_MCP_TRANSPORT", "stdio")
    host = args.host or runtime_config_value("OBSIDIAN_MCP_HOST", DEFAULT_HTTP_HOST)
    port = args.port if args.port is not None else runtime_config_int("OBSIDIAN_MCP_PORT", DEFAULT_HTTP_PORT)
    http_path = args.path or runtime_config_value("OBSIDIAN_MCP_PATH", DEFAULT_HTTP_PATH)
    token_env = args.token_env or runtime_config_value("OBSIDIAN_MCP_TOKEN_ENV", DEFAULT_TOKEN_ENV)
    try:
        if transport == "stdio":
            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()
            return 0
        if transport != "streamable-http":
            raise RuntimeError(f"unsupported MCP transport: {transport}")
        run_http_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"],
            host=host,
            port=port,
            token_env=token_env,
            streamable_http_path=http_path,
            allowed_hosts=args.allowed_hosts,
            allowed_origins=args.allowed_origins,
        )
    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())