195 lines
7.7 KiB
Python
195 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""mcp_client_demo.py — rtl-signal-mcp 验收脚本。
|
|
|
|
不做真 MCP 握手库，直接以子进程方式拉起 mcp_server.py，按行写 JSON-RPC、按行读响应，
|
|
模拟一个 agent 的完整调用序列：
|
|
|
|
initialize -> notifications/initialized -> tools/list
|
|
-> build_index -> list_modules -> module_summary -> find_signal
|
|
-> trace_drivers -> trace_loads -> get_source -> hierarchy
|
|
|
|
用法：
|
|
python3 mcp_client_demo.py /data/repos/e203_hbirdv2/rtl/e203
|
|
python3 mcp_client_demo.py # 缺省用 E203 路径（不存在则报错提示）
|
|
|
|
退出码 0 = 全部步骤通过；非 0 = 某步失败（打印到 stderr）。
|
|
"""
|
|
import json, os, subprocess, sys
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
SERVER = os.path.join(HERE, "mcp_server.py")
|
|
|
|
DEFAULT_ROOTS = [
|
|
"/data/repos/e203_hbirdv2/rtl/e203", # 本地缺省
|
|
]
|
|
|
|
|
|
def start_server(roots):
|
|
env = dict(os.environ)
|
|
env["RTL_MCP_ROOTS"] = ":".join(roots)
|
|
return subprocess.Popen(
|
|
[sys.executable, SERVER],
|
|
stdin=subprocess.PIPE, stdout=subprocess.PIPE,
|
|
env=env, text=True, encoding="utf-8", bufsize=1)
|
|
|
|
|
|
def call(proc, method, params=None, msg_id=[0]):
|
|
"""写一条请求并读回响应（跳过通知）。返回 result 或抛 RuntimeError。"""
|
|
msg_id[0] += 1
|
|
req = dict(jsonrpc="2.0", id=msg_id[0], method=method, params=params or {})
|
|
proc.stdin.write(json.dumps(req) + "\n")
|
|
proc.stdin.flush()
|
|
while True:
|
|
line = proc.stdout.readline()
|
|
if not line:
|
|
raise RuntimeError(f"server closed during {method}")
|
|
resp = json.loads(line)
|
|
if resp.get("id") == msg_id[0]:
|
|
if "error" in resp:
|
|
raise RuntimeError(f"{method} error: {resp['error']}")
|
|
return resp["result"]
|
|
|
|
|
|
def notify(proc, method, params=None):
|
|
req = dict(jsonrpc="2.0", method=method, params=params or {})
|
|
proc.stdin.write(json.dumps(req) + "\n")
|
|
proc.stdin.flush()
|
|
|
|
|
|
def show(title, text, max_lines=12):
|
|
print(f"\n=== {title} ===")
|
|
lines = text.splitlines()
|
|
for ln in lines[:max_lines]:
|
|
print(" ", ln)
|
|
if len(lines) > max_lines:
|
|
print(f" ... ({len(lines) - max_lines} more lines)")
|
|
|
|
|
|
def fail(step, msg):
|
|
raise RuntimeError(f"step {step} failed: {msg}")
|
|
|
|
|
|
def main():
|
|
roots = [a for a in sys.argv[1:] if not a.startswith("-")] or DEFAULT_ROOTS
|
|
missing = [r for r in roots if not os.path.exists(r)]
|
|
if missing:
|
|
print(f"roots not found: {missing}\n"
|
|
f"usage: python3 mcp_client_demo.py <rtl_dir> [...]", file=sys.stderr)
|
|
return 2
|
|
|
|
proc = start_server(roots)
|
|
try:
|
|
# 1. MCP 握手
|
|
r = call(proc, "initialize", dict(
|
|
protocolVersion="2024-11-05",
|
|
capabilities={},
|
|
clientInfo=dict(name="demo-client", version="0.1")))
|
|
ver = r["protocolVersion"]
|
|
print(f"[1] initialize OK server={r['serverInfo']['name']} "
|
|
f"v{r['serverInfo']['version']} protocol={ver}")
|
|
notify(proc, "notifications/initialized")
|
|
|
|
# 2. 工具清单
|
|
r = call(proc, "tools/list")
|
|
tools = [t["name"] for t in r["tools"]]
|
|
print(f"[2] tools/list OK {len(tools)} tools: {', '.join(tools)}")
|
|
|
|
# 3. 建索引（启动时已用 RTL_MCP_ROOTS 预建过一次，这里验证重复调用幂等）
|
|
r = call(proc, "tools/call", dict(name="build_index",
|
|
arguments=dict(roots=roots)))
|
|
st = r["structuredContent"]
|
|
print(f"[3] build_index OK modules={st['modules']} "
|
|
f"cross_edges={st['cross_edges']}")
|
|
|
|
# 4. 模块列表（取第一个模块做后续实验对象）
|
|
r = call(proc, "tools/call", dict(name="list_modules", arguments={}))
|
|
mods = r["structuredContent"]["modules"]
|
|
print(f"[4] list_modules OK total={r['structuredContent']['total']} "
|
|
f"first={mods[0]['module']}")
|
|
|
|
# 5. 挑一个"实例数较多"的模块做摘要
|
|
target = max(mods, key=lambda m: m["insts"])["module"]
|
|
r = call(proc, "tools/call", dict(name="module_summary",
|
|
arguments=dict(module=target)))
|
|
show(f"module_summary {target}", r["content"][0]["text"], 8)
|
|
|
|
# 6. 找一个跨模块信号：用摘要里第一个输出端口
|
|
r2 = call(proc, "tools/call", dict(name="module_summary",
|
|
arguments=dict(module=target)))
|
|
sc = r2["structuredContent"]
|
|
port = None
|
|
for p in sc.get("ports", []):
|
|
# 端口格式是 "Output nice_xs_off" / "Input csr_ena" 这样的字符串
|
|
if str(p).startswith("Output"):
|
|
port = str(p).split(None, 1)[1]
|
|
break
|
|
if port is None: # 摘要结构兜底：直接试常见信号
|
|
port = "clk"
|
|
r = call(proc, "tools/call", dict(name="find_signal",
|
|
arguments=dict(pattern=port, limit=5)))
|
|
hits = r["structuredContent"]["hits"]
|
|
print(f"[6] find_signal '{port}' OK {len(hits)} hits")
|
|
|
|
# 7. 追驱动链（核心能力）
|
|
r = call(proc, "tools/call", dict(name="trace_drivers",
|
|
arguments=dict(module=target,
|
|
signal=port, depth=4)))
|
|
sc = r["structuredContent"]
|
|
n = len(sc.get("result", []))
|
|
print(f"[7] trace_drivers {target}.{port} OK {n} rows")
|
|
show("driver chain", r["content"][0]["text"], 12)
|
|
if n == 0 and port != "clk":
|
|
# 输出端口可能由子模块驱动也可能本就是顶层输入，空不算失败
|
|
print(" (empty chain — top-level input or tb-driven)")
|
|
|
|
# 8. 追负载
|
|
r = call(proc, "tools/call", dict(name="trace_loads",
|
|
arguments=dict(module=target,
|
|
signal=port, depth=3)))
|
|
n = len(r["structuredContent"].get("result", []))
|
|
print(f"[8] trace_loads {target}.{port} OK {n} rows")
|
|
|
|
# 9. 从驱动链里拿一个 file:line 看小段源码
|
|
r = call(proc, "tools/call", dict(name="trace_drivers",
|
|
arguments=dict(module=target,
|
|
signal=port, depth=2)))
|
|
rows = r["structuredContent"].get("result", [])
|
|
loc = None
|
|
import re
|
|
for row in rows:
|
|
# where 形如 "e203_exu_csr (e203_exu_csr.v:210)"
|
|
m = re.search(r"([\w./]+):(\d+)\)\s*$", row.get("where", ""))
|
|
if row.get("kind") == "assign" and m:
|
|
loc = (m.group(1), int(m.group(2)))
|
|
break
|
|
if loc:
|
|
fname, lineno = loc
|
|
r = call(proc, "tools/call", dict(
|
|
name="get_source",
|
|
arguments=dict(file=fname, start=max(1, lineno - 5),
|
|
end=lineno + 5)))
|
|
show(f"get_source {fname}:{lineno - 5}-{lineno + 5}",
|
|
r["content"][0]["text"], 12)
|
|
if "error" in r["structuredContent"]:
|
|
fail("get_source", r["structuredContent"]["error"])
|
|
else:
|
|
print("[9] get_source skipped (no assign location in chain)")
|
|
|
|
# 10. 层次树
|
|
r = call(proc, "tools/call", dict(name="hierarchy",
|
|
arguments=dict(module=target, depth=2)))
|
|
show("hierarchy", r["content"][0]["text"], 10)
|
|
|
|
print("\nALL CHECKS PASSED")
|
|
return 0
|
|
finally:
|
|
try:
|
|
proc.stdin.close()
|
|
except Exception:
|
|
pass
|
|
proc.terminate()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|