execute.py

"""Context Engine — One-Click Setup

Run from the Plugin UI to deploy all services and index your codebase.

Supports two deployment modes:
  - Container mode (default when inside Docker): installs Python deps
    and runs MCP servers directly with embedded Qdrant (no Docker needed)
  - Host mode (when Docker is available): uses Docker Compose for
    full service stack with dedicated Qdrant + Redis containers
"""

import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

_PLUGIN_DIR = Path(__file__).parent
_SERVER_DIR = _PLUGIN_DIR / "server"
_CONFIG_FILE = _PLUGIN_DIR / "config.json"
_COMPOSE_PROJECT = "a0-context-engine"
_PID_DIR = _PLUGIN_DIR / ".pids"

_DOCKER_SEARCH_PATHS = [
    "/usr/local/bin", "/opt/homebrew/bin", "/usr/bin",
    os.path.expanduser("~/.docker/bin"),
    "/Applications/Docker.app/Contents/Resources/bin",
]

# Deps that need to be installed in the A0 container for embedded mode.
# uvicorn and starlette are already available via A0's deps.
_EMBEDDED_DEPS = [
    "qdrant-client[fastembed]",
    "fastembed",
    "uvicorn",
    "starlette",
    "tree-sitter",
    "tree-sitter-python",
    "tree-sitter-javascript",
    "tree-sitter-typescript",
    "tree-sitter-go",
    "tree-sitter-rust",
    "tree-sitter-java",
    "tree-sitter-c",
    "tree-sitter-cpp",
    "tree-sitter-ruby",
    "tree-sitter-bash",
    "xxhash",
    "real-ladybug",
]


def _load_config() -> dict:
    if _CONFIG_FILE.exists():
        return json.loads(_CONFIG_FILE.read_text())
    return {}


def _is_inside_container() -> bool:
    """Detect if we're running inside a Docker container."""
    if os.path.exists("/.dockerenv"):
        return True
    if os.environ.get("DOCKER_CONTAINER") or os.environ.get("container"):
        return True
    try:
        with open("/proc/1/cgroup", "r") as f:
            return "docker" in f.read() or "kubepods" in f.read()
    except (FileNotFoundError, PermissionError):
        pass
    return False


def _find_docker() -> str | None:
    """Find the docker binary, searching common install paths."""
    found = shutil.which("docker")
    if found:
        return found
    for d in _DOCKER_SEARCH_PATHS:
        candidate = os.path.join(d, "docker")
        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    return None


# ── Embedded mode (inside Docker container) ──────────────────────────

def _install_deps() -> bool:
    """Install Python dependencies for the MCP servers."""
    print("\nInstalling dependencies...")
    print("  (This may take 1-2 minutes on first run)")

    try:
        proc = subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", "--no-warn-script-location"]
            + _EMBEDDED_DEPS,
            capture_output=True, text=True, timeout=300,
        )
        if proc.returncode != 0:
            print(f"✗ pip install failed:\n{proc.stderr[:500]}")
            return False
        print("✓ Dependencies installed")
        return True
    except subprocess.TimeoutExpired:
        print("✗ Dependency installation timed out")
        return False


def _start_embedded_server(name: str, script: str, port: int, env: dict) -> bool:
    """Start an MCP server as a background process."""
    _PID_DIR.mkdir(parents=True, exist_ok=True)
    pid_file = _PID_DIR / f"{name}.pid"

    # Kill existing process if running
    if pid_file.exists():
        try:
            old_pid = int(pid_file.read_text().strip())
            os.kill(old_pid, 9)
            print(f"  Stopped old {name} (pid {old_pid})")
        except (ValueError, ProcessLookupError, PermissionError):
            pass
        pid_file.unlink(missing_ok=True)

    script_path = _SERVER_DIR / script
    if not script_path.exists():
        print(f"✗ Server script not found: {script_path}")
        return False

    proc = subprocess.Popen(
        [sys.executable, str(script_path)],
        env={**os.environ, **env},
        stdout=open(_PID_DIR / f"{name}.log", "w"),
        stderr=subprocess.STDOUT,
        cwd=str(_SERVER_DIR),
    )
    pid_file.write_text(str(proc.pid))
    print(f"  ✓ {name} started (pid {proc.pid}, port {port})")
    return True


def _run_embedded(config: dict) -> bool:
    """Deploy in embedded mode: install deps + start servers directly."""
    print("\nDeploying in embedded mode (inside container)...")

    if not _install_deps():
        return False

    # Determine paths — resolve to container path
    configured_path = os.path.expanduser(config.get("host_index_path", "/a0"))

    # In container mode, host_index_path is a HOST path that doesn't exist here.
    # The A0 codebase is mounted at /a0 inside the container.
    if os.path.isdir(configured_path) and len(os.listdir(configured_path)) > 2:
        work_root = configured_path
    elif os.path.isdir("/a0"):
        work_root = "/a0"
        print(f"  Note: Configured path '{configured_path}' not accessible in container.")
        print(f"  Using /a0 (Agent Zero root) instead.")
    else:
        work_root = os.getcwd()
        print(f"  Note: Using current directory as workspace: {work_root}")

    print(f"  Workspace: {work_root}")

    data_dir = os.path.join(work_root, ".codebase")
    os.makedirs(data_dir, exist_ok=True)

    qdrant_path = os.path.join(data_dir, "qdrant_storage")
    collection = config.get("collection_name", "codebase")

    mem_port = "8002"
    idx_port = "8003"
    try:
        mem_port = config.get("memory_endpoint", "http://localhost:8002/mcp").split(":")[-1].split("/")[0]
        idx_port = config.get("indexer_endpoint", "http://localhost:8003/mcp").split(":")[-1].split("/")[0]
    except Exception:
        pass

    # Shared env vars for embedded servers — use local Qdrant (no server)
    base_env = {
        "QDRANT_URL": f"path:{qdrant_path}",
        "COLLECTION_NAME": collection,
        "WORK_ROOT": work_root,
        "FASTMCP_HOST": "0.0.0.0",
        "FASTMCP_TRANSPORT": "http",
        "LOG_LEVEL": "INFO",
    }

    print("\nStarting MCP servers...")

    ok = _start_embedded_server("memory", "memory_server.py", int(mem_port), {
        **base_env,
        "FASTMCP_PORT": mem_port,
    })
    if not ok:
        return False

    ok = _start_embedded_server("indexer", "indexer_server.py", int(idx_port), {
        **base_env,
        "FASTMCP_PORT": idx_port,
        "FASTMCP_INDEXER_PORT": idx_port,
        "REDIS_URL": "",  # Skip Redis in embedded mode
    })
    if not ok:
        return False

    # Give servers a moment to start (or crash)
    print("\n  Waiting for servers to initialize...")
    time.sleep(5)

    # Check if processes are still alive
    for name in ("memory", "indexer"):
        pid_file = _PID_DIR / f"{name}.pid"
        log_file = _PID_DIR / f"{name}.log"
        if pid_file.exists():
            try:
                pid = int(pid_file.read_text().strip())
                os.kill(pid, 0)  # Signal 0 = check if alive
            except (ProcessLookupError, ValueError):
                print(f"  ✗ {name} server crashed. Log:")
                if log_file.exists():
                    print(log_file.read_text()[-1000:])
                return False

    return True


# ── Docker Compose mode (host deployment) ────────────────────────────

def _check_docker() -> bool:
    """Verify Docker is installed and running."""
    print("Checking Docker...")
    docker = _find_docker()
    if not docker:
        return False

    print(f"  Found: {docker}")
    try:
        result = subprocess.run([docker, "info"], capture_output=True, timeout=15)
        if result.returncode != 0:
            return False
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return False

    print("✓ Docker is running")
    return True


def _generate_env(config: dict) -> None:
    """Write .env file for docker-compose."""
    host_path = os.path.expanduser(config.get("host_index_path", "~"))
    mem_port = "8002"
    idx_port = "8003"
    try:
        mem_port = config.get("memory_endpoint", "http://localhost:8002/mcp").split(":")[-1].split("/")[0]
        idx_port = config.get("indexer_endpoint", "http://localhost:8003/mcp").split(":")[-1].split("/")[0]
    except Exception:
        pass
    env = f"HOST_INDEX_PATH={host_path}\nCE_MEMORY_PORT={mem_port}\nCE_INDEXER_PORT={idx_port}\nLOG_LEVEL=INFO\n"
    (_SERVER_DIR / ".env").write_text(env)


def _deploy_docker() -> bool:
    """Build and start Docker containers."""
    print("\nDeploying Context Engine services via Docker...")
    print("  (First run builds images — this may take 3-5 minutes)")
    docker = _find_docker()
    try:
        proc = subprocess.run(
            [docker, "compose", "-p", _COMPOSE_PROJECT, "up", "-d", "--build"],
            cwd=str(_SERVER_DIR), timeout=600,
        )
        if proc.returncode != 0:
            print(f"✗ Deploy failed. Run: docker compose -p {_COMPOSE_PROJECT} logs")
            return False
    except subprocess.TimeoutExpired:
        print("✗ Deploy timed out after 10 minutes.")
        return False
    print("✓ Services deployed")
    return True


def _wait_for_healthy_docker(timeout: int = 120) -> bool:
    """Wait for Docker services to become healthy."""
    print("\nWaiting for services...")
    docker = _find_docker()
    start = time.time()
    while time.time() - start < timeout:
        try:
            result = subprocess.run(
                [docker, "compose", "-p", _COMPOSE_PROJECT, "ps", "--format", "json"],
                capture_output=True, text=True, cwd=str(_SERVER_DIR), timeout=10,
            )
            if result.returncode == 0 and result.stdout.strip():
                raw = result.stdout.strip()
                try:
                    services = json.loads(raw)
                    if not isinstance(services, list):
                        services = [services]
                except json.JSONDecodeError:
                    services = []
                    for line in raw.splitlines():
                        if line.strip():
                            try:
                                services.append(json.loads(line.strip()))
                            except json.JSONDecodeError:
                                pass
                if services:
                    running = sum(1 for s in services if s.get("State") == "running")
                    total = len(services)
                    print(f"  {running}/{total} services running...", end="\r")
                    if running == total:
                        print(f"\n✓ All {total} services running")
                        return True
        except Exception:
            pass
        time.sleep(3)
    print("\n⚠ Some services may still be starting.")
    return True


# ── Shared ───────────────────────────────────────────────────────────

def _test_connection(config: dict) -> bool:
    """Test that MCP endpoints are responding."""
    print("\nTesting connection...")
    import urllib.request

    endpoints = {
        "Memory": config.get("memory_endpoint", "http://localhost:8002/mcp"),
        "Indexer": config.get("indexer_endpoint", "http://localhost:8003/mcp"),
    }
    all_ok = True
    for name, url in endpoints.items():
        try:
            req = urllib.request.Request(
                url, method="POST",
                data=b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}',
                headers={"Content-Type": "application/json"},
            )
            urllib.request.urlopen(req, timeout=10)
            print(f"  ✓ {name} responding at {url}")
        except Exception as e:
            print(f"  ✗ {name} not responding: {e}")
            all_ok = False
    return all_ok


def main():
    print("=" * 50)
    print("  Context Engine — One-Click Setup")
    print("=" * 50)
    print()

    config = _load_config()
    host_path = config.get("host_index_path", ".")
    collection = config.get("collection_name", "codebase")
    in_container = _is_inside_container()

    print(f"Project path: {os.path.expanduser(host_path)}")
    print(f"Collection:   {collection}")
    print(f"Environment:  {'Container (embedded mode)' if in_container else 'Host (Docker mode)'}")

    if in_container:
        # Embedded mode: install deps + run servers directly
        if not _run_embedded(config):
            return 1
    else:
        # Docker mode: compose up
        if not _check_docker():
            print("✗ Docker is not available.")
            print("  Install Docker Desktop: https://www.docker.com/products/docker-desktop/")
            return 1
        _generate_env(config)
        if not _deploy_docker():
            return 1
        _wait_for_healthy_docker()

    # Wait for servers to initialize
    time.sleep(3)
    _test_connection(config)

    print()
    print("=" * 50)
    print("  Setup complete!")
    print("=" * 50)
    print()
    print("Next steps:")
    print("  1. Open the Context Engine dashboard (code_blocks icon in sidebar)")
    print("  2. Go to the 'Index' tab and click 'Index Full Workspace'")
    print("  3. Once indexed, use the 'Search' tab to search your code")
    print("  4. Use the 'Graph' tab to visualize symbol relationships")
    print()
    return 0


if __name__ == "__main__":
    sys.exit(main())