# server.py
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("Production MCP Server")

@mcp.tool()
def greet(name: str) -> str:
    """向用户打招呼"""
    return f"你好，{name}！欢迎使用生产级 MCP 服务器。"

@mcp.tool()
def add(a: int, b: int) -> int:
    """计算两个数字之和"""
    return a + b

# 启动入口
if __name__ == "__main__":
    mcp.run(transport="http")

# auth.py
import os
from functools import wraps
from fastapi import HTTPException, Security
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

# 从环境变量读取 API Keys（用逗号分隔支持多 Key）
VALID_API_KEYS = set(
    os.getenv("MCP_API_KEYS", "dev-key-123").split(",")
)

security_scheme = HTTPBearer(auto_error=False)

def verify_api_key(credentials: HTTPAuthorizationCredentials = Security(security_scheme)):
    """验证 API Key"""
    if not credentials:
        raise HTTPException(
            status_code=401,
            detail="缺少 Authorization 头，请提供 API Key"
        )

    token = credentials.credentials

    if token not in VALID_API_KEYS:
        raise HTTPException(
            status_code=403,
            detail="无效的 API Key"
        )

    return token

# server.py
from mcp.server.fastmcp import FastMCP
from auth import verify_api_key

mcp = FastMCP("Production MCP Server")

@mcp.tool()
def secure_greet(name: str) -> str:
    """需要 API Key 才能调用的工具"""
    return f"你好，{name}！"

# 修改启动入口，挂载鉴权 middleware
from fastapi import FastAPI
from starlette.middleware.base import BaseHTTPMiddleware

app = FastAPI()

# 鉴权 middleware
class AuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        # 健康检查和指标端点不需要鉴权
        if request.url.path in ["/health", "/metrics"]:
            return await call_next(request)

        # 验证 API Key
        auth_header = request.headers.get("Authorization")
        if not auth_header or not auth_header.startswith("Bearer "):
            from fastapi.responses import JSONResponse
            return JSONResponse(
                status_code=401,
                content={"error": "缺少 Authorization 头"}
            )

        token = auth_header.replace("Bearer ", "")
        if token not in VALID_API_KEYS:
            from fastapi.responses import JSONResponse
            return JSONResponse(
                status_code=403,
                content={"error": "无效的 API Key"}
            )

        return await call_next(request)

app.add_middleware(AuthMiddleware)

# 挂载 MCP SSE 端点
app.mount("/", mcp.sse_app())

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# jwt_auth.py
import os
import time
import jwt
from typing import Optional

# 从环境变量读取 JWT Secret
JWT_SECRET = os.getenv("MCP_JWT_SECRET", "change-me-in-production")
JWT_ALGORITHM = "HS256"

def create_token(user_id: str, role: str = "user", expiry_hours: int = 24) -> str:
    """生成 JWT Token"""
    payload = {
        "sub": user_id,
        "role": role,
        "iat": int(time.time()),
        "exp": int(time.time()) + expiry_hours * 3600,
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)

def verify_jwt(token: str) -> Optional[dict]:
    """验证 JWT Token，返回 payload"""
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

# tools.py
from jwt_auth import verify_jwt
from functools import wraps

def require_role(role: str):
    """装饰器：要求特定角色才能调用"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 从上下文中获取用户信息（需要注入 context）
            user_role = kwargs.get("_user_role", "anonymous")
            if user_role != role and user_role != "admin":
                raise PermissionError(f"需要 {role} 角色，当前为 {user_role}")
            return await func(*args, **kwargs)
        return wrapper
    return decorator

@mcp.tool()
@require_role("admin")
def admin_only_tool() -> str:
    """仅管理员可调用"""
    return "这是管理员专属工具"

# logger.py
import structlog
import logging
import sys

def setup_logging():
    """配置结构化日志"""
    structlog.configure(
        processors=[
            structlog.stdlib.filter_by_level,
            structlog.stdlib.add_logger_name,
            structlog.stdlib.add_log_level,
            structlog.stdlib.PositionalArgumentsFormatter(),
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.StackInfoRenderer(),
            structlog.processors.format_exc_info,
            structlog.processors.UnicodeDecoder(),
            # JSON 输出，适合生产环境
            structlog.processors.JSONRenderer()
        ],
        context_class=dict,
        logger_factory=structlog.stdlib.LoggerFactory(),
        cache_logger_on_first_use=True,
    )

    # 设置 root logger
    root_logger = logging.getLogger()
    handler = logging.StreamHandler(sys.stdout)
    root_logger.addHandler(handler)
    root_logger.setLevel(logging.INFO)

    return structlog.get_logger()

# 使用
logger = setup_logging()
logger.info("mcp_server_started", port=8000, workers=4, transport="http")

# server.py
from logger import logger

@mcp.tool()
def greet(name: str) -> str:
    logger.info("greet_called", name=name, source_ip=request.client.host)
    return f"你好，{name}！"

# metrics.py
from prometheus_client import Counter, Histogram, Gauge, generate_latest, CONTENT_TYPE_LATEST
from starlette.responses import Response
import time

# 定义指标
TOOL_CALLS_TOTAL = Counter(
    "mcp_tool_calls_total",
    "MCP 工具调用总数",
    ["tool_name", "status"]  # label: 工具名 + 成功/失败
)

TOOL_CALL_DURATION = Histogram(
    "mcp_tool_call_duration_seconds",
    "MCP 工具调用耗时（秒）",
    ["tool_name"],
    buckets=(0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0)
)

ACTIVE_CONNECTIONS = Gauge(
    "mcp_active_connections",
    "当前活跃的 SSE 连接数"
)

ACTIVE_TOOLS = Gauge(
    "mcp_registered_tools_total",
    "注册的工具总数"
)

def metrics_endpoint(request):
    """Prometheus metrics 端点"""
    return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)

def track_tool_metrics(tool_name: str):
    """工具调用耗时追踪装饰器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            start = time.time()
            try:
                result = func(*args, **kwargs)
                TOOL_CALLS_TOTAL.labels(tool_name=tool_name, status="success").inc()
                return result
            except Exception as e:
                TOOL_CALLS_TOTAL.labels(tool_name=tool_name, status="error").inc()
                raise
            finally:
                duration = time.time() - start
                TOOL_CALL_DURATION.labels(tool_name=tool_name).observe(duration)
        return wrapper
    return decorator

# server.py
from metrics import metrics_endpoint, track_tool_metrics, ACTIVE_TOOLS

# 在启动时登记工具数量
ACTIVE_TOOLS.set(len(mcp._tool_manager.list_tools()))

# 在 ASGI app 中挂载 metrics
app.mount("/metrics", metrics_endpoint)

@mcp.tool()
@track_tool_metrics("greet")
def greet(name: str) -> str:
    return f"你好，{name}！"

# tenant.py
import os
import threading
from contextvars import ContextVar

# 使用 ContextVar 实现线程/协程级租户隔离
current_tenant: ContextVar[str] = ContextVar("current_tenant", default="default")

class TenantRouter:
    """根据租户 ID 路由到不同的资源配置"""

    def __init__(self):
        self._tenants = {}  # tenant_id -> config

    def register_tenant(self, tenant_id: str, config: dict):
        self._tenants[tenant_id] = config

    def get_config(self, key: str, default=None):
        tenant = current_tenant.get()
        config = self._tenants.get(tenant, {})
        return config.get(key, default)

# 全局路由
tenant_router = TenantRouter()

# 初始化租户
tenant_router.register_tenant("acme-corp", {
    "database_url": "postgresql://acme:pass@db:5432/acme",
    "api_key": "ak-acme-secret",
    "rate_limit": 100,   # 每分钟允许的请求数
})
tenant_router.register_tenant("startup-inc", {
    "database_url": "postgresql://startup:pass@db:5432/startup",
    "api_key": "ak-startup-secret",
    "rate_limit": 20,
})

@mcp.tool()
def query_tenant_data(query: str) -> str:
    """查询当前租户的数据"""
    db_url = tenant_router.get_config("database_url")
    # 连接到该租户的数据库
    # ... 执行查询
    return f"查询完成（租户：{current_tenant.get()}）"

# tenant_manager.py
import subprocess
import os
import signal

class TenantProcessManager:
    """管理每个租户的独立 MCP 服务器进程"""

    def __init__(self, base_port: int = 9000):
        self.base_port = base_port
        self._processes: dict[str, subprocess.Popen] = {}
        self._ports: dict[str, int] = {}

    def start_tenant(self, tenant_id: str, config: dict):
        """为租户启动一个独立的 MCP 服务器进程"""
        if tenant_id in self._processes:
            return self._ports[tenant_id]

        port = self.base_port + len(self._processes)
        env = os.environ.copy()
        env.update({
            "MCP_TENANT_ID": tenant_id,
            "MCP_DATABASE_URL": config["database_url"],
            "MCP_API_KEYS": config["api_key"],
            "MCP_PORT": str(port),
        })

        proc = subprocess.Popen(
            ["uvicorn", "server:mcp.sse_app",
             "--host", "0.0.0.0",
             "--port", str(port),
             "--workers", "2"],
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        self._processes[tenant_id] = proc
        self._ports[tenant_id] = port
        return port

    def stop_tenant(self, tenant_id: str):
        """停止租户进程"""
        if tenant_id in self._processes:
            self._processes[tenant_id].terminate()
            self._processes[tenant_id].wait()
            del self._processes[tenant_id]
            del self._ports[tenant_id]

    def stop_all(self):
        """停止所有租户进程"""
        for tenant_id in list(self._processes.keys()):
            self.stop_tenant(tenant_id)

# connection_pool.py
import aiohttp
import asyncio
from typing import Optional

class ConnectionPool:
    """全局 HTTP 连接池"""

    _instance: Optional["ConnectionPool"] = None
    _session: Optional[aiohttp.ClientSession] = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    async def get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            connector = aiohttp.TCPConnector(
                limit=100,           # 最大并发连接数
                limit_per_host=20,   # 每主机最大连接数
                ttl_dns_cache=300,   # DNS 缓存 5 分钟
                enable_cleanup_closed=True,
            )
            timeout = aiohttp.ClientTimeout(
                total=30,           # 总超时
                connect=5,          # 连接超时
                sock_read=30,       # 读取超时
            )
            self._session = aiohttp.ClientSession(
                connector=connector,
                timeout=timeout,
            )
        return self._session

    async def close(self):
        if self._session and not self._session.closed:
            await self._session.close()

# 应用关闭时清理
pool = ConnectionPool()
import atexit
atexit.register(lambda: asyncio.run(pool.close()))

# rate_limiter.py
import time
import asyncio
from collections import defaultdict

class TokenBucket:
    """令牌桶限流器"""

    def __init__(self, rate: float, capacity: int):
        """
        rate: 每秒新增令牌数
        capacity: 桶容量（最大突发）
        """
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.last_refill = time.monotonic()

    def consume(self, tokens: int = 1) -> bool:
        """消费令牌，返回是否允许通过"""
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
        self.last_refill = now

        if self.tokens >= tokens:
            self.tokens -= tokens
            return True
        return False

class RateLimiter:
    """多租户限流器"""

    def __init__(self, default_rate: float = 10, default_capacity: int = 20):
        self.default_rate = default_rate
        self.default_capacity = default_capacity
        self._buckets: dict[str, TokenBucket] = {}

    def check(self, key: str, rate: Optional[float] = None, capacity: Optional[int] = None) -> bool:
        """检查是否限流"""
        if key not in self._buckets:
            self._buckets[key] = TokenBucket(
                rate or self.default_rate,
                capacity or self.default_capacity
            )
        return self._buckets[key].consume()

    def get_wait_time(self, key: str) -> float:
        """获取需要等待的秒数"""
        bucket = self._buckets.get(key)
        if not bucket or bucket.tokens > 0:
            return 0
        return (1 - bucket.tokens / bucket.capacity) / bucket.rate

# 全局限流器
rate_limiter = RateLimiter()

@mcp.tool()
def rate_limited_tool(name: str) -> str:
    """受限流保护的工具"""
    tenant = current_tenant.get()
    if not rate_limiter.check(tenant):
        wait = rate_limiter.get_wait_time(tenant)
        raise RateLimitError(retry_after=int(wait))
    return f"你好，{name}！"

# timeout.py
import asyncio
from functools import wraps

def with_timeout(seconds: float):
    """带超时的工具装饰器"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            try:
                return await asyncio.wait_for(
                    func(*args, **kwargs),
                    timeout=seconds
                )
            except asyncio.TimeoutError:
                raise TimeoutError(f"工具执行超时（{seconds}秒）")
        return wrapper
    return decorator

# 使用
@mcp.tool()
@with_timeout(30.0)
async def slow_data_fetch(query: str) -> str:
    """可能耗时较长的数据查询"""
    await asyncio.sleep(1)  # 模拟耗时操作
    return f"查询结果：{query}"

# locustfile.py
from locust import HttpUser, task, between
import json

class MCPUser(HttpUser):
    wait_time = between(0.5, 2)

    def on_start(self):
        """每个模拟用户先建立 SSE 连接"""
        self.client.headers = {
            "Authorization": "Bearer sk-test-key",
            "Content-Type": "application/json"
        }

    @task(3)
    def list_tools(self):
        payload = {
            "jsonrpc": "2.0",
            "method": "tools/list",
            "params": {},
            "id": 1
        }
        with self.client.post(
            "/mcp/message",
            json=payload,
            catch_response=True
        ) as response:
            if response.status_code != 200:
                response.failure(f"Status: {response.status_code}")

    @task(7)
    def call_tool(self):
        payload = {
            "jsonrpc": "2.0",
            "method": "tools/call",
            "params": {
                "name": "greet",
                "arguments": {"name": "test"}
            },
            "id": 2
        }
        with self.client.post(
            "/mcp/message",
            json=payload,
            catch_response=True
        ) as response:
            if response.status_code != 200:
                response.failure(f"Status: {response.status_code}")

# 客户端自动重连示例
import asyncio
import sseclient

async def connect_with_retry(url: str, max_retries: int = 5):
    for attempt in range(max_retries):
        try:
            response = requests.get(url, stream=True)
            client = sseclient.SSEClient(response)
            for event in client.events():
                process_event(event)
            break
        except (ConnectionError, requests.RequestException) as e:
            wait = 2 ** attempt  # 指数退避
            print(f"连接断开，{wait}s 后重试（{attempt+1}/{max_retries}）")
            await asyncio.sleep(wait)