File size: 2,803 Bytes
b2c86fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import re
from pathlib import Path
DEFAULT_MCP_ROOT = (
Path(__file__).resolve().parent
/ "biomni_web"
/ "backend"
/ "data"
/ "mcp_generated"
)
TOOL_DECORATOR_RE = re.compile(r"^(\s*)@mcp\.tool\(\)\s*$")
COMMENTED_TOOL_DECORATOR_RE = re.compile(r"^(\s*)#\s*@mcp\.tool\(\)\s*$")
def build_server_name(server_file: Path) -> str:
stem = server_file.stem
if stem.endswith("_server"):
stem = stem[: -len("_server")]
return f"local_{stem.replace('-', '_').replace('.', '_')}"
def patch_file(server_file: Path) -> bool:
text = server_file.read_text(encoding="utf-8")
original = text
if "_shim_server.py" in server_file.name:
return False
if "mcp = FastMCP(" in text and 'mcp.run(transport="stdio")' in text:
return False
lines = text.splitlines()
# Enable commented decorators first so the file becomes a real MCP server.
lines = [COMMENTED_TOOL_DECORATOR_RE.sub(r"\1@mcp.tool()", line) for line in lines]
first_decorator_idx = None
for idx, line in enumerate(lines):
if TOOL_DECORATOR_RE.match(line):
first_decorator_idx = idx
break
if first_decorator_idx is None:
return False
if "from mcp.server.fastmcp import FastMCP" not in "\n".join(lines):
insertion = [
"from mcp.server.fastmcp import FastMCP",
"",
f"SERVER_NAME = {build_server_name(server_file)!r}",
"mcp = FastMCP(SERVER_NAME)",
"",
]
lines[first_decorator_idx:first_decorator_idx] = insertion
patched = "\n".join(lines).rstrip() + "\n"
if 'mcp.run(transport="stdio")' not in patched:
patched += '\nif __name__ == "__main__":\n mcp.run(transport="stdio")\n'
if patched != original:
server_file.write_text(patched, encoding="utf-8")
return True
return False
def main() -> None:
parser = argparse.ArgumentParser(
description="Patch raw MCP server files to be directly runnable with FastMCP."
)
parser.add_argument(
"--mcp-root",
default=str(DEFAULT_MCP_ROOT),
help="Root directory containing mcp_* subdirectories.",
)
args = parser.parse_args()
mcp_root = Path(args.mcp_root).resolve()
if not mcp_root.is_dir():
raise SystemExit(f"MCP root does not exist or is not a directory: {mcp_root}")
changed = 0
scanned = 0
for server_file in sorted(mcp_root.glob("mcp_*/app/*_server.py")):
scanned += 1
if patch_file(server_file):
changed += 1
print(f"Scanned {scanned} server files.")
print(f"Patched {changed} raw MCP server files.")
if __name__ == "__main__":
main()
|