agent功能开发增加MCP后端

mcp_server.py

@@ -140,6 +140,7 @@ class AgentChatRequest(BaseModel):
     user_id: Optional[str] = None  # 用户ID
     user_nickname: Optional[str] = None  # 用户昵称
     user_avatar: Optional[str] = None  # 用户头像URL
+    subscription_type: Optional[str] = None  # 用户订阅类型（free/pro/max）
     session_id: Optional[str] = None  # 会话ID（如果为空则创建新会话）
 
 # ==================== MCP工具定义 ====================
@@ -1549,22 +1550,36 @@ async def agent_chat(request: AgentChatRequest):
     logger.info(f"Agent chat: {request.message} (user: {request.user_id})")
 
     # ==================== 权限检查 ====================
-    # 仅允许 max 用户使用（支持多种格式：字符串 "max"、数字 1、或 nickname 为 "max"）
-    is_max_user = (
-        request.user_id == "max" or
-        request.user_id == "1" or
-        request.user_id == 1 or
-        request.user_nickname == "max"
-    )
+    # 订阅等级判断函数（与 app.py 保持一致）
+    def subscription_level(sub_type: str) -> int:
+        """将订阅类型映射到等级数值，free=0, pro=1, max=2"""
+        mapping = {'free': 0, 'pro': 1, 'max': 2}
+        return mapping.get((sub_type or 'free').lower(), 0)
 
-    if not is_max_user:
-        logger.warning(f"权限检查失败 - user_id: {request.user_id}, nickname: {request.user_nickname}")
+    # 获取用户订阅类型（默认为 free）
+    user_subscription = (request.subscription_type or 'free').lower()
+    required_level = 'max'
+
+    # 权限检查：仅允许 max 用户访问（与传导链分析权限保持一致）
+    has_access = subscription_level(user_subscription) >= subscription_level(required_level)
+
+    if not has_access:
+        logger.warning(
+            f"权限检查失败 - user_id: {request.user_id}, "
+            f"nickname: {request.user_nickname}, "
+            f"subscription_type: {user_subscription}, "
+            f"required: {required_level}"
+        )
         raise HTTPException(
             status_code=403,
-            detail="很抱歉，「价小前投研」功能目前仅对特定用户开放。如需使用，请联系管理员。"
+            detail="很抱歉，「价小前投研」功能仅对 Max 订阅用户开放。请升级您的订阅以使用此功能。"
         )
 
-    logger.info(f"权限检查通过 - user_id: {request.user_id}, nickname: {request.user_nickname}")
+    logger.info(
+        f"权限检查通过 - user_id: {request.user_id}, "
+        f"nickname: {request.user_nickname}, "
+        f"subscription_type: {user_subscription}"
+    )
 
     # ==================== 会话管理 ====================
     # 如果没有提供 session_id，创建新会话
@@ -1644,6 +1659,38 @@ async def agent_chat_stream(request: AgentChatRequest):
     """智能代理对话端点（流式 SSE）"""
     logger.info(f"Agent chat stream: {request.message}")
 
+    # ==================== 权限检查 ====================
+    # 订阅等级判断函数（与 app.py 保持一致）
+    def subscription_level(sub_type: str) -> int:
+        """将订阅类型映射到等级数值，free=0, pro=1, max=2"""
+        mapping = {'free': 0, 'pro': 1, 'max': 2}
+        return mapping.get((sub_type or 'free').lower(), 0)
+
+    # 获取用户订阅类型（默认为 free）
+    user_subscription = (request.subscription_type or 'free').lower()
+    required_level = 'max'
+
+    # 权限检查：仅允许 max 用户访问（与传导链分析权限保持一致）
+    has_access = subscription_level(user_subscription) >= subscription_level(required_level)
+
+    if not has_access:
+        logger.warning(
+            f"[Stream] 权限检查失败 - user_id: {request.user_id}, "
+            f"nickname: {request.user_nickname}, "
+            f"subscription_type: {user_subscription}, "
+            f"required: {required_level}"
+        )
+        raise HTTPException(
+            status_code=403,
+            detail="很抱歉，「价小前投研」功能仅对 Max 订阅用户开放。请升级您的订阅以使用此功能。"
+        )
+
+    logger.info(
+        f"[Stream] 权限检查通过 - user_id: {request.user_id}, "
+        f"nickname: {request.user_nickname}, "
+        f"subscription_type: {user_subscription}"
+    )
+
     # 获取工具列表
     tools = [tool.dict() for tool in TOOLS]