mcp-softBV

mcp/main.py

@@ -10,17 +10,18 @@ from starlette.routing import Mount
 
 from system_tools import create_system_mcp
 from materialproject_mcp import create_materials_mcp
-
+from softBV import create_softbv_mcp
 # 创建 MCP 实例
 system_mcp = create_system_mcp()
 materials_mcp = create_materials_mcp()
-
+softbv_mcp = create_softbv_mcp()
 # 在 Starlette 的 lifespan 中启动 MCP 的 session manager
 @contextlib.asynccontextmanager
 async def lifespan(app: Starlette):
     async with contextlib.AsyncExitStack() as stack:
         await stack.enter_async_context(system_mcp.session_manager.run())
         await stack.enter_async_context(materials_mcp.session_manager.run())
+        await stack.enter_async_context(softbv_mcp.session_manager.run())
         yield  # 服务器运行期间
         # 退出时自动清理
 
@@ -30,6 +31,7 @@ app = Starlette(
     routes=[
         Mount("/system", app=system_mcp.streamable_http_app()),
         Mount("/materials", app=materials_mcp.streamable_http_app()),
+        Mount("/softBV", app=softbv_mcp.streamable_http_app()),
     ],
 )
 
@@ -38,5 +40,5 @@ app = Starlette(
 # 访问：
 #   http://localhost:8000/system
 #   http://localhost:8000/materials
-#
+#   http://localhost:8000/softBV
 # 如果需要浏览器客户端访问（CORS 暴露 Mcp-Session-Id），请参考 README 中的 CORS 配置示例 [1]

mcp/softBV.py

@@ -0,0 +1,464 @@
+# softbv_mcp.py
+#
+# 在远程服务器上激活 softBV 环境并执行计算（支持 --md 与 --gen-cube 两个专用工具）。
+# - 生命周期：建立 SSH 连接并注入上下文
+# - 激活环境：source /cluster/home/koko125/script/softBV.sh
+# - 工作目录：/cluster/home/koko125/sandbox
+# - 可执行文件：/cluster/home/koko125/tool/softBV-GUI_linux/bin/softBV.x
+#
+# 依赖：
+#   pip install "mcp[cli]" asyncssh pydantic
+#
+# 用法（Starlette 挂载示例见你现有 main.py，导入 create_softbv_mcp 即可）：
+#   from softbv_mcp import create_softbv_mcp
+#   softbv_mcp = create_softbv_mcp()
+#   Mount("/softbv", app=softbv_mcp.streamable_http_app())
+#
+# 可通过环境变量覆盖连接信息：
+#   REMOTE_HOST, REMOTE_USER, PRIVATE_KEY_PATH, SOFTBV_PROFILE, SOFTBV_BIN, DEFAULT_WORKDIR
+
+import os
+import posixpath
+import asyncio
+from dataclasses import dataclass
+from socket import socket
+from typing import Any
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import asyncssh
+from pydantic import BaseModel, Field
+
+from mcp.server.fastmcp import FastMCP, Context
+from mcp.server.session import ServerSession
+
+from lifespan import REMOTE_HOST, REMOTE_USER, PRIVATE_KEY_PATH
+
+def shell_quote(arg: str) -> str:
+    """安全地把字符串作为单个 shell 参数（POSIX）。"""
+    return "'" + str(arg).replace("'", "'\"'\"'") + "'"
+
+# 如果你已经定义了这些常量与路径，可保持复用；也可按需改为你自己的配置源
+
+
+# 固定的 softBV 环境信息（可用环境变量覆盖）
+SOFTBV_PROFILE = os.getenv("SOFTBV_PROFILE", "/cluster/home/koko125/script/softBV.sh")
+SOFTBV_BIN = os.getenv("SOFTBV_BIN", "/cluster/home/koko125/tool/softBV-GUI_linux/bin/softBV.x")
+DEFAULT_WORKDIR = os.getenv("DEFAULT_WORKDIR", "/cluster/home/koko125/sandbox")
+
+@dataclass
+class SoftBVContext:
+    ssh_connection: asyncssh.SSHClientConnection
+    workdir: str
+    profile: str
+    bin_path: str
+
+@asynccontextmanager
+async def softbv_lifespan(_server) -> AsyncIterator[SoftBVContext]:
+    """
+    FastMCP 生命周期：建立 SSH 连接并注入 softBV 上下文。
+    - 不再做额外的 DNS 解析或自定义异步步骤，避免 socket.SOCK_STREAM 的环境覆盖问题
+    - 仅负责连接与清理；工具中通过 ctx.request_context.lifespan_context 访问该上下文 [1]
+    """
+    # 允许用环境变量覆盖连接信息
+    host = os.getenv("REMOTE_HOST", REMOTE_HOST)
+    user = os.getenv("REMOTE_USER", REMOTE_USER)
+    key_path = os.getenv("PRIVATE_KEY_PATH", PRIVATE_KEY_PATH)
+
+    conn: asyncssh.SSHClientConnection | None = None
+    try:
+        conn = await asyncssh.connect(
+            host,
+            username=user,
+            client_keys=[key_path],
+            known_hosts=None,      # 如需主机指纹校验，可移除此参数
+            connect_timeout=15,    # 避免长时间挂起
+        )
+        yield SoftBVContext(
+            ssh_connection=conn,
+            workdir=DEFAULT_WORKDIR,
+            profile=SOFTBV_PROFILE,
+            bin_path=SOFTBV_BIN,
+        )
+    finally:
+        if conn:
+            conn.close()
+            await conn.wait_closed()
+async def run_in_softbv_env(
+    conn: asyncssh.SSHClientConnection,
+    profile_path: str,
+    cmd: str,
+    cwd: str | None = None,
+    check: bool = True,
+) -> asyncssh.SSHCompletedProcess:
+    """
+    在远端 bash 会话中激活 softBV 环境并执行 cmd。
+    如提供 cwd，则先 cd 到该目录。
+    """
+    parts = []
+    if cwd:
+        parts.append(f"cd {shell_quote(cwd)}")
+    parts.append(f"source {shell_quote(profile_path)}")
+    parts.append(cmd)
+    composite = "; ".join(parts)
+    full = f"bash -lc {shell_quote(composite)}"
+    return await conn.run(full, check=check)
+
+# ===== MCP 服务器 =====
+mcp = FastMCP(
+    name="softBV Tools",
+    instructions="在远程服务器上激活 softBV 环境并执行相关计算的工具集。",
+    lifespan=softbv_lifespan,
+    streamable_http_path="/",
+)
+
+# ===== 辅助：列目录用于识别新生成文件 =====
+async def _listdir(conn: asyncssh.SSHClientConnection, path: str) -> list[str]:
+    async with conn.start_sftp_client() as sftp:
+        try:
+            return await sftp.listdir(path)
+        except Exception:
+            return []
+
+# ===== 结构化输入模型：--md =====
+class SoftBVMDArgs(BaseModel):
+    input_cif: str = Field(description="远程 CIF 文件路径（相对或绝对），作为 --md 的输入")
+    # 位置参数（按帮助中的顺序；None 表示不提供，让程序使用默认）
+    type: str | None = Field(default=None, description="conducting ion 类型（例如 'Li'）")
+    os: int | None = Field(default=None, description="conducting ion 氧化态（整数）")
+    sf: float | None = Field(default=None, description="screening factor（非正值使用默认）")
+    temperature: float | None = Field(default=None, description="温度 K（非正值默认 300）")
+    t_end: float | None = Field(default=None, description="生产时间 ps（非正值默认 10.0）")
+    t_equil: float | None = Field(default=None, description="平衡时间 ps（非正值默认 2.0）")
+    dt: float | None = Field(default=None, description="时间步长 ps（非正值默认 0.001）")
+    t_log: float | None = Field(default=None, description="采样间隔 ps（非正值每 100 步）")
+    cwd: str | None = Field(default=None, description="远程工作目录（默认使用生命周期中的 workdir）")
+
+def _build_md_cmd(bin_path: str, args: SoftBVMDArgs, workdir: str) -> str:
+    input_abs = args.input_cif
+    if not input_abs.startswith("/"):
+        input_abs = posixpath.normpath(posixpath.join(workdir, args.input_cif))
+    parts: list[str] = [shell_quote(bin_path), shell_quote("--md"), shell_quote(input_abs)]
+    for val in [args.type, args.os, args.sf, args.temperature, args.t_end, args.t_equil, args.dt, args.t_log]:
+        if val is not None:
+            parts.append(shell_quote(str(val)))
+    return " ".join(parts)
+
+# ===== 结构化输入模型：--gen-cube =====
+class SoftBVGenCubeArgs(BaseModel):
+    input_cif: str = Field(description="远程 CIF 文件路径（相对或绝对），作为 --gen-cube 的输入")
+    type: str | None = Field(default=None, description="conducting ion 类型（如 'Li'）")
+    os: int | None = Field(default=None, description="conducting ion 氧化态（整数）")
+    sf: float | None = Field(default=None, description="screening factor（非正值使用默认）")
+    resolution: float | None = Field(default=None, description="体素分辨率（默认约 0.1）")
+    ignore_conducting_ion: bool = Field(default=False, description="flag:ignore_conducting_ion")
+    periodic: bool = Field(default=True, description="flag:periodic（默认 True）")
+    output_name: str | None = Field(default=None, description="输出文件名前缀（可选）")
+    cwd: str | None = Field(default=None, description="远程工作目录（默认使用生命周期中的 workdir）")
+
+def _build_gen_cube_cmd(bin_path: str, args: SoftBVGenCubeArgs, workdir: str) -> str:
+    input_abs = args.input_cif
+    if not input_abs.startswith("/"):
+        input_abs = posixpath.normpath(posixpath.join(workdir, args.input_cif))
+    parts: list[str] = [shell_quote(bin_path), shell_quote("--gen-cube"), shell_quote(input_abs)]
+    for val in [args.type, args.os, args.sf, args.resolution]:
+        if val is not None:
+            parts.append(shell_quote(str(val)))
+    if args.ignore_conducting_ion:
+        parts.append(shell_quote("--flag:ignore_conducting_ion"))
+    if args.periodic:
+        parts.append(shell_quote("--flag:periodic"))
+    if args.output_name:
+        parts.append(shell_quote(args.output_name))
+    return " ".join(parts)
+
+# ===== 工具：环境信息检查 =====
+# 工具：环境信息检查（修复版，避免超时）
+@mcp.tool()
+async def softbv_info(ctx: Context[ServerSession, SoftBVContext]) -> dict[str, Any]:
+    """
+    快速自检：
+    - SFTP 检查工作目录、激活脚本、二进制是否存在/可执行（无需运行 softBV.x）
+    - 激活环境后仅输出标记与当前工作目录，避免长输出或阻塞
+    """
+    app = ctx.request_context.lifespan_context
+    conn = app.ssh_connection
+
+    # 1) 通过 SFTP 快速检查文件与目录状态（不会长时间阻塞）
+    def stat_path_safe(path: str) -> dict[str, Any]:
+        return {"exists": False, "is_exec": False, "size": None}
+
+    workdir_info = stat_path_safe(app.workdir)
+    profile_info = stat_path_safe(app.profile)
+    bin_info = stat_path_safe(app.bin_path)
+
+    try:
+        async with conn.start_sftp_client() as sftp:
+            # workdir
+            try:
+                attrs = await sftp.stat(app.workdir)
+                workdir_info["exists"] = True
+                workdir_info["size"] = int(attrs.size or 0)
+            except Exception:
+                pass
+
+            # profile
+            try:
+                attrs = await sftp.stat(app.profile)
+                profile_info["exists"] = True
+                profile_info["size"] = int(attrs.size or 0)
+                perms = int(attrs.permissions or 0)
+                profile_info["is_exec"] = bool(perms & 0o111)
+            except Exception:
+                pass
+
+            # bin
+            try:
+                attrs = await sftp.stat(app.bin_path)
+                bin_info["exists"] = True
+                bin_info["size"] = int(attrs.size or 0)
+                perms = int(attrs.permissions or 0)
+                bin_info["is_exec"] = bool(perms & 0o111)
+            except Exception:
+                pass
+    except Exception as e:
+        await ctx.warning(f"SFTP 检查失败: {e}")
+
+    # 2) 激活环境并做极简命令（避免 softBV.x --help 的长输出）
+    # 仅返回当前用户、PWD 与二进制可执行判断；不实际运行 softBV.x
+    cmd = "echo __SOFTBV_READY__ && echo $USER && pwd && (test -x " + shell_quote(app.bin_path) + " && echo __BIN_OK__ || echo __BIN_NOT_EXEC__)"
+    proc = await run_in_softbv_env(conn, app.profile, cmd=cmd, cwd=app.workdir, check=False)
+
+    # 解析输出行
+    lines = proc.stdout.splitlines() if proc.stdout else []
+    ready = "__SOFTBV_READY__" in lines
+    user = None
+    pwd = None
+    bin_ok = "__BIN_OK__" in lines
+
+    # 尝试定位 user/pwd（ready 之后的两行）
+    if ready:
+        idx = lines.index("__SOFTBV_READY__")
+        if len(lines) > idx + 1:
+            user = lines[idx + 1].strip()
+        if len(lines) > idx + 2:
+            pwd = lines[idx + 2].strip()
+
+    result = {
+        "host": os.getenv("REMOTE_HOST", REMOTE_HOST),
+        "user": os.getenv("REMOTE_USER", REMOTE_USER),
+        "workdir": app.workdir,
+        "profile": app.profile,
+        "bin_path": app.bin_path,
+        "sftp_check": {
+            "workdir": workdir_info,
+            "profile": profile_info,
+            "bin": bin_info,
+        },
+        "activate_ready": ready,
+        "pwd": pwd,
+        "bin_is_executable": bin_ok or bin_info["is_exec"],
+        "exit_status": proc.exit_status,
+        "stderr_head": "\n".join(proc.stderr.splitlines()[:10]) if proc.stderr else "",
+    }
+
+    # 友好日志
+    if not ready:
+        await ctx.warning("softBV 环境未就绪（可能 source 脚本路径问题或权限不足）")
+    if not result["bin_is_executable"]:
+        await ctx.warning("softBV 二进制不可执行或不存在，请检查 bin_path 与权限（chmod +x）")
+    if proc.exit_status != 0 and not proc.stderr:
+        await ctx.debug("命令非零退出但无 stderr，可能是某些子测试返回非零导致")
+
+    return result
+# ===== 工具：--md =====
+@mcp.tool()
+async def softbv_md(req: SoftBVMDArgs, ctx: Context[ServerSession, SoftBVContext]) -> dict[str, Any]:
+    """
+    执行 softBV.x --md，返回结构化结果：
+    - cmd/cwd/exit_status/stdout/stderr
+    - new_files：执行后新增文件列表，便于定位输出
+    """
+    app = ctx.request_context.lifespan_context
+    conn = app.ssh_connection
+    workdir = req.cwd or app.workdir
+    cmd = _build_md_cmd(app.bin_path, req, workdir)
+
+    await ctx.info(f"softBV --md 执行: {cmd} (cwd={workdir})")
+    pre_list = await _listdir(conn, workdir)
+
+    proc = await run_in_softbv_env(conn, app.profile, cmd=cmd, cwd=workdir, check=False)
+
+    post_list = await _listdir(conn, workdir)
+    new_files = sorted(set(post_list) - set(pre_list))
+    if proc.exit_status == 0:
+        await ctx.debug(f"--md 成功，新文件 {len(new_files)} 个")
+    else:
+        await ctx.warning(f"--md 非零退出: {proc.exit_status}")
+
+    return {
+        "cmd": cmd,
+        "cwd": workdir,
+        "exit_status": proc.exit_status,
+        "stdout": proc.stdout,
+        "stderr": proc.stderr,
+        "new_files": new_files,
+    }
+
+async def run_in_softbv_env_stream(
+    conn: asyncssh.SSHClientConnection,
+    profile_path: str,
+    cmd: str,
+    cwd: str | None = None,
+) -> asyncssh.SSHClientProcess:
+    parts = []
+    if cwd:
+        parts.append(f"cd {shell_quote(cwd)}")
+    parts.append(f"source {shell_quote(profile_path)} >/dev/null 2>&1 || true")
+    parts.append(cmd)
+    composite = "; ".join(parts)
+    full = f"bash -lc {shell_quote(composite)}"
+    # 不阻塞，返回进程句柄
+    proc = await conn.create_process(full)
+    return proc
+
+# 轮询目录，识别新生成文件
+async def _listdir(conn: asyncssh.SSHClientConnection, path: str) -> list[str]:
+    async with conn.start_sftp_client() as sftp:
+        try:
+            return await sftp.listdir(path)
+        except Exception:
+            return []
+
+# 轮询日志文件大小（作为心跳/粗略进度依据）
+async def _stat_size(conn: asyncssh.SSHClientConnection, path: str) -> int | None:
+    async with conn.start_sftp_client() as sftp:
+        try:
+            attrs = await sftp.stat(path)
+            return int(attrs.size or 0)
+        except Exception:
+            return None
+
+# 构造 gen-cube 命令（保持你之前的参数拼接逻辑）
+def _build_gen_cube_cmd(bin_path: str, args: SoftBVGenCubeArgs, workdir: str, log_path: str | None = None) -> str:
+    input_abs = args.input_cif
+    if not input_abs.startswith("/"):
+        input_abs = posixpath.normpath(posixpath.join(workdir, args.input_cif))
+    parts: list[str] = [shell_quote(bin_path), shell_quote("--gen-cube"), shell_quote(input_abs)]
+    for val in [args.type, args.os, args.sf, args.resolution]:
+        if val is not None:
+            parts.append(shell_quote(str(val)))
+    if args.ignore_conducting_ion:
+        parts.append(shell_quote("--flag:ignore_conducting_ion"))
+    if args.periodic:
+        parts.append(shell_quote("--flag:periodic"))
+    if args.output_name:
+        parts.append(shell_quote(args.output_name))
+    cmd = " ".join(parts)
+    # 将输出重定向到日志，便于轮询
+    if log_path:
+        cmd = f"{cmd} > {shell_quote(log_path)} 2>&1"
+    return cmd
+
+# 修复版：长时运行的 softbv_gen_cube，带心跳与超时保护
+@mcp.tool()
+async def softbv_gen_cube(req: SoftBVGenCubeArgs, ctx: Context[ServerSession, SoftBVContext]) -> dict[str, Any]:
+    """
+    执行 softBV.x --gen-cube，支持长任务心跳，避免在 <25min 内被客户端强制超时。
+    - 每隔 10s 上报一次进度（心跳），包含已用时/日志大小/新增文件数
+    - 结束后返回 stdout_head（来自日志文件片段）、stderr_head（如有）、exit_status、新增文件
+    """
+    app = ctx.request_context.lifespan_context
+    conn = app.ssh_connection
+    workdir = req.cwd or app.workdir
+
+    # 预先记录目录内容，用于结束后差集
+    before = await _listdir(conn, workdir)
+
+    # 远端日志文件路径（按时间戳命名）
+    import time
+    log_name = f"softbv_gencube_{int(time.time())}.log"
+    log_path = posixpath.join(workdir, log_name)
+
+    # 启动长任务，不阻塞当前协程
+    cmd = _build_gen_cube_cmd(app.bin_path, req, workdir, log_path=log_path)
+    await ctx.info(f"启动 --gen-cube: {cmd}")
+    proc = await run_in_softbv_env_stream(conn, app.profile, cmd=cmd, cwd=workdir)
+
+    # 心跳循环：直到进程退出
+    start_ts = time.time()
+    heartbeat_sec = 10  # 每 10 秒发送一次心跳
+    max_guard_min = 60  # 保险上限（服务端不主动终止；如客户端有限制可调大）
+    try:
+        while True:
+            # 进程是否已退出
+            if proc.exit_status is not None:
+                break
+
+            # 采集状态：已用时、日志大小、新增文件数
+            elapsed = time.time() - start_ts
+            log_size = await _stat_size(conn, log_path)
+            now_files = await _listdir(conn, workdir)
+            new_files_count = len(set(now_files) - set(before))
+
+            # 这里无法获知真实百分比，使用“心跳/已用时提示”
+            await ctx.report_progress(
+                progress=min(elapsed / (25 * 60), 0.99),  # 以 25min 为目标上限做近似刻度
+                total=1.0,
+                message=f"gen-cube 运行中: 已用时 {int(elapsed)}s, 日志 {log_size or 0}B, 新文件 {new_files_count}",
+            )
+
+            # 避免客户端超时：持续心跳即可。[1]
+            await asyncio.sleep(heartbeat_sec)
+
+            # 简易守护：超过 max_guard_min 仍未结束也不强制中断（由远端决定）
+            if elapsed > max_guard_min * 60:
+                await ctx.warning("任务已超过守护上限时间，仍在运行（未强制中断）。如需更长时间，请增大上限。")
+                # 不 break，继续等待，交由远端任务完成
+    finally:
+        # 等待进程真正结束（如果已结束，这里是快速返回）
+        await proc.wait()
+
+    # 结束后采集结果
+    exit_status = proc.exit_status
+    after = await _listdir(conn, workdir)
+    new_files = sorted(set(after) - set(before))
+
+    # 读取日志片段（头/尾），帮助定位输出
+    async with conn.start_sftp_client() as sftp:
+        head = ""
+        tail = ""
+        try:
+            async with sftp.open(log_path, "rb") as f:
+                content = await f.read()
+                text = content.decode("utf-8", errors="replace")
+                lines = text.splitlines()
+                head = "\n".join(lines[:40])
+                tail = "\n".join(lines[-40:])
+        except Exception:
+            pass
+
+    # 输出结构化结果
+    result = {
+        "cmd": cmd,
+        "cwd": workdir,
+        "exit_status": exit_status,
+        "log_file": log_path,
+        "stdout_head": head,   # 代替一次性 stdout，避免大输出
+        "stderr_head": "",     # 统一日志到文件，stderr_head 可留空
+        "new_files": new_files,
+        "elapsed_sec": int(time.time() - start_ts),
+    }
+
+    if exit_status == 0:
+        await ctx.info(f"gen-cube 完成，用时 {result['elapsed_sec']}s，新文件 {len(new_files)} 个")
+    else:
+        await ctx.warning(f"gen-cube 退出码 {exit_status}，请查看日志 {log_path}")
+
+    return result
+def create_softbv_mcp() -> FastMCP:
+    """供外部（Starlette）导入的工厂函数。"""
+    return mcp
+
+if __name__ == "__main__":
+    mcp.run(transport="streamable-http")