WeKnora/mcp-server/main.py

#!/usr/bin/env python3
"""
WeKnora MCP Server 主入口点

这个文件提供了一个统一的入口点来启动 WeKnora MCP 服务器。
可以通过多种方式运行：
1. python main.py
2. python -m weknora_mcp_server
3. weknora-mcp-server (安装后)

注意：在 stdio 传输下，stdout 是 JSON-RPC 通道，所有诊断/提示信息必须写入
stderr，否则会破坏 MCP 协议流导致客户端判定"启动失败"。本文件所有 print
均通过 stderr 输出。
"""

import argparse
import asyncio
import os
import sys
from pathlib import Path


def setup_environment():
    """设置环境和路径"""
    # 确保当前目录在 Python 路径中
    current_dir = Path(__file__).parent.absolute()
    if str(current_dir) not in sys.path:
        sys.path.insert(0, str(current_dir))


def check_dependencies():
    """检查依赖是否已安装"""
    try:
        import mcp
        import requests

        return True
    except ImportError as e:
        print(f"缺少依赖: {e}", file=sys.stderr)
        print("请运行: pip install -r requirements.txt", file=sys.stderr)
        return False


def check_environment_variables():
    """检查环境变量配置"""
    base_url = os.getenv("WEKNORA_BASE_URL")
    api_key = os.getenv("WEKNORA_API_KEY")

    print("=== WeKnora MCP Server 环境检查 ===", file=sys.stderr)
    print(f"Base URL: {base_url or 'http://localhost:8080/api/v1 (默认)'}", file=sys.stderr)
    print(f"API Key: {'已设置' if api_key else '未设置 (警告)'}", file=sys.stderr)

    if not base_url:
        print("提示: 可以设置 WEKNORA_BASE_URL 环境变量", file=sys.stderr)

    if not api_key:
        print("警告: 建议设置 WEKNORA_API_KEY 环境变量", file=sys.stderr)

    print("=" * 40, file=sys.stderr)
    return True


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="WeKnora MCP Server - Model Context Protocol server for WeKnora API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python main.py                    # 使用默认配置启动
  python main.py --check-only       # 仅检查环境，不启动服务器
  python main.py --verbose          # 启用详细日志
  
环境变量:
  WEKNORA_BASE_URL       WeKnora API 基础 URL (默认: http://localhost:8080/api/v1)
  WEKNORA_API_KEY        WeKnora API 密钥
  MCP_SERVER_AUTH_TOKEN  SSE/HTTP 传输必填，客户端通过 Authorization: Bearer 传递
        """,
    )

    parser.add_argument(
        "--check-only", action="store_true", help="仅检查环境配置，不启动服务器"
    )

    parser.add_argument("--verbose", "-v", action="store_true", help="启用详细日志输出")

    parser.add_argument(
        "--version", action="version", version="WeKnora MCP Server 1.1.0"
    )

    parser.add_argument(
        "--transport",
        choices=["stdio", "sse", "http"],
        default=os.getenv("MCP_TRANSPORT", "stdio"),
        help="Transport type: stdio (default), sse, or http",
    )
    parser.add_argument(
        "--host",
        default=os.getenv("MCP_HOST", "127.0.0.1"),
        help="Bind host for network transports (default: 127.0.0.1)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("MCP_PORT", "8000")),
        help="Bind port for network transports (default: 8000)",
    )

    return parser.parse_args()


async def main():
    """主函数"""
    args = parse_arguments()

    # 设置环境
    setup_environment()

    # 检查依赖
    if not check_dependencies():
        sys.exit(1)

    # 检查环境变量
    check_environment_variables()

    # 如果只是检查环境，则退出
    if args.check_only:
        print("环境检查完成。", file=sys.stderr)
        return

    # 设置日志级别
    if args.verbose:
        import logging

        logging.basicConfig(level=logging.DEBUG)
        print("已启用详细日志模式", file=sys.stderr)

    try:
        print(f"正在启动 WeKnora MCP Server (transport={args.transport})...", file=sys.stderr)

        from weknora_mcp_server import run_stdio, run_sse, run_http

        # Select transport mode based on CLI argument or MCP_TRANSPORT env var
        # - stdio: Default, used by VS Code Copilot for local integration
        # - sse: Server-Sent Events over HTTP, suitable for cloud/remote deployments
        # - http: Streamable HTTP sessions (MCP 2025-03-26 spec), compatible with REST clients
        if args.transport == "stdio":
            # Stdio mode: communication via stdin/stdout pipes (typical for CLI integrations)
            await run_stdio()
        elif args.transport == "sse":
            # SSE mode: HTTP server with Server-Sent Events for bidirectional streaming
            await run_sse(args.host, args.port)
        elif args.transport == "http":
            # HTTP mode: HTTP REST server with request/response model
            await run_http(args.host, args.port)

    except ImportError as e:
        print(f"导入错误: {e}", file=sys.stderr)
        print("请确保所有文件都在正确的位置", file=sys.stderr)
        sys.exit(1)
    except KeyboardInterrupt:
        print("\n服务器已停止", file=sys.stderr)
    except Exception as e:
        print(f"服务器运行错误: {e}", file=sys.stderr)
        if args.verbose:
            import traceback

            traceback.print_exc()
        sys.exit(1)


def sync_main():
    """同步版本的主函数，用于 entry_points"""
    asyncio.run(main())


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