PlantAI-rag-new/tooling/cap/cap_selfcheck.py

#!/usr/bin/env python3
"""tooling/cap/cap_selfcheck.py —— 机制端到端自检（假 MCP server + 热加载证明）。

流程（全程同一 QA 进程，不重启）：
  1. 加载现有 capabilities.yaml → registry 工具数
  2. 临时注册假 MCP echo 工具（manifest 落盘）→ reload → 工具表出现 echo
  3. 实际调用 echo（stdio 子进程）→ 返回正确
  4. 停用（移除 manifest）→ reload → 工具表消失、MCP 连接关闭
  5. 恢复原 capabilities.yaml

用法：
  QA_ENV_FILE=<prod.env> python3 cap_selfcheck.py <qa-config-dir>

<qa-config-dir> = deploy 树根（含 configs/capabilities.yaml），
即 $GENE_FAMILY_ROOT 的值。测试中的 yaml 落盘用其真实路径，因此会短暂修改
配置文件；本脚本保证恢复。
"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

TOOL_DIR = Path(__file__).parent
FAKE_MCP = TOOL_DIR / "fake_mcp_server.py"

FAKE_MANIFEST = """
capabilities:
  - name: fake_echo
    provider: mcp
    permission: read_only
    description: 自检用 echo 工具
    session_default: true
    schema:
      text: {type: string, required: true, description: 要回显的文本}
    mcp:
      command: {python}
      args: ["{fake_mcp}"]
      tool: echo
"""


def _python() -> str:
    return os.environ.get(
        "QA_PYTHON", "/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python"
    )


def _qa_import(root: Path):
    """在目标 QA 树环境里跑一段 python（用 deploy 树的 venv + PYTHONPATH）。"""
    env = dict(os.environ)
    env["PYTHONPATH"] = f"{root}:{root}/src" + (
        f":{env['PYTHONPATH']}" if env.get("PYTHONPATH") else ""
    )
    return env


def run_step(root: Path, code: str) -> str:
    env = _qa_import(root)
    proc = subprocess.run(
        [_python(), "-c", code],
        capture_output=True,
        text=True,
        timeout=120,
        env=env,
        cwd=str(root),
    )
    if proc.returncode != 0:
        raise RuntimeError(f"step failed: {proc.stderr[-2000:]}")
    return proc.stdout.strip()


def main() -> int:
    root = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else None
    if root is None or not (root / "configs").is_dir():
        print(__doc__)
        return 2
    caps_yaml = root / "configs" / "capabilities.yaml"
    original = caps_yaml.read_text(encoding="utf-8") if caps_yaml.is_file() else "capabilities: []\n"

    manifest = FAKE_MANIFEST.format(python=_python(), fake_mcp=FAKE_MCP)
    results = {}
    try:
        # 1. 初始工具表
        out = run_step(root, (
            "from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
            "r=get_registry();"
            "print(json.dumps({'tools':[t.name for t in r.list_all()],"
            "'warnings':r.load_warnings}, ensure_ascii=False))"
        ))
        initial = json.loads(out)
        results["initial_tools"] = initial["tools"]
        assert "fake_echo" not in initial["tools"], "预置状态异常：fake_echo 已在工具表"

        # 2. 注册 + reload（热加载）
        caps_yaml.write_text(manifest, encoding="utf-8")
        out = run_step(root, (
            "from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
            "w,d=get_registry().reload();"
            "print(json.dumps({'warnings':w,'diff':d}, ensure_ascii=False))"
        ))
        reloaded = json.loads(out)
        assert "fake_echo" in reloaded["diff"]["added"], "热加载未新增 fake_echo"
        results["hot_reload_added"] = reloaded["diff"]["added"]

        # 3. 实际调用（MCP stdio 子进程）
        out = run_step(root, (
            "from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
            "from src.rag.qa_v1_1_73.capabilities.providers import get_dispatcher;"
            "m=get_registry().get('fake_echo').manifest;"
            "r=get_dispatcher().execute(m, {'text':'你好机制'});"
            "print(r)"
        ))
        assert "你好机制" in out, f"echo 返回异常: {out}"
        results["mcp_call_ok"] = True

        # 4. 停用 + reload（热卸载）
        caps_yaml.write_text("capabilities: []\n", encoding="utf-8")
        out = run_step(root, (
            "from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
            "from src.rag.qa_v1_1_73.capabilities.providers import get_dispatcher;"
            "w,d=get_registry().reload();"
            "get_dispatcher().shutdown_connections(d['removed']);"
            "print(json.dumps({'warnings':w,'diff':d}, ensure_ascii=False))"
        ))
        unloaded = json.loads(out)
        assert "fake_echo" in unloaded["diff"]["removed"], "热卸载未移除 fake_echo"
        results["hot_reload_removed"] = unloaded["diff"]["removed"]

        # 5. 恢复原文件
        caps_yaml.write_text(original, encoding="utf-8")
        run_step(root, (
            "from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
            "get_registry().reload();print('restored')"
        ))
        results["restored"] = True

        print("SELFCHECK PASS:", json.dumps(results, ensure_ascii=False, indent=2))
        return 0
    except Exception as exc:
        print(f"SELFCHECK FAIL: {exc}")
        return 1
    finally:
        caps_yaml.write_text(original, encoding="utf-8")


if __name__ == "__main__":
    raise SystemExit(main())