setup.py

"""Build configuration hook for the iai-mcp wheel.

All project metadata lives in pyproject.toml.  This file exists solely to
register a custom build_py subclass that:
- compiles the TypeScript MCP wrapper and stages the resulting JS files into
  the wheel build directory (build_lib) before the wheel is assembled;
- stages the native-extension type stubs (*.pyi, py.typed) from the Rust
  workspace beside the compiled extension in the wheel.

Both operations write into build_lib, never into the source tree, so an
editable checkout stays clean.

Editable installs (pip install -e .) skip the npm build entirely.  The
install script (scripts/install.sh) builds the wrapper separately; the
resolver falls back to mcp-wrapper/dist/ on an editable install.
"""

from __future__ import annotations

import shutil
import subprocess
from pathlib import Path

from setuptools import setup
from setuptools.command.build_py import build_py as _OrigBuildPy

_REPO_ROOT = Path(__file__).parent
_WRAPPER_SRC = _REPO_ROOT / "mcp-wrapper"

# Tracked type stubs for the native extension; staged flat beside the
# compiled .so in the wheel so that `Path(iai_mcp_native.__file__).parent`
# finds them after installation.
_NATIVE_STUB_SRC = _REPO_ROOT / "rust" / "iai_mcp_native" / "iai_mcp_native"
_NATIVE_STUB_FILES = [
    # Primary stub for the flat-layout wheel (importable as iai_mcp_native.pyi).
    ("__init__.pyi", "iai_mcp_native.pyi"),
    ("embed.pyi", "embed.pyi"),
    ("graph.pyi", "graph.pyi"),
    ("py.typed", "py.typed"),
]


class _BuildWithWrapper(_OrigBuildPy):
    """build_py subclass that compiles the TS wrapper and stages native stubs.

    At wheel-build time: collects the package into build_lib (via the parent
    ``build_py``), runs ``npm ci && npm run build`` inside mcp-wrapper/, then
    stages the resulting JS files into build_lib/iai_mcp/_wrapper/ and stages
    the native extension type stubs flat into build_lib/ so they ship beside
    the compiled extension.  The source tree is never touched.

    At editable-install time (``pip install -e .``): returns immediately without
    touching npm.  The editable resolver finds the wrapper via mcp-wrapper/dist/
    and the stubs via the maturin editable package directory.
    """

    def run(self) -> None:
        # Editable installs must never trigger npm. The install script builds
        # the wrapper as a separate step.
        if self.editable_mode:
            super().run()
            return

        # Collect the package into build_lib FIRST so build_lib/iai_mcp/ exists,
        # THEN stage the freshly compiled JS into build_lib/iai_mcp/_wrapper/
        # and the native stubs flat into build_lib/.
        # Staging into build_lib (not the source tree) keeps the checkout clean.
        super().run()
        self._build_ts_wrapper()
        self._stage_native_stubs()

    def _build_ts_wrapper(self) -> None:
        """Compile the TypeScript wrapper and stage the JS output into build_lib."""
        if not (_WRAPPER_SRC / "package.json").exists():
            raise RuntimeError(
                "mcp-wrapper/package.json not found.  The TypeScript source must be "
                "present to build the MCP wrapper.  If you are building from an sdist, "
                "ensure MANIFEST.in includes the mcp-wrapper source."
            )

        npm_exe = shutil.which("npm")
        if npm_exe is None:
            raise RuntimeError(
                "Node.js/npm is required to build the MCP wrapper.  "
                "Install Node.js >=18 and ensure 'npm' is on your PATH, then retry."
            )

        # Install exact locked dependencies.
        subprocess.run(
            [npm_exe, "ci", "--prefer-offline"],
            cwd=str(_WRAPPER_SRC),
            check=True,
        )

        # Compile TypeScript → dist/*.js
        subprocess.run(
            [npm_exe, "run", "build"],
            cwd=str(_WRAPPER_SRC),
            check=True,
        )

        dist_dir = _WRAPPER_SRC / "dist"
        if not dist_dir.exists():
            raise RuntimeError(
                f"Expected {dist_dir} after 'npm run build' but the directory is absent.  "
                "Check the TypeScript compiler output above for errors."
            )

        # Stage the freshly built JS into the wheel build directory
        # (build_lib/iai_mcp/_wrapper/) — never into the source tree, so an
        # editable checkout stays clean and its resolver falls back to
        # mcp-wrapper/dist/ as intended.
        # Copy *.js only — source maps (.js.map) are excluded from the wheel.
        wrapper_dest = Path(self.build_lib) / "iai_mcp" / "_wrapper"
        if wrapper_dest.exists():
            shutil.rmtree(wrapper_dest)
        wrapper_dest.mkdir(parents=True)

        js_files = sorted(dist_dir.glob("*.js"))
        if not js_files:
            raise RuntimeError(
                f"No *.js files found in {dist_dir}.  "
                "The 'npm run build' step produced no output."
            )
        for js_file in js_files:
            shutil.copy2(js_file, wrapper_dest / js_file.name)

    def _stage_native_stubs(self) -> None:
        """Stage the native extension type stubs flat into build_lib.

        setuptools-rust places the compiled extension flat at the build_lib
        root (i.e. ``build_lib/iai_mcp_native.cpython-*.so``), so stubs must
        land at the same level to be found via
        ``Path(iai_mcp_native.__file__).parent`` after installation.
        """
        build_lib_root = Path(self.build_lib)
        for src_name, dest_name in _NATIVE_STUB_FILES:
            src = _NATIVE_STUB_SRC / src_name
            if src.exists():
                shutil.copy2(src, build_lib_root / dest_name)


setup(cmdclass={"build_py": _BuildWithWrapper})