elicitation版本，需要客户端实现 elicitation_handler 支持

2025-07-29 16:38:23 +08:00
parent 4243814fb1
commit b2b1093c6d
1 changed files with 275 additions and 0 deletions
--- a/src/rustfs_s3_toolkit/mcp_server_elicitation.py
+++ b/src/rustfs_s3_toolkit/mcp_server_elicitation.py
@@ -0,0 +1,275 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+"""
+@file      : mcp_server.py
+@desc      : MCP 2.0 服务端，支持多 session，自动会话配置，兼容 Studio/streamable-http 客户端 #!User Elicitation: 需要客户端实现 elicitation_handler 支持
+@version   : 0.3.0
+@author    : lyzeng
+@email     : pylyzeng@gmail.com
+"""
+
+from typing import Dict, Optional, Any
+from fastmcp import FastMCP, Context
+from dataclasses import dataclass
+from rustfs_s3_toolkit.s3_client import S3StorageToolkit
+import os
+
+# ==== MCP 服务端实例 ====
+mcp = FastMCP(
+    name="RustFS-MCP",
+    version="0.3.0",
+    instructions="""
+RustFS 对象存储工具（S3 兼容，多 session 隔离）
+
+- 支持文件上传、下载、列举、删除、目录操作。
+- 推荐所有操作在请求参数或 headers 里包含唯一 session_id（或 uuid），用于自动检索/创建你的 S3 连接实例。
+- 首次调用同一 session_id 时会自动引导填写 S3 配置信息（endpoint_url, access_key_id, secret_access_key, bucket_name）。
+- 同一 session_id 后续调用自动复用配置，无需重复输入。
+- 支持 MCP Studio、streamable-http、LangChain、LlamaIndex 等多客户端自动发现与调用。
+"""
+)
+
+# ==== S3 配置结构体，每个 session 独立一份 ====
+@dataclass
+class RustFSMCPConfig:
+    endpoint_url: str
+    access_key_id: str
+    secret_access_key: str
+    bucket_name: str
+    region_name: str = "us-east-1"
+
+# ==== 配置缓存：session_id -> S3 配置 ====
+client_config_cache: Dict[str, RustFSMCPConfig] = {}
+
+def get_config_for_client(session_id: str) -> Optional[RustFSMCPConfig]:
+    """根据 session_id 获取 S3 配置"""
+    return client_config_cache.get(session_id)
+
+def save_config_for_client(session_id: str, config: RustFSMCPConfig):
+    """保存/更新 session_id 对应的 S3 配置"""
+    client_config_cache[session_id] = config
+
+# ==== 通用 session_id/uuid 提取工具，兼容 MCP/streamable-http 多客户端 ====
+def get_session_id(ctx: Context) -> str:
+    """
+    优先提取 ctx.session_id，兼容 headers/arguments/message 等参数名变体
+    若未提供则 fallback 为 default
+    """
+    # MCP Studio 等标准写法
+    if hasattr(ctx, "session_id") and ctx.session_id:
+        return str(ctx.session_id)
+    # headers/arguments 各类别名兼容
+    for key in ("uuid", "session_id", "sessionId"):
+        if hasattr(ctx, "headers") and ctx.headers and key in ctx.headers:
+            return str(ctx.headers[key])
+        if hasattr(ctx, "arguments") and ctx.arguments and key in ctx.arguments:
+            return str(ctx.arguments[key])
+    # JSON-RPC message body 参数
+    if hasattr(ctx, "message") and ctx.message:
+        for key in ("uuid", "session_id", "sessionId"):
+            if key in ctx.message:
+                return str(ctx.message[key])
+    # fallback: 默认会话
+    return "default"
+
+# ==== MCP 工具方法 ====
+
+@mcp.tool  
+async def rustfs_upload_file(file_path: str, object_key: str, ctx: Context) -> str:  
+    """  
+    上传文件到 S3 存储  
+      
+    Args:  
+        file_path: 本地文件路径  
+        object_key: S3 对象键名  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+  
+    # 首次使用时配置 S3 连接  
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+  
+    # 直接使用传入的参数，无需再次征询  
+    try:  
+        toolkit = S3StorageToolkit(**vars(config))  
+        check = toolkit.test_connection()  
+        if not check.get("success"):  
+            return f"连接失败: {check.get('message')}"  
+        upload_result = toolkit.upload_file(file_path, object_key)  
+        if upload_result.get("success"):  
+            return f"上传成功，访问URL: {upload_result.get('public_url')}"  
+        return f"上传失败: {upload_result.get('error', '')}"  
+    except Exception as e:  
+        return f"上传异常: {str(e)}"
+
+@mcp.tool  
+async def rustfs_download_file(object_key: str, local_path: str, ctx: Context) -> str:  
+    """  
+    下载文件  
+      
+    Args:  
+        object_key: 远端文件路径  
+        local_path: 本地保存路径  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+      
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+      
+    try:  
+        toolkit = S3StorageToolkit(**vars(config))  
+        dl = toolkit.download_file(object_key, local_path)  
+        if dl.get("success"):  
+            return f"下载成功，保存路径: {local_path}"  
+        return f"下载失败: {dl.get('error', '')}"  
+    except Exception as e:  
+        return f"下载异常: {str(e)}"
+
+@mcp.tool  
+async def rustfs_list_files(prefix: str, ctx: Context) -> Any:  
+    """  
+    列举 bucket 文件  
+      
+    Args:  
+        prefix: 文件前缀（使用""显示默认说有文件）  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+      
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+      
+    toolkit = S3StorageToolkit(**vars(config))  
+    result = toolkit.list_files(prefix=prefix)  
+    if not result.get("success"):  
+        return f"列举失败: {result.get('error', '')}"  
+    files = result.get("files", [])  
+    if not files:  
+        return "该前缀下没有文件"  
+    # 只展示前20个文件，避免输出过长  
+    return "\n".join(  
+        f"{f['key']} | {f['size']}B | {f['public_url']}" for f in files[:20]  
+    ) + ("\n...等" if len(files) > 20 else "")
+
+@mcp.tool  
+async def rustfs_delete_file(object_key: str, ctx: Context) -> str:  
+    """  
+    删除文件  
+      
+    Args:  
+        object_key: 要删除的文件键名  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+      
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+      
+    toolkit = S3StorageToolkit(**vars(config))  
+    result = toolkit.delete_file(object_key)  
+    if result.get("success"):  
+        return f"删除成功: {object_key}"  
+    return f"删除失败: {result.get('error', '')}"
+
+@mcp.tool  
+async def rustfs_create_folder(folder_path: str, ctx: Context) -> str:  
+    """  
+    创建目录  
+      
+    Args:  
+        folder_path: 要创建的文件夹路径（建议以 / 结尾）  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+      
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+      
+    toolkit = S3StorageToolkit(**vars(config))  
+    result = toolkit.create_folder(folder_path)  
+    if result.get("success"):  
+        return f"目录创建成功: {folder_path}"  
+    return f"创建失败: {result.get('error', '')}"
+
+@mcp.tool  
+async def rustfs_delete_directory(remote_prefix: str, ctx: Context) -> str:  
+    """  
+    删除目录  
+      
+    Args:  
+        remote_prefix: 要删除的目录前缀（如 foo/bar/）  
+    """  
+    session_id = get_session_id(ctx)  
+    config = get_config_for_client(session_id)  
+      
+    if not config:  
+        res = await ctx.elicit(  
+            "首次使用此 session，请填写 S3 连接参数",   
+            response_type=RustFSMCPConfig  
+        )  
+        if res.action != "accept":  
+            return "已取消"  
+        config = res.data  
+        save_config_for_client(session_id, config)  
+      
+    toolkit = S3StorageToolkit(**vars(config))  
+    result = toolkit.delete_directory(remote_prefix)  
+    if result.get("success"):  
+        return f"目录删除成功，共 {result.get('deleted_count', 0)} 项"  
+    return f"删除失败: {result.get('error', '')}"
+
+# ==== 推荐标准中间件（日志/限流/异常/性能）====
+from fastmcp.server.middleware.logging import LoggingMiddleware
+from fastmcp.server.middleware.timing import TimingMiddleware
+from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
+from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
+
+mcp.add_middleware(ErrorHandlingMiddleware())  # 全局异常兜底
+mcp.add_middleware(RateLimitingMiddleware(max_requests_per_second=30, burst_capacity=60))
+mcp.add_middleware(TimingMiddleware())
+mcp.add_middleware(LoggingMiddleware(include_payloads=True, max_payload_length=500))
+
+# ==== MCP 服务端启动 ====
+if __name__ == "__main__":
+    mcp.run(
+        transport=os.getenv("TRANSPORT", "streamable-http"),  # 支持 stdio/http/streamable-http
+        host=os.getenv("HOST", "0.0.0.0"),
+        port=int(os.getenv("PORT", 9009))
+    )