topoteretes--cognee/cognee/api/v1/ui/ui.py

import os
import platform
import signal
import socket
import subprocess
import threading
import time
import webbrowser
import zipfile
import requests
from pathlib import Path
from typing import Callable, Optional, Tuple, List
import tempfile
import shutil

from cognee.shared.logging_utils import get_logger
from cognee.version import get_cognee_version
from .node_setup import check_node_npm, get_nvm_dir, get_nvm_sh_path
from .npm_utils import run_npm_command

logger = get_logger()


def _check_docker_available() -> Tuple[bool, str]:
    """
    Check if the Docker daemon is reachable by running `docker info`.

    Returns:
        Tuple of (is_available: bool, message: str) where message provides
        actionable guidance when Docker is not available.
    """
    docker_path = shutil.which("docker")
    if docker_path is None:
        return False, (
            "Docker CLI not found on PATH.\n"
            "Install Docker Desktop (https://www.docker.com/products/docker-desktop/) "
            "or Colima (https://github.com/abiosoft/colima) and ensure the `docker` "
            "binary is on your PATH."
        )

    try:
        result = subprocess.run(
            ["docker", "info"],
            capture_output=True,
            timeout=15,
        )
        if result.returncode == 0:
            return True, "Docker daemon is running."

        stderr = result.stderr.decode("utf-8", errors="replace").strip()
        return False, (
            f"Docker CLI is installed but the daemon is not responding.\n"
            f"  docker info stderr: {stderr}\n\n"
            f"Possible fixes:\n"
            f"  - Docker Desktop: open the Docker Desktop application and wait for it to start.\n"
            f"  - Colima: run `colima start` (add `--network-address` for host.docker.internal support).\n"
            f"  - Linux: run `sudo systemctl start docker` or `sudo service docker start`."
        )
    except subprocess.TimeoutExpired:
        return False, (
            "Docker daemon did not respond within 15 seconds.\n"
            "The daemon may be starting up. Wait a moment and try again, or:\n"
            "  - Docker Desktop: open the Docker Desktop application.\n"
            "  - Colima: run `colima start`."
        )
    except (OSError, subprocess.SubprocessError) as e:
        # Covers FileNotFoundError (binary vanished after the shutil.which check),
        # PermissionError / "Exec format error" (TOCTOU race, noexec mount, foreign
        # arch binary), and any other exec-time failure. Returning a tuple here keeps
        # start_ui's graceful-skip guarantee total instead of crashing the UI launch.
        return False, (
            f"Docker CLI could not be executed ({type(e).__name__}: {e}).\n"
            "Install Docker Desktop (https://www.docker.com/products/docker-desktop/) "
            "or Colima (https://github.com/abiosoft/colima) and ensure the `docker` "
            "binary is on your PATH and executable."
        )


def _stream_process_output(
    process: subprocess.Popen, stream_name: str, prefix: str, color_code: str = ""
) -> threading.Thread:
    """
    Stream output from a process with a prefix to identify the source.

    Args:
        process: The subprocess to monitor
        stream_name: 'stdout' or 'stderr'
        prefix: Text prefix for each line (e.g., '[BACKEND]', '[FRONTEND]')
        color_code: ANSI color code for the prefix (optional)

    Returns:
        Thread that handles the streaming
    """

    def stream_reader():
        stream = getattr(process, stream_name)
        if stream is None:
            return

        reset_code = "\033[0m" if color_code else ""

        try:
            for line in iter(stream.readline, b""):
                if line:
                    line_text = line.decode("utf-8").rstrip()
                    if line_text:
                        print(f"{color_code}{prefix}{reset_code} {line_text}", flush=True)
        except Exception:
            pass
        finally:
            if stream:
                stream.close()

    thread = threading.Thread(target=stream_reader, daemon=True)
    thread.start()
    return thread


def _is_port_available(port: int) -> bool:
    """
    Check if a port is available on localhost.
    Returns True if the port is available, False otherwise.
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(1)  # 1 second timeout
            result = sock.connect_ex(("localhost", port))
            return result != 0  # Port is available if connection fails
    except Exception:
        return False


def _check_required_ports(ports_to_check: List[Tuple[int, str]]) -> Tuple[bool, List[str]]:
    """
    Check if all required ports are available on localhost.

    Args:
        ports_to_check: List of (port, service_name) tuples

    Returns:
        Tuple of (all_available: bool, unavailable_services: List[str])
    """
    unavailable = []

    for port, service_name in ports_to_check:
        if not _is_port_available(port):
            unavailable.append(f"{service_name} (port {port})")
            logger.error(f"Port {port} is already in use for {service_name}")

    return len(unavailable) == 0, unavailable


def normalize_version_for_comparison(version: str) -> str:
    """
    Normalize version string for comparison.
    Handles development versions and edge cases.
    """
    # Remove common development suffixes for comparison
    normalized = (
        version.replace("-local", "").replace("-dev", "").replace("-alpha", "").replace("-beta", "")
    )
    return normalized.strip()


def get_frontend_cache_dir() -> Path:
    """
    Get the directory where downloaded frontend assets are cached.
    Uses user's home directory to persist across package updates.
    Each cached frontend is version-specific and will be re-downloaded
    when the cognee package version changes.
    """
    cache_dir = Path.home() / ".cognee" / "ui-cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def get_frontend_download_info() -> Tuple[str, str]:
    """
    Get the download URL and version for the actual cognee-frontend source.
    Downloads the real frontend from GitHub releases, matching the installed version.
    """
    version = get_cognee_version()

    # Clean up version string (remove -local suffix for development)
    clean_version = version.replace("-local", "")

    # Download from specific release tag to ensure version compatibility
    download_url = f"https://github.com/topoteretes/cognee/archive/refs/tags/v{clean_version}.zip"

    return download_url, version


def download_frontend_assets(force: bool = False) -> bool:
    """
    Download and cache frontend assets.

    Args:
        force: If True, re-download even if already cached

    Returns:
        bool: True if successful, False otherwise
    """
    cache_dir = get_frontend_cache_dir()
    frontend_dir = cache_dir / "frontend"
    version_file = cache_dir / "version.txt"

    # Check if already downloaded and up to date
    if not force and frontend_dir.exists() and version_file.exists():
        try:
            cached_version = version_file.read_text().strip()
            current_version = get_cognee_version()

            # Compare normalized versions to handle development versions
            cached_normalized = normalize_version_for_comparison(cached_version)
            current_normalized = normalize_version_for_comparison(current_version)

            if cached_normalized == current_normalized:
                logger.debug(f"Frontend assets already cached for version {current_version}")
                return True
            else:
                logger.info(
                    f"Version mismatch detected: cached={cached_version}, current={current_version}"
                )
                logger.info("Updating frontend cache to match current cognee version...")
                # Clear the old cached version
                if frontend_dir.exists():
                    shutil.rmtree(frontend_dir)
                if version_file.exists():
                    version_file.unlink()
        except Exception as e:
            logger.debug(f"Error checking cached version: {e}")
            # Clear potentially corrupted cache
            if frontend_dir.exists():
                shutil.rmtree(frontend_dir)
            if version_file.exists():
                version_file.unlink()

    download_url, version = get_frontend_download_info()

    logger.info(f"Downloading cognee frontend assets for version {version}...")
    logger.info("This will be cached and reused until the cognee version changes.")

    try:
        # Create a temporary directory for download
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir)
            archive_path = temp_path / "cognee-main.zip"

            # Download the actual cognee repository from releases
            logger.info(
                f"Downloading cognee v{version.replace('-local', '')} from GitHub releases..."
            )
            logger.info(f"URL: {download_url}")
            response = requests.get(download_url, stream=True, timeout=60)
            response.raise_for_status()

            with open(archive_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)

            # Extract the archive and find the cognee-frontend directory
            if frontend_dir.exists():
                shutil.rmtree(frontend_dir)

            with zipfile.ZipFile(archive_path, "r") as zip_file:
                # Extract to temp directory first
                extract_dir = temp_path / "extracted"
                zip_file.extractall(extract_dir)

                # Find the cognee-frontend directory in the extracted content
                # The archive structure will be: cognee-{version}/cognee-frontend/
                cognee_frontend_source = None
                for root, dirs, files in os.walk(extract_dir):
                    if "cognee-frontend" in dirs:
                        cognee_frontend_source = Path(root) / "cognee-frontend"
                        break

                if not cognee_frontend_source or not cognee_frontend_source.exists():
                    logger.error(
                        "Could not find cognee-frontend directory in downloaded release archive"
                    )
                    logger.error("This might indicate a version mismatch or missing release.")
                    return False

                # Copy the cognee-frontend to our cache
                shutil.copytree(cognee_frontend_source, frontend_dir)
                logger.debug(f"Frontend extracted to: {frontend_dir}")

            # Write version info for future cache validation
            version_file.write_text(version)
            logger.debug(f"Cached frontend for cognee version: {version}")

            logger.info(
                f"✓ Cognee frontend v{version.replace('-local', '')} downloaded and cached successfully!"
            )
            return True

    except requests.exceptions.RequestException as e:
        if "404" in str(e):
            logger.error(f"Release v{version.replace('-local', '')} not found on GitHub.")
            logger.error(
                "This version might not have been released yet, or you're using a development version."
            )
            logger.error("Try using a stable release version of cognee.")
        else:
            logger.error(f"Failed to download from GitHub: {str(e)}")
        logger.error("You can still use cognee without the UI functionality.")
        return False
    except Exception as e:
        logger.error(f"Failed to download frontend assets: {str(e)}")
        logger.error("You can still use cognee without the UI functionality.")
        return False


def find_frontend_path() -> Optional[Path]:
    """
    Find the cognee-frontend directory.
    Checks both development location and cached download location.
    """
    current_file = Path(__file__)

    # First, try development paths (for contributors/developers)
    dev_search_paths = [
        current_file.parents[4] / "cognee-frontend",  # from cognee/api/v1/ui/ui.py to project root
        current_file.parents[3] / "cognee-frontend",  # fallback path
        current_file.parents[2] / "cognee-frontend",  # another fallback
    ]

    for path in dev_search_paths:
        if path.exists() and (path / "package.json").exists():
            logger.debug(f"Found development frontend at: {path}")
            return path

    # Then try cached download location (for pip-installed users)
    cache_dir = get_frontend_cache_dir()
    cached_frontend = cache_dir / "frontend"

    if cached_frontend.exists() and (cached_frontend / "package.json").exists():
        logger.debug(f"Found cached frontend at: {cached_frontend}")
        return cached_frontend

    return None


def install_frontend_dependencies(frontend_path: Path) -> bool:
    """
    Install frontend dependencies if node_modules doesn't exist.
    This is needed for both development and downloaded frontends since both use npm run dev.
    """

    node_modules = frontend_path / "node_modules"
    if node_modules.exists():
        logger.debug("Frontend dependencies already installed")
        return True

    logger.info("Installing frontend dependencies (this may take a few minutes)...")

    try:
        result = run_npm_command(["npm", "install"], frontend_path, timeout=300)

        if result.returncode == 0:
            logger.info("Frontend dependencies installed successfully")
            return True
        else:
            logger.error(f"Failed to install dependencies: {result.stderr}")
            return False

    except subprocess.TimeoutExpired:
        logger.error("Timeout installing frontend dependencies")
        return False
    except Exception as e:
        logger.error(f"Error installing frontend dependencies: {str(e)}")
        return False


def is_development_frontend(frontend_path: Path) -> bool:
    """
    Check if this is a development frontend (has Next.js) vs downloaded assets.
    """
    package_json_path = frontend_path / "package.json"
    if not package_json_path.exists():
        return False

    try:
        import json

        with open(package_json_path) as f:
            package_data = json.load(f)

        # Development frontend has Next.js as dependency
        dependencies = package_data.get("dependencies", {})
        dev_dependencies = package_data.get("devDependencies", {})

        return "next" in dependencies or "next" in dev_dependencies
    except Exception:
        return False


def prompt_user_for_download() -> bool:
    """
    Ask user if they want to download the frontend assets.
    Returns True if user consents, False otherwise.
    """
    try:
        print("\n" + "=" * 60)
        print("🎨 Cognee UI Setup Required")
        print("=" * 60)
        print("The cognee frontend is not available on your system.")
        print("This is required to use the web interface.")
        print("\nWhat will happen:")
        print("• Download the actual cognee-frontend from GitHub")
        print("• Cache it in your home directory (~/.cognee/ui-cache/)")
        print("• Install dependencies with npm (requires Node.js)")
        print("• This is a one-time setup per cognee version")
        print("\nThe frontend will then be available offline for future use.")

        response = input("\nWould you like to download the frontend now? (y/N): ").strip().lower()
        return response in ["y", "yes"]
    except (KeyboardInterrupt, EOFError):
        print("\nOperation cancelled by user.")
        return False


def start_ui(
    pid_callback: Callable[[int], None],
    port: int = 3000,
    open_browser: bool = True,
    auto_download: bool = False,
    start_backend: bool = False,
    backend_port: int = 8000,
    start_mcp: bool = False,
    mcp_port: int = 8001,
) -> Optional[subprocess.Popen]:
    """
    Start the cognee frontend UI server, optionally with the backend API server and MCP server.

    This function will:
    1. Optionally start the cognee backend API server
    2. Optionally start the cognee MCP server
    3. Find the cognee-frontend directory (development) or download it (pip install)
    4. Check if Node.js and npm are available (for development mode)
    5. Install dependencies if needed (development mode)
    6. Start the frontend server
    7. Optionally open the browser

    Args:
        pid_callback: Callback to notify with PID of each spawned process
        port: Port to run the frontend server on (default: 3000)
        open_browser: Whether to open the browser automatically (default: True)
        auto_download: If True, download frontend without prompting (default: False)
        start_backend: If True, also start the cognee API backend server (default: False)
        backend_port: Port to run the backend server on (default: 8000)
        start_mcp: If True, also start the cognee MCP server (default: False)
        mcp_port: Port to run the MCP server on (default: 8001)

    Returns:
        subprocess.Popen object representing the running frontend server, or None if failed
        Note: If backend and/or MCP server are started, they run in separate processes
        that will be cleaned up when the frontend process is terminated.

    Example:
        >>> import cognee
        >>> def dummy_callback(pid): pass
        >>> # Start just the frontend
        >>> server = cognee.start_ui(dummy_callback)
        >>>
        >>> # Start both frontend and backend
        >>> server = cognee.start_ui(dummy_callback, start_backend=True)
        >>> # UI will be available at http://localhost:3000
        >>> # API will be available at http://localhost:8000
        >>>
        >>> # Start frontend with MCP server
        >>> server = cognee.start_ui(dummy_callback, start_mcp=True)
        >>> # UI will be available at http://localhost:3000
        >>> # MCP server will be available at http://127.0.0.1:8001/sse
        >>> # To stop all servers later:
        >>> server.terminate()
    """
    logger.info("Starting cognee UI...")

    ports_to_check = [(port, "Frontend UI")]

    if start_backend:
        ports_to_check.append((backend_port, "Backend API"))

    if start_mcp:
        ports_to_check.append((mcp_port, "MCP Server"))

    logger.info("Checking port availability...")
    all_ports_available, unavailable_services = _check_required_ports(ports_to_check)

    if not all_ports_available:
        error_msg = f"Cannot start cognee UI: The following services have ports already in use: {', '.join(unavailable_services)}"
        logger.error(error_msg)
        logger.error("Please stop the conflicting services or change the port configuration.")
        return None

    logger.info("✓ All required ports are available")
    backend_process = None

    if start_mcp:
        logger.info("Starting Cognee MCP server with Docker...")

        # Preflight: verify the Docker daemon is reachable before pulling images
        docker_ok, docker_msg = _check_docker_available()
        if not docker_ok:
            logger.error(f"Docker preflight check failed:\n{docker_msg}")
            logger.error(
                "The MCP server requires a running Docker (or Colima) daemon. "
                "Skipping MCP server startup."
            )
            # Continue without MCP — the UI and backend can still work
            start_mcp = False

    if start_mcp:
        try:
            image = "cognee/cognee-mcp:main"
            # Bound the pull so a reachable-but-stalled daemon / registry can't hang
            # start_ui indefinitely. On timeout, degrade to skipping MCP like the
            # preflight failure path rather than blocking the whole UI launch.
            try:
                subprocess.run(["docker", "pull", image], check=True, timeout=300)
            except subprocess.TimeoutExpired:
                logger.error(
                    f"Timed out after 300s pulling Docker image '{image}'. "
                    "The daemon is reachable but the pull stalled (slow network or "
                    "registry). Skipping MCP server startup."
                )
                raise

            import uuid

            container_name = f"cognee-mcp-{uuid.uuid4().hex[:8]}"

            docker_cmd = [
                "docker",
                "run",
                "--name",
                container_name,
                "-p",
                f"{mcp_port}:8000",
                "--rm",
                "-e",
                "TRANSPORT_MODE=sse",
            ]

            if start_backend:
                docker_cmd.extend(
                    [
                        "-e",
                        f"API_URL=http://localhost:{backend_port}",
                    ]
                )
                logger.info(
                    f"Configuring MCP to connect to backend API at http://localhost:{backend_port}"
                )
                logger.info("(localhost will be auto-converted to host.docker.internal)")
            else:
                cwd = os.getcwd()
                env_file = os.path.join(cwd, ".env")
                docker_cmd.extend(["--env-file", env_file])

            docker_cmd.append(image)

            mcp_process = subprocess.Popen(
                docker_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                preexec_fn=os.setsid if hasattr(os, "setsid") else None,
            )

            _stream_process_output(mcp_process, "stdout", "[MCP]", "\033[34m")  # Blue
            _stream_process_output(mcp_process, "stderr", "[MCP]", "\033[34m")  # Blue

            # Pass both PID and container name using a tuple
            pid_callback((mcp_process.pid, container_name))

            mode_info = "API mode" if start_backend else "direct mode"
            logger.info(
                f"✓ Cognee MCP server starting on http://127.0.0.1:{mcp_port}/sse ({mode_info})"
            )
        except Exception as e:
            logger.error(f"Failed to start MCP server with Docker: {str(e)}")
    # Start backend server if requested
    if start_backend:
        logger.info("Starting cognee backend API server...")
        try:
            import sys

            backend_process = subprocess.Popen(
                [
                    sys.executable,
                    "-m",
                    "uvicorn",
                    "cognee.api.client:app",
                    "--host",
                    "localhost",
                    "--port",
                    str(backend_port),
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                preexec_fn=os.setsid if hasattr(os, "setsid") else None,
            )

            # Start threads to stream backend output with prefix
            _stream_process_output(backend_process, "stdout", "[BACKEND]", "\033[32m")  # Green
            _stream_process_output(backend_process, "stderr", "[BACKEND]", "\033[32m")  # Green

            pid_callback(backend_process.pid)

            # Give the backend a moment to start
            time.sleep(2)

            if backend_process.poll() is not None:
                logger.error("Backend server failed to start - process exited early")
                return None

            logger.info(f"✓ Backend API started at http://localhost:{backend_port}")

        except Exception as e:
            logger.error(f"Failed to start backend server: {str(e)}")
            return None

    # Find frontend directory
    frontend_path = find_frontend_path()

    if not frontend_path:
        logger.info("Frontend not found locally. This is normal for pip-installed cognee.")

        # Offer to download the frontend
        if auto_download or prompt_user_for_download():
            if download_frontend_assets():
                frontend_path = find_frontend_path()
                if not frontend_path:
                    logger.error(
                        "Download succeeded but frontend still not found. This is unexpected."
                    )
                    return None
            else:
                logger.error("Failed to download frontend assets.")
                return None
        else:
            logger.info("Frontend download declined. UI functionality not available.")
            logger.info("You can still use all other cognee features without the web interface.")
            return None

    # Check Node.js and npm
    node_available, node_message = check_node_npm()
    if not node_available:
        logger.error(f"Cannot start UI: {node_message}")
        logger.error("Please install Node.js from https://nodejs.org/ to use the UI functionality")
        return None

    logger.debug(f"Environment check passed: {node_message}")

    # Install dependencies if needed
    if not install_frontend_dependencies(frontend_path):
        logger.error("Failed to install frontend dependencies")
        return None

    # Prepare environment variables
    env = os.environ.copy()
    env["HOST"] = "localhost"
    env["PORT"] = str(port)
    # The shared frontend defaults to cloud mode; this launcher always runs
    # against the local backend, so opt into local mode unless overridden.
    env.setdefault("NEXT_PUBLIC_IS_CLOUD_ENVIRONMENT", "false")

    # If nvm is installed, ensure it's available in the environment
    nvm_path = get_nvm_sh_path()
    if platform.system() != "Windows" and nvm_path.exists():
        # Add nvm to PATH for the subprocess
        nvm_dir = get_nvm_dir()
        # Find the latest Node.js version installed via nvm
        nvm_versions = nvm_dir / "versions" / "node"
        if nvm_versions.exists():
            versions = sorted(nvm_versions.iterdir(), reverse=True)
            if versions:
                latest_node_bin = versions[0] / "bin"
                if latest_node_bin.exists():
                    current_path = env.get("PATH", "")
                    env["PATH"] = f"{latest_node_bin}:{current_path}"

    # Start the development server
    logger.info(f"Starting frontend server at http://localhost:{port}")
    logger.info("This may take a moment to compile and start...")

    try:
        # Create frontend in its own process group for clean termination
        # Use shell=True on Windows for npm commands
        if platform.system() == "Windows":
            process = subprocess.Popen(
                ["npm", "run", "dev"],
                cwd=frontend_path,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                shell=True,
            )
        else:
            # On Unix-like systems, use bash with nvm sourced if available
            if nvm_path.exists():
                # Use bash to source nvm and run npm
                process = subprocess.Popen(
                    ["bash", "-c", f"source {nvm_path} && npm run dev"],
                    cwd=frontend_path,
                    env=env,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    preexec_fn=os.setsid if hasattr(os, "setsid") else None,
                )
            else:
                process = subprocess.Popen(
                    ["npm", "run", "dev"],
                    cwd=frontend_path,
                    env=env,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    preexec_fn=os.setsid if hasattr(os, "setsid") else None,
                )

        # Start threads to stream frontend output with prefix
        _stream_process_output(process, "stdout", "[FRONTEND]", "\033[33m")  # Yellow
        _stream_process_output(process, "stderr", "[FRONTEND]", "\033[33m")  # Yellow

        pid_callback(process.pid)

        # Give it a moment to start up
        time.sleep(3)

        # Check if process is still running
        if process.poll() is not None:
            logger.error("Frontend server failed to start - check the logs above for details")
            return None

        # Open browser if requested
        if open_browser:

            def open_browser_delayed():
                time.sleep(5)  # Give Next.js time to fully start
                try:
                    webbrowser.open(f"http://localhost:{port}")
                except Exception as e:
                    logger.warning(f"Could not open browser automatically: {e}")

            browser_thread = threading.Thread(target=open_browser_delayed, daemon=True)
            browser_thread.start()

        logger.info("✓ Cognee UI is starting up...")
        logger.info(f"✓ Open your browser to: http://localhost:{port}")
        logger.info("✓ The UI will be available once Next.js finishes compiling")

        return process

    except Exception as e:
        logger.error(f"Failed to start frontend server: {str(e)}")
        # Clean up backend process if it was started
        if backend_process:
            logger.info("Cleaning up backend process due to frontend failure...")
            try:
                backend_process.terminate()
                backend_process.wait(timeout=5)
            except (subprocess.TimeoutExpired, OSError, ProcessLookupError):
                try:
                    backend_process.kill()
                    backend_process.wait()
                except (OSError, ProcessLookupError):
                    pass
        return None