aift_mcp

  1"""Optional AIFT Model Context Protocol server entry point.
  2
  3Usage::
  4
  5    python aift_mcp.py --transport stdio
  6"""
  7
  8from __future__ import annotations
  9
 10import argparse
 11import ipaddress
 12import logging
 13import sys
 14from collections.abc import Sequence
 15
 16from runtime_compat import UnsupportedPythonVersionError, assert_supported_python_version
 17
 18_MCP_INSTALL_MESSAGE = (
 19    "AIFT MCP support requires the 'mcp' package. "
 20    "Install it with: pip install -r requirements.txt"
 21)
 22_DEFAULT_HTTP_HOST = "127.0.0.1"
 23_DEFAULT_HTTP_PORT = 8765
 24_LOG_LEVELS = ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL")
 25
 26
 27class MCPStartupError(RuntimeError):
 28    """Raised for startup failures that should be reported cleanly."""
 29
 30
 31class _StderrArgumentParser(argparse.ArgumentParser):
 32    """ArgumentParser that never writes help or errors to stdout."""
 33
 34    def _print_message(self, message: str, file: object | None = None) -> None:
 35        if message:
 36            super()._print_message(message, sys.stderr)
 37
 38
 39def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace:
 40    """Parse command-line arguments for the MCP entry point."""
 41    parser = _StderrArgumentParser(
 42        description="Run the optional AIFT MCP server.",
 43    )
 44    parser.add_argument(
 45        "--transport",
 46        choices=("stdio", "streamable-http"),
 47        default="stdio",
 48        help=(
 49            "MCP transport to use. Default: stdio. Streamable HTTP is for "
 50            "local integrations only and is unsupported for remote exposure "
 51            "unless --allow-remote is deliberately set."
 52        ),
 53    )
 54    parser.add_argument(
 55        "--host",
 56        default=_DEFAULT_HTTP_HOST,
 57        help=(
 58            "Host for --transport streamable-http. Default: 127.0.0.1. "
 59            "Non-loopback hosts require --allow-remote and are unsupported "
 60            "by default."
 61        ),
 62    )
 63    parser.add_argument(
 64        "--port",
 65        type=_port_value,
 66        default=_DEFAULT_HTTP_PORT,
 67        help=f"Port for --transport streamable-http. Default: {_DEFAULT_HTTP_PORT}.",
 68    )
 69    parser.add_argument(
 70        "--allow-remote",
 71        action="store_true",
 72        help=(
 73            "Allow Streamable HTTP to bind a non-loopback host. Remote exposure "
 74            "is unsupported by default."
 75        ),
 76    )
 77    parser.add_argument(
 78        "--log-level",
 79        choices=_LOG_LEVELS,
 80        default="WARNING",
 81        help="Python logging level. Logs always go to stderr. Default: WARNING.",
 82    )
 83    args = parser.parse_args(argv)
 84    args.host = args.host.strip()
 85    if args.transport == "streamable-http" and not args.host:
 86        parser.error("--host must not be empty for --transport streamable-http.")
 87    if (
 88        args.transport == "streamable-http"
 89        and not args.allow_remote
 90        and not _is_loopback_host(args.host)
 91    ):
 92        parser.error(
 93            "--transport streamable-http binds to loopback hosts only by default; "
 94            "use --allow-remote to deliberately bind a non-loopback host."
 95        )
 96    return args
 97
 98
 99def _port_value(value: str) -> int:
100    """Return a valid TCP port value for argparse."""
101    try:
102        port = int(value)
103    except ValueError as exc:
104        raise argparse.ArgumentTypeError("port must be an integer") from exc
105    if not 1 <= port <= 65535:
106        raise argparse.ArgumentTypeError("port must be between 1 and 65535")
107    return port
108
109
110def _is_loopback_host(host: str) -> bool:
111    """Return whether a host string identifies a loopback bind address."""
112    normalized = host.strip().lower()
113    if normalized in {"localhost", "localhost."}:
114        return True
115    try:
116        return ipaddress.ip_address(normalized).is_loopback
117    except ValueError:
118        return False
119
120
121def _configure_logging(level_name: str) -> None:
122    """Configure Python logging so entry-point diagnostics never use stdout."""
123    logging.basicConfig(
124        level=getattr(logging, level_name),
125        format="%(levelname)s:%(name)s:%(message)s",
126        stream=sys.stderr,
127        force=True,
128    )
129
130
131def _is_missing_mcp_import(error: ImportError) -> bool:
132    """Return whether an ImportError appears to be the optional MCP SDK."""
133    missing_name = getattr(error, "name", None)
134    return bool(missing_name == "mcp" or str(missing_name).startswith("mcp."))
135
136
137def _build_and_run_server(
138    transport: str,
139    *,
140    host: str = _DEFAULT_HTTP_HOST,
141    port: int = _DEFAULT_HTTP_PORT,
142) -> None:
143    """Build and run the MCP server using the requested transport."""
144    if transport not in {"stdio", "streamable-http"}:
145        raise MCPStartupError(f"Unsupported MCP transport: {transport}")
146
147    try:
148        from app.automation.mcp_server import MissingMCPDependencyError, build_mcp_server
149    except ImportError as exc:
150        if _is_missing_mcp_import(exc):
151            raise MCPStartupError(_MCP_INSTALL_MESSAGE) from exc
152        raise MCPStartupError(f"Failed to import AIFT MCP server: {exc}") from exc
153
154    try:
155        build_kwargs = {}
156        if transport == "streamable-http":
157            build_kwargs = {"transport_host": host, "transport_port": port}
158        server = build_mcp_server(**build_kwargs)
159    except MissingMCPDependencyError as exc:
160        raise MCPStartupError(str(exc)) from exc
161
162    if transport == "stdio":
163        server.run(transport="stdio")
164    else:
165        server.run(transport="streamable-http")
166
167
168def main(argv: Sequence[str] | None = None) -> int:
169    """Validate the runtime, then start the optional MCP server."""
170    try:
171        assert_supported_python_version()
172        args = _parse_args(argv)
173        _configure_logging(args.log_level)
174        _build_and_run_server(args.transport, host=args.host, port=args.port)
175    except UnsupportedPythonVersionError as exc:
176        print(str(exc), file=sys.stderr)
177        return 1
178    except MCPStartupError as exc:
179        print(str(exc), file=sys.stderr)
180        return 1
181    except SystemExit as exc:
182        return int(exc.code or 0)
183    except Exception as exc:
184        print(f"AIFT MCP startup failed: {exc}", file=sys.stderr)
185        return 1
186    return 0
187
188
189if __name__ == "__main__":
190    raise SystemExit(main())