mcp-forge/src/mcp_forge/builder/uv_installer.py

"""UV-based package installer for custom environments."""

import tempfile
import shutil
from typing import List
from pathlib import Path


class UVInstaller:
    """Manages UV-based package installation in containers."""
    
    def __init__(self, cache_path: Path):
        """
        Initialize UV installer.
        
        Args:
            cache_path: Path to UV cache directory
        """
        self.cache_path = Path(cache_path)
        self._ensure_cache_dir()
    
    def generate_requirements(self, packages: List[str]) -> str:
        """
        Generate requirements.txt content.
        
        One package per line with version specifiers preserved.
        
        Args:
            packages: List of package specifications
        
        Returns:
            requirements.txt content
        """
        if not packages:
            return ""
        
        return '\n'.join(packages)
    
    def generate_containerfile(
        self,
        base_image: str,
        packages: List[str],
        python_version: str = "3.11"
    ) -> str:
        """
        Generate Containerfile for building custom environment.
        
        Containerfile structure optimizes layer caching:
        1. Base image
        2. Install UV (cached layer)
        3. Create non-root user
        4. Copy requirements.txt (cache-friendly)
        5. Install packages with UV
        6. Set working directory
        
        Args:
            base_image: Base container image
            packages: List of package specifications
            python_version: Python version (for reference)
        
        Returns:
            Containerfile content
        """
        has_packages = bool(packages)
        
        containerfile = f"""FROM {base_image}

# Install UV for fast package installation
RUN pip install --no-cache-dir uv

# Create non-root user for security
RUN useradd -m -u 1000 -s /bin/bash forge && \\
    mkdir -p /home/forge/.cache/uv && \\
    chown -R forge:forge /home/forge

# Switch to non-root user
USER forge
WORKDIR /home/forge

# Copy requirements for layer caching
COPY --chown=forge:forge requirements.txt /home/forge/requirements.txt

"""
        
        if has_packages:
            containerfile += """# Install packages with UV
RUN uv pip install --system -r requirements.txt

"""
        
        containerfile += """# Set working directory
WORKDIR /home/forge/workspace

# Default command
CMD ["/bin/bash"]
"""
        
        return containerfile
    
    def create_build_context(
        self,
        base_image: str,
        packages: List[str],
        python_version: str = "3.11"
    ) -> Path:
        """
        Create temporary build context directory.
        
        Contains:
        - Containerfile
        - requirements.txt
        
        Args:
            base_image: Base container image
            packages: List of package specifications
            python_version: Python version
        
        Returns:
            Path to build context directory (caller must cleanup)
        """
        # Create temporary directory
        context_dir = Path(tempfile.mkdtemp(prefix="mcp-forge-build-"))
        
        try:
            # Generate Containerfile
            containerfile_content = self.generate_containerfile(
                base_image, packages, python_version
            )
            containerfile_path = context_dir / "Containerfile"
            containerfile_path.write_text(containerfile_content)
            
            # Generate requirements.txt
            requirements_content = self.generate_requirements(packages)
            requirements_path = context_dir / "requirements.txt"
            requirements_path.write_text(requirements_content)
            
            return context_dir
            
        except Exception:
            # Cleanup on error
            shutil.rmtree(context_dir, ignore_errors=True)
            raise
    
    def get_cache_volume_mount(self) -> dict:
        """
        Get volume mount configuration for UV cache.
        
        Returns:
            Volume mount dict for Podman
        """
        return {
            "bind": "/home/forge/.cache/uv",
            "mode": "rw"
        }
    
    def _ensure_cache_dir(self) -> None:
        """Ensure UV cache directory exists."""
        self.cache_path.mkdir(parents=True, exist_ok=True)