344 lines
12 KiB
Python
344 lines
12 KiB
Python
"""
|
||
Backend-integrated MCP Streamable HTTP server.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from datetime import datetime
|
||
import hmac
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List
|
||
import uuid
|
||
|
||
from fastapi import HTTPException, Response
|
||
from sqlalchemy import select
|
||
|
||
try:
|
||
from mcp.server.fastmcp import FastMCP
|
||
except ImportError: # pragma: no cover - runtime dependency
|
||
FastMCP = None
|
||
|
||
from app.core.database import AsyncSessionLocal
|
||
from app.models.mcp_bot import MCPBot
|
||
from app.models.project import Project, ProjectMember
|
||
from app.models.user import User
|
||
from app.schemas.project import ProjectResponse
|
||
from app.services.project_file_service import project_file_service
|
||
from app.services.storage import storage_service
|
||
from app.services.project_service import (
|
||
count_project_documents,
|
||
require_project_read_access,
|
||
require_project_write_access,
|
||
)
|
||
from app.core.config import settings
|
||
from app.mcp.context import MCPRequestContext, current_mcp_request
|
||
|
||
|
||
mcp = (
|
||
FastMCP(
|
||
"NexDocs MCP",
|
||
host=settings.HOST,
|
||
port=settings.PORT,
|
||
stateless_http=True,
|
||
json_response=True,
|
||
streamable_http_path="/",
|
||
)
|
||
if FastMCP
|
||
else None
|
||
)
|
||
|
||
|
||
async def _get_current_user(db) -> User:
|
||
ctx = current_mcp_request.get()
|
||
if ctx is None:
|
||
raise RuntimeError("MCP request context is missing.")
|
||
|
||
result = await db.execute(select(User).where(User.id == ctx.user_id, User.status == 1))
|
||
user = result.scalar_one_or_none()
|
||
if not user:
|
||
raise RuntimeError("Authenticated MCP user does not exist or is disabled.")
|
||
return user
|
||
|
||
|
||
async def _get_project_with_write_access(project_id: int, current_user: User, db):
|
||
project, _ = await require_project_write_access(db, project_id, current_user)
|
||
return project
|
||
|
||
|
||
def _ensure_file_exists(file_path: Path, path: str) -> None:
|
||
if not file_path.exists():
|
||
raise HTTPException(status_code=404, detail=f"文件不存在: {path}")
|
||
if not file_path.is_file():
|
||
raise HTTPException(status_code=400, detail=f"目标不是文件: {path}")
|
||
|
||
|
||
def _ensure_file_not_exists(file_path: Path, path: str) -> None:
|
||
if file_path.exists():
|
||
raise HTTPException(status_code=400, detail=f"文件已存在: {path}")
|
||
|
||
if mcp is not None:
|
||
@mcp.tool(name="list_created_projects", description="获取当前用户创建的项目列表。")
|
||
async def list_created_projects(keyword: str = "", limit: int = 100) -> List[Dict[str, Any]]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
result = await db.execute(
|
||
select(Project).where(Project.owner_id == current_user.id, Project.status == 1)
|
||
)
|
||
projects = result.scalars().all()
|
||
|
||
items = []
|
||
keyword_lower = keyword.strip().lower()
|
||
for project in projects:
|
||
project_dict = ProjectResponse.from_orm(project).dict()
|
||
project_dict["doc_count"] = count_project_documents(project.storage_key)
|
||
if keyword_lower:
|
||
haystack = f"{project.name} {project.description or ''}".lower()
|
||
if keyword_lower not in haystack:
|
||
continue
|
||
items.append(project_dict)
|
||
|
||
return items[: max(limit, 0)]
|
||
|
||
|
||
@mcp.tool(name="create_project", description="创建新项目（按项目名称创建，创建者自动成为项目管理员）。")
|
||
async def create_project(name: str, description: str = "") -> Dict[str, Any]:
|
||
"""按项目名称创建新项目，返回项目 ID 与存储标识。"""
|
||
name = (name or "").strip()
|
||
if not name:
|
||
raise HTTPException(status_code=400, detail="项目名称不能为空")
|
||
if len(name) > 100:
|
||
raise HTTPException(status_code=400, detail="项目名称不能超过 100 个字符")
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
|
||
# 生成 UUID 作为存储键
|
||
storage_key = str(uuid.uuid4())
|
||
db_project = Project(
|
||
name=name,
|
||
description=(description or "").strip() or None,
|
||
storage_key=storage_key,
|
||
owner_id=current_user.id,
|
||
is_public=0,
|
||
status=1,
|
||
)
|
||
db.add(db_project)
|
||
await db.commit()
|
||
await db.refresh(db_project)
|
||
|
||
# 创建物理文件夹结构，失败则回滚数据库记录
|
||
try:
|
||
storage_service.create_project_structure(storage_key)
|
||
except Exception as exc: # noqa: BLE001
|
||
await db.delete(db_project)
|
||
await db.commit()
|
||
raise HTTPException(status_code=500, detail=f"项目文件夹创建失败: {exc}")
|
||
|
||
# 项目创建者自动成为管理员成员
|
||
db_member = ProjectMember(
|
||
project_id=db_project.id,
|
||
user_id=current_user.id,
|
||
role="admin",
|
||
)
|
||
db.add(db_member)
|
||
await db.commit()
|
||
|
||
# 记录操作日志（MCP 无 HTTP 请求对象，跳过 request 字段）
|
||
try:
|
||
from app.core.enums import OperationType
|
||
from app.services.log_service import log_service
|
||
|
||
await log_service.log_project_operation(
|
||
db=db,
|
||
operation_type=OperationType.CREATE_PROJECT,
|
||
project_id=db_project.id,
|
||
user=current_user,
|
||
detail={"project_name": name, "source": "mcp"},
|
||
)
|
||
except Exception as exc: # noqa: BLE001
|
||
pass
|
||
|
||
return {
|
||
"message": "项目创建成功",
|
||
"project_id": db_project.id,
|
||
"name": db_project.name,
|
||
"storage_key": db_project.storage_key,
|
||
}
|
||
|
||
|
||
@mcp.tool(name="get_project_tree", description="获取指定项目的目录树。")
|
||
async def get_project_tree(project_id: int) -> Dict[str, Any]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
project, user_role = await require_project_read_access(db, project_id, current_user)
|
||
project_root = storage_service.get_secure_path(project.storage_key)
|
||
tree = storage_service.generate_tree(project_root)
|
||
|
||
return {
|
||
"tree": [item.model_dump() for item in tree],
|
||
"user_role": user_role,
|
||
"project_name": project.name,
|
||
"project_description": project.description,
|
||
}
|
||
|
||
|
||
@mcp.tool(name="get_file", description="读取指定项目中的文件内容。")
|
||
async def get_file(project_id: int, path: str) -> Dict[str, Any]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
project, _ = await require_project_read_access(db, project_id, current_user)
|
||
file_path = storage_service.get_secure_path(project.storage_key, path)
|
||
_ensure_file_exists(file_path, path)
|
||
content = await storage_service.read_file(file_path)
|
||
return {"path": path, "content": content}
|
||
|
||
|
||
@mcp.tool(name="create_file", description="在指定项目的路径下创建新文件。")
|
||
async def create_file(project_id: int, path: str, content: str = "") -> Dict[str, Any]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
project = await _get_project_with_write_access(project_id, current_user, db)
|
||
file_path = storage_service.get_secure_path(project.storage_key, path)
|
||
_ensure_file_not_exists(file_path, path)
|
||
await project_file_service.operate_file(
|
||
db,
|
||
project_id,
|
||
project,
|
||
"create_file",
|
||
path,
|
||
current_user,
|
||
content=content,
|
||
source="mcp",
|
||
)
|
||
|
||
return {
|
||
"message": "文件创建成功",
|
||
"project_id": project_id,
|
||
"path": path,
|
||
}
|
||
|
||
|
||
@mcp.tool(name="update_file", description="更新指定项目中已有文件的内容。")
|
||
async def update_file(project_id: int, path: str, content: str) -> Dict[str, Any]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
project = await _get_project_with_write_access(project_id, current_user, db)
|
||
file_path = storage_service.get_secure_path(project.storage_key, path)
|
||
_ensure_file_exists(file_path, path)
|
||
await project_file_service.save_file(
|
||
db,
|
||
project_id,
|
||
project,
|
||
path,
|
||
content,
|
||
current_user,
|
||
source="mcp",
|
||
)
|
||
|
||
return {
|
||
"message": "文件更新成功",
|
||
"project_id": project_id,
|
||
"path": path,
|
||
}
|
||
|
||
|
||
@mcp.tool(name="delete_file", description="删除指定项目中的文件。")
|
||
async def delete_file(project_id: int, path: str) -> Dict[str, Any]:
|
||
async with AsyncSessionLocal() as db:
|
||
current_user = await _get_current_user(db)
|
||
project = await _get_project_with_write_access(project_id, current_user, db)
|
||
file_path = storage_service.get_secure_path(project.storage_key, path)
|
||
_ensure_file_exists(file_path, path)
|
||
await project_file_service.operate_file(
|
||
db,
|
||
project_id,
|
||
project,
|
||
"delete",
|
||
path,
|
||
current_user,
|
||
source="mcp",
|
||
)
|
||
|
||
return {
|
||
"message": "文件删除成功",
|
||
"project_id": project_id,
|
||
"path": path,
|
||
}
|
||
|
||
|
||
def create_mcp_http_app():
|
||
"""Return the MCP streamable HTTP ASGI app."""
|
||
if mcp is None:
|
||
raise RuntimeError("Package 'mcp' is required to run the MCP endpoint.")
|
||
return mcp.streamable_http_app()
|
||
|
||
|
||
def get_mcp_session_manager():
|
||
"""Return the MCP streamable HTTP session manager."""
|
||
if mcp is None:
|
||
raise RuntimeError("Package 'mcp' is required to run the MCP endpoint.")
|
||
return mcp.session_manager
|
||
|
||
|
||
class MCPHeaderAuthApp:
|
||
"""ASGI wrapper that authenticates incoming MCP requests via bot headers."""
|
||
|
||
def __init__(self, app):
|
||
self.app = app
|
||
|
||
async def __call__(self, scope, receive, send):
|
||
if scope["type"] != "http":
|
||
await self.app(scope, receive, send)
|
||
return
|
||
|
||
headers = {
|
||
key.decode("latin-1").lower(): value.decode("latin-1")
|
||
for key, value in scope.get("headers", [])
|
||
}
|
||
bot_id = headers.get("x-bot-id", "").strip()
|
||
bot_secret = headers.get("x-bot-secret", "").strip()
|
||
|
||
if not bot_id or not bot_secret:
|
||
response = Response(
|
||
content='{"error":"Missing X-Bot-Id or X-Bot-Secret"}',
|
||
status_code=401,
|
||
media_type="application/json",
|
||
)
|
||
await response(scope, receive, send)
|
||
return
|
||
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(MCPBot, User)
|
||
.join(User, User.id == MCPBot.user_id)
|
||
.where(MCPBot.bot_id == bot_id, MCPBot.status == 1, User.status == 1)
|
||
)
|
||
row = result.first()
|
||
|
||
if not row:
|
||
response = Response(
|
||
content='{"error":"Invalid MCP bot"}',
|
||
status_code=403,
|
||
media_type="application/json",
|
||
)
|
||
await response(scope, receive, send)
|
||
return
|
||
|
||
mcp_bot, user = row
|
||
if not hmac.compare_digest(mcp_bot.bot_secret, bot_secret):
|
||
response = Response(
|
||
content='{"error":"Invalid MCP secret"}',
|
||
status_code=403,
|
||
media_type="application/json",
|
||
)
|
||
await response(scope, receive, send)
|
||
return
|
||
|
||
mcp_bot.last_used_at = datetime.utcnow()
|
||
await db.commit()
|
||
|
||
token = current_mcp_request.set(MCPRequestContext(bot_id=bot_id, user_id=user.id))
|
||
try:
|
||
await self.app(scope, receive, send)
|
||
finally:
|
||
current_mcp_request.reset(token)
|