仿生人MCP接入点
This commit is contained in:
BBIT-Kai
@@ -0,0 +1,258 @@
|
||||
"""
|
||||
MCP Endpoint Server
|
||||
主服务器文件
|
||||
"""
|
||||
|
||||
import json
|
||||
import signal
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from urllib.parse import quote
|
||||
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
from core.connection_manager import connection_manager
|
||||
from handlers.websocket_handler import websocket_handler
|
||||
from utils import __version__
|
||||
from utils.aes_utils import decrypt, encrypt
|
||||
from utils.config import config
|
||||
from utils.jsonrpc import (
|
||||
JSONRPCProtocol,
|
||||
)
|
||||
from utils.logger import get_logger
|
||||
from utils.util import get_local_ip
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
async def validate_token_and_get_agent_id(websocket: WebSocket) -> str:
|
||||
"""
|
||||
验证token并获取agentId的公共方法
|
||||
|
||||
Args:
|
||||
websocket: WebSocket连接对象
|
||||
|
||||
Returns:
|
||||
str: 验证成功返回agentId，失败返回None
|
||||
"""
|
||||
token = websocket.query_params.get("token")
|
||||
if not token:
|
||||
logger.error("缺少token参数")
|
||||
await websocket.close(code=1008, reason="缺少token参数")
|
||||
return None
|
||||
|
||||
data = decrypt(config.get("server", "key", ""), token)
|
||||
if not data:
|
||||
logger.error(f"token解密失败: {token}")
|
||||
await websocket.close(code=1008, reason="token解密失败")
|
||||
return None
|
||||
|
||||
try:
|
||||
data = json.loads(data)
|
||||
agent_id = data.get("agentId")
|
||||
if not agent_id:
|
||||
logger.error("无对应agentId")
|
||||
await websocket.close(code=1008, reason="无对应agentId")
|
||||
return None
|
||||
return agent_id
|
||||
except json.JSONDecodeError:
|
||||
logger.error("token数据格式错误")
|
||||
await websocket.close(code=1008, reason="token数据格式错误")
|
||||
return None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""应用生命周期管理"""
|
||||
# 启动时
|
||||
logger.info("MCP Endpoint Server 正在启动...")
|
||||
logger.info(f"=====下面的地址分别是智控台/单模块MCP接入点地址====")
|
||||
local_ip = get_local_ip()
|
||||
logger.info(
|
||||
f"智控台MCP参数配置: http://{local_ip}:{config.getint('server', 'port', 8004)}/mcp_endpoint/health?key={config.get('server', 'key', '')}"
|
||||
)
|
||||
encrypted_token = encrypt(
|
||||
config.get("server", "key", ""), '{"agentId":"single_module"}'
|
||||
)
|
||||
token = quote(encrypted_token)
|
||||
logger.info(
|
||||
f"单模块部署MCP接入点: ws://{local_ip}:{config.getint('server', 'port', 8004)}/mcp_endpoint/mcp/?token={token}"
|
||||
)
|
||||
logger.info(
|
||||
"=====请根据具体部署选择使用，请勿泄露给任何人======",
|
||||
)
|
||||
yield
|
||||
# 关闭时
|
||||
logger.info("MCP Endpoint Server 已关闭")
|
||||
|
||||
|
||||
# 创建FastAPI应用
|
||||
app = FastAPI(
|
||||
title="MCP Endpoint Server",
|
||||
description="高效的WebSocket中转服务器",
|
||||
version=__version__,
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# 配置CORS
|
||||
if config.getboolean("security", "enable_cors", True):
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=[config.get("security", "allowed_origins", "*")],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def redirect_root():
|
||||
"""根路径重定向到 /mcp_endpoint/"""
|
||||
return RedirectResponse(url="/mcp_endpoint/")
|
||||
|
||||
|
||||
@app.get("/mcp_endpoint/")
|
||||
async def root():
|
||||
"""根路径"""
|
||||
response = JSONRPCProtocol.create_success_response(
|
||||
result={
|
||||
"message": "MCP Endpoint Server",
|
||||
"version": __version__,
|
||||
"status": "running",
|
||||
}
|
||||
)
|
||||
return JSONRPCProtocol.to_dict(response)
|
||||
|
||||
|
||||
@app.get("/mcp_endpoint/health")
|
||||
async def health_check(key: str = None):
|
||||
"""健康检查"""
|
||||
# 验证key参数
|
||||
expected_key = config.get("server", "key", "")
|
||||
if not key or key != expected_key:
|
||||
response = JSONRPCProtocol.create_error_response(
|
||||
error_code=JSONRPCProtocol.AUTHENTICATION_ERROR,
|
||||
error_message="密钥验证失败",
|
||||
error_data={"details": "提供的密钥无效或缺失"},
|
||||
)
|
||||
return JSONRPCProtocol.to_dict(response)
|
||||
|
||||
stats = connection_manager.get_connection_stats()
|
||||
response = JSONRPCProtocol.create_success_response(
|
||||
result={"status": "success", "connections": stats}
|
||||
)
|
||||
return JSONRPCProtocol.to_dict(response)
|
||||
|
||||
|
||||
@app.websocket("/mcp_endpoint/mcp/")
|
||||
async def websocket_tool_endpoint(websocket: WebSocket):
|
||||
"""工具端WebSocket端点"""
|
||||
await websocket.accept()
|
||||
|
||||
# 获取agentId参数
|
||||
agent_id = await validate_token_and_get_agent_id(websocket)
|
||||
if not agent_id:
|
||||
return
|
||||
|
||||
try:
|
||||
# 注册连接
|
||||
await connection_manager.register_tool_connection(agent_id, websocket)
|
||||
logger.info(f"工具端连接已建立: {agent_id}")
|
||||
|
||||
# 处理消息
|
||||
while True:
|
||||
try:
|
||||
message = await websocket.receive_text()
|
||||
await websocket_handler._handle_tool_message(agent_id, message)
|
||||
except WebSocketDisconnect:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"处理工具端消息时发生错误: {e}")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理工具端连接时发生错误: {e}")
|
||||
finally:
|
||||
await connection_manager.unregister_tool_connection(agent_id)
|
||||
logger.info(f"工具端连接已关闭: {agent_id}")
|
||||
|
||||
|
||||
@app.websocket("/mcp_endpoint/call/")
|
||||
async def websocket_robot_endpoint(websocket: WebSocket):
|
||||
"""小智端WebSocket端点"""
|
||||
await websocket.accept()
|
||||
|
||||
# 获取agentId参数
|
||||
agent_id = await validate_token_and_get_agent_id(websocket)
|
||||
if not agent_id:
|
||||
return
|
||||
|
||||
try:
|
||||
# 注册连接并获取UUID
|
||||
connection_uuid = await connection_manager.register_robot_connection(
|
||||
agent_id, websocket
|
||||
)
|
||||
logger.info(f"小智端连接已建立: {agent_id} (UUID: {connection_uuid})")
|
||||
|
||||
# 处理消息
|
||||
while True:
|
||||
try:
|
||||
message = await websocket.receive_text()
|
||||
await websocket_handler._handle_robot_message(
|
||||
agent_id, message, connection_uuid
|
||||
)
|
||||
except WebSocketDisconnect:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"处理小智端消息时发生错误: {e}")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理小智端连接时发生错误: {e}")
|
||||
finally:
|
||||
await connection_manager.unregister_robot_connection(connection_uuid)
|
||||
logger.info(f"小智端连接已关闭: {agent_id} (UUID: {connection_uuid})")
|
||||
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
"""信号处理器"""
|
||||
logger.info(f"收到信号 {signum}，正在关闭服务器...")
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
# 设置uvicorn日志拦截
|
||||
from utils.logger import logger_manager
|
||||
|
||||
logger_manager.setup_uvicorn_logging()
|
||||
|
||||
# 注册信号处理器
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
|
||||
# 获取配置
|
||||
host = config.get("server", "host", "0.0.0.0")
|
||||
port = config.getint("server", "port", 8004)
|
||||
debug = config.getboolean("server", "debug", False)
|
||||
|
||||
logger.info(f"启动MCP Endpoint Server: {host}:{port}")
|
||||
|
||||
# 启动服务器
|
||||
uvicorn.run(
|
||||
"app:app",
|
||||
host=host,
|
||||
port=port,
|
||||
reload=debug,
|
||||
log_level=config.get("server", "log_level", "INFO").lower(),
|
||||
access_log=False,
|
||||
log_config=None, # 禁用uvicorn默认日志配置
|
||||
use_colors=True, # 启用颜色支持
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user