sanguo_llmwiki/mcp_server/main.py

"""
MCP Server - Wiki MCP Server 主入口

实现 MCP 协议解析、工具注册和路由。

参考设计文档：第 2.1 节
"""

import asyncio
import json
import logging
import sys
from typing import Any, Callable, Dict, Optional
from pathlib import Path

# MCP SDK
try:
    from mcp.server.models import InitializationOptions
    from mcp.server import Server, NotificationOptions
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    stdio_server = None
    logger = logging.getLogger(__name__)
    logger.warning("MCP SDK not available, using mock implementation")

# 本地模块
from .storage import Database, fix_dirty_states
from .services import CacheService, QueryService, ParserService, IndexerService, GraphService
from .tools import (
    WikiQueryTool, MemoryBridgeTool, WikiStatusTool,
    WikiLintTool, CrossLinkerTool, TagTaxonomyTool,
    WikiSynthesizeTool, DailyUpdateTool
)

# 配置
DEFAULT_WIKI_VAULT = "/Volumes/KnowledgeBase/wiki-vault"
DEFAULT_INDEX_PATH = "~/.sanguo-llmwiki/index.db"
DEFAULT_LOG_LEVEL = "INFO"

# 配置日志
logging.basicConfig(
    level=getattr(logging, DEFAULT_LOG_LEVEL),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class MCPServer:
    """MCP 服务器封装类"""

    def __init__(self, wiki_vault_path: str, index_path: str):
        self.wiki_vault_path = wiki_vault_path
        self.index_path = index_path
        self.db: Optional[Database] = None
        self.tools: Dict[str, Callable] = {}
        self.tool_schemas: Dict[str, dict] = {}

        # 如果 MCP SDK 可用，创建 Server 实例
        if MCP_AVAILABLE:
            self.server = Server("wiki-mcp")
            self._setup_handlers()
        else:
            self.server = None
            logger.warning("Running in mock mode (MCP SDK not available)")

    async def initialize(self) -> None:
        """初始化数据库和服务"""
        logger.info("Initializing Wiki MCP Server...")

        # 初始化数据库
        self.db = Database(self.index_path)
        await self.db.connect()

        # 恢复脏状态
        await fix_dirty_states(self.db)

        # 初始化服务
        self.cache = CacheService(max_size=1000)
        self.query_service = QueryService(self.db, self.cache)
        self.parser_service = ParserService()
        self.indexer_service = IndexerService(self.db, self.parser_service, self.wiki_vault_path)
        self.graph_service = GraphService(self.db)

        # 初始化工具
        self._init_tools()

        logger.info("Wiki MCP Server initialized")

    def _init_tools(self) -> None:
        """初始化所有工具"""
        wiki_query = WikiQueryTool(self.query_service)
        memory_bridge = MemoryBridgeTool(self.query_service)
        wiki_status = WikiStatusTool(self.query_service, self.indexer_service, self.graph_service)
        wiki_lint = WikiLintTool(self.parser_service, self.query_service, self.graph_service, self.wiki_vault_path)
        cross_linker = CrossLinkerTool(self.graph_service)
        tag_taxonomy = TagTaxonomyTool(self.query_service)
        wiki_synthesize = WikiSynthesizeTool(self.query_service, self.graph_service)
        daily_update = DailyUpdateTool(self.query_service, self.indexer_service, self.wiki_vault_path)

        # 注册工具
        self.tools = {
            "wiki_query": wiki_query.handle,
            "memory_bridge": memory_bridge.handle,
            "wiki_status": wiki_status.handle,
            "wiki_lint": wiki_lint.handle,
            "cross_linker": cross_linker.handle,
            "tag_taxonomy": tag_taxonomy.handle,
            "wiki_synthesize": wiki_synthesize.handle,
            "daily_update": daily_update.handle
        }

        # 注册工具 schema
        self.tool_schemas = {
            "wiki_query": wiki_query.get_schema(),
            "memory_bridge": memory_bridge.get_schema(),
            "wiki_status": wiki_status.get_schema(),
            "wiki_lint": wiki_lint.get_schema(),
            "cross_linker": cross_linker.get_schema(),
            "tag_taxonomy": tag_taxonomy.get_schema(),
            "wiki_synthesize": wiki_synthesize.get_schema(),
            "daily_update": daily_update.get_schema()
        }

        logger.info(f"Registered {len(self.tools)} tools")

    def _setup_handlers(self) -> None:
        """设置 MCP 处理器"""
        if not self.server:
            return

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """列出所有可用工具"""
            return [
                Tool(
                    name=schema["name"],
                    description=schema["description"],
                    inputSchema=schema["inputSchema"]
                )
                for schema in self.tool_schemas.values()
            ]

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict) -> list[TextContent]:
            """调用工具"""
            if name not in self.tools:
                raise ValueError(f"Unknown tool: {name}")

            try:
                result = await self.tools[name](**arguments)
                return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))]
            except Exception as e:
                # 记录完整堆栈用于调试
                logger.exception(f"Error calling tool {name} with arguments {arguments}")
                # 返回用户友好的错误消息
                return [TextContent(type="text", text=json.dumps({
                    "success": False,
                    "error": str(e),
                    "tool": name
                }, ensure_ascii=False))]

    async def run(self) -> None:
        """运行 MCP 服务器"""
        if not self.server or not stdio_server:
            logger.error("MCP Server not available (MCP SDK not installed)")
            return

        # 运行服务器 - stdio 模式
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name="wiki-mcp",
                    server_version="1.0.0",
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={}
                    )
                )
            )

    async def close(self) -> None:
        """关闭服务器"""
        if self.db:
            await self.db.close()
            logger.info("Database connection closed")


async def main():
    """主入口"""
    import os

    # 从环境变量读取配置
    wiki_vault_path = os.environ.get("WIKI_VAULT_PATH", DEFAULT_WIKI_VAULT)
    index_path = os.path.expanduser(os.environ.get("WIKI_INDEX_PATH", DEFAULT_INDEX_PATH))
    log_level = os.environ.get("LOG_LEVEL", DEFAULT_LOG_LEVEL)

    # 设置日志级别
    logging.getLogger().setLevel(getattr(logging, log_level))

    # 检查 wiki vault 路径
    if not Path(wiki_vault_path).exists():
        logger.error(f"Wiki vault path does not exist: {wiki_vault_path}")
        sys.exit(1)

    # 创建并运行服务器
    server = MCPServer(wiki_vault_path, index_path)

    try:
        await server.initialize()
        await server.run()
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    finally:
        await server.close()


if __name__ == "__main__":
    asyncio.run(main())