entrypoint.py

"""Container entrypoint for the M3GNet FastMCP service.

The script searches for ``start_mcp.py`` in the extracted project tree and
executes it with sensible defaults. This avoids hard-coding the deployment
layout, which is helpful when only the ``mcp_output`` folder is uploaded to a
Service (e.g., HuggingFace Space).
"""

from __future__ import annotations

import os
import sys
from pathlib import Path


def find_start_script(root: Path) -> Path | None:
    """Return the first ``start_mcp.py`` found within ``root`` (depth <= 3)."""

    max_depth = 3
    queue = [(root, 0)]
    while queue:
        current, depth = queue.pop(0)
        if not current.is_dir():
            continue
        # Look for the script in the current directory first
        candidate = current / "start_mcp.py"
        if candidate.is_file():
            return candidate
        if depth >= max_depth:
            continue
        for child in current.iterdir():
            if child.is_dir():
                queue.append((child, depth + 1))
    return None


def main() -> None:
    root = Path("/app")
    script = find_start_script(root)
    if script is None:
        print("[ERROR] 未找到 start_mcp.py。/app 目录内容如下：", file=sys.stderr)
        for path in root.rglob("*"):
            if path.is_file():
                print(f" - {path.relative_to(root)}", file=sys.stderr)
        sys.exit(1)

    host = os.environ.get("MCP_HOST", "0.0.0.0")
    port = os.environ.get("MCP_PORT", "7860")
    mode = os.environ.get("MCP_MODE", "sse")

    print(f"[INFO] 启动 FastMCP 服务：{script} (mode={mode}, host={host}, port={port})")
    os.execv(sys.executable, [sys.executable, str(script), "--mode", mode, "--host", host, "--port", port])


if __name__ == "__main__":
    main()