insights/bin/sandbox-entrypoint.py

#!/usr/bin/env python3
"""
Sandbox container entrypoint.

Four modes, dispatched at the bottom of this file:

  1. Root phase (UID 0): create sandbox user, seed SSH/git config, re-exec
     as the sandbox user. Workspace is already populated by bin/sandbox.
  2. User phase (default): launch tmux with claude immediately in window 1
     and a setup window in window 2; PID 1 poll-blocks on the tmux session.
  3. Setup phase (SANDBOX_MODE=setup): runs in tmux window 2. Installs deps
     (Python/Node/Rust in parallel), migrates, seeds demo data, spawns the
     phrocs window, then exec's into bash -l so the window stays usable.
  4. Cache-init phase (SANDBOX_MODE=cache-init): one-shot cache warming for
     `bin/sandbox create` — installs deps, migrates, generates demo data,
     pre-builds Rust, then exits so the host can snapshot databases.
"""

from __future__ import annotations

import os
import sys
import json
import stat
import time
import shutil
import socket
import traceback
import subprocess
import urllib.request
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from textwrap import dedent

WORKSPACE = Path("/workspace")
# The sandbox user's $HOME. The base image creates this user at runtime (it's
# UID-agnostic and shared); the per-user tools image creates it at build time.
# Both must agree on this path, so it is mirrored by SANDBOX_HOME_IN_CONTAINER
# in bin/sandbox_tools.py — keep the two in sync.
SANDBOX_HOME = Path("/home/sandbox")
PROGRESS_FILE = Path("/tmp/sandbox-progress")

# Shown in tmux status bar, polled every 2s by tmux.sandbox.conf.
STATUS_FILE = Path("/tmp/sandbox-status")

# The host's detached restore writes its status here: "done" once the (cache)
# Datastore RESTORE has fully landed, or a failure reason if it dies before the
# database is ready. The setup window waits for "done" before migrating and
# aborts on a failure reason instead of timing out. Written by bin/sandbox.
RESTORE_STATUS_FILE = Path("/tmp/sandbox-restore-status")

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _ts() -> str:
    return time.strftime("%H:%M:%S", time.gmtime())


def log(msg: str) -> None:
    """Print to container stdout only (root phase)."""
    print(f"[{_ts()}] ==> {msg}", flush=True)  # noqa: T201


def info(msg: str) -> None:
    """Log to stdout and write to progress file for the host script."""
    print(f"[{_ts()}] ==> {msg}", flush=True)  # noqa: T201
    with PROGRESS_FILE.open("a") as f:
        f.write(f"[{_ts()}] ==> {msg}\n")


def run(cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, check=True, **kwargs)


def run_quiet(cmd: list[str]) -> subprocess.CompletedProcess:
    return subprocess.run(cmd, check=False, capture_output=True)


def write_file(path: Path, content: str, mode: int | None = None) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content)
    if mode is not None:
        path.chmod(mode)


def write_file_if_missing(path: Path, content: str) -> None:
    if not path.exists():
        write_file(path, content)


def _write_status(step: str) -> None:
    """Update the tmux status bar label. Swallows OSError — purely cosmetic."""
    try:
        STATUS_FILE.write_text(step)
    except OSError:
        pass


# ---------------------------------------------------------------------------
# Root phase — runs as UID 0
# ---------------------------------------------------------------------------


def create_sandbox_user(uid: int, gid: int) -> None:
    """Create passwd/group/shadow entries so tools resolve the UID."""
    if run_quiet(["getent", "passwd", str(uid)]).returncode == 0:
        return

    with open("/etc/passwd", "a") as f:
        f.write(f"sandbox:x:{uid}:{gid}:sandbox:{SANDBOX_HOME}:/bin/bash\n")
    if run_quiet(["getent", "group", str(gid)]).returncode != 0:
        with open("/etc/group", "a") as f:
            f.write(f"sandbox:x:{gid}:\n")
    with open("/etc/shadow", "a") as f:
        f.write("sandbox:*:19000:0:99999:7:::\n")


def export_environment(uid: int, gid: int) -> None:
    """Write env vars to /etc/environment and /etc/profile.d for SSH sessions."""
    extra_vars = {
        "HOME": str(SANDBOX_HOME),
        # flox sets VIRTUAL_ENV by sourcing the venv's activate; the sandbox
        # never activates. bin/insightscli and bin/ruff.sh prefer it over PATH. Equal
        # to UV_PROJECT_ENVIRONMENT, so uv sync emits no mismatch warning.
        "VIRTUAL_ENV": "/cache/python",
        "UV_CACHE_DIR": "/cache/uv",
        "UV_LINK_MODE": "copy",
        "XDG_CACHE_HOME": "/tmp/sandbox-cache",
        "CARGO_TARGET_DIR": "/cache/cargo-target",
        "npm_config_store_dir": "/cache/pnpm",
        "COREPACK_ENABLE_AUTO_PIN": "0",
        "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0",
    }
    skip = {"HOSTNAME", "TERM", "SHELL", "PWD", "SHLVL", "_", "OLDPWD", "HOME", "USER", "LOGNAME"}

    lines = [f"{k}={v}" for k, v in os.environ.items() if k not in skip]
    lines.extend(f"{k}={v}" for k, v in extra_vars.items())

    Path("/etc/environment").write_text("\n".join(lines) + "\n")
    Path("/etc/profile.d/sandbox-env.sh").write_text("\n".join(f"export {line}" for line in lines) + "\n")


def _configure_ssh_agent_env() -> None:
    """Set SSH_AUTH_SOCK only if /ssh-agent is a real socket (not /dev/null fallback)."""
    if Path("/ssh-agent").exists() and stat.S_ISSOCK(Path("/ssh-agent").lstat().st_mode):
        os.environ["SSH_AUTH_SOCK"] = "/ssh-agent"
    else:
        os.environ.pop("SSH_AUTH_SOCK", None)


def configure_user_ssh(uid: int, gid: int) -> None:
    """Seed ~/.ssh/config so git push doesn't hang on a host-key prompt in tmux."""
    ssh_dir = SANDBOX_HOME / ".ssh"
    ssh_dir.mkdir(parents=True, exist_ok=True)
    config = ssh_dir / "config"
    config.write_text(
        dedent("""\
            Host github.com
                StrictHostKeyChecking accept-new
                UserKnownHostsFile ~/.ssh/known_hosts
            """)
    )
    ssh_dir.chmod(0o700)
    config.chmod(0o600)
    run(["chown", "-R", f"{uid}:{gid}", str(ssh_dir)])


def start_sshd(uid: int, gid: int) -> None:
    """Start sshd if authorized keys are present."""
    keys = Path("/tmp/sandbox-authorized-keys")
    if not keys.exists() or keys.stat().st_size == 0:
        return

    ssh_dir = SANDBOX_HOME / ".ssh"
    ssh_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy2(keys, ssh_dir / "authorized_keys")
    ssh_dir.chmod(0o700)
    (ssh_dir / "authorized_keys").chmod(0o600)
    run(["chown", "-R", f"{uid}:{gid}", str(ssh_dir)])

    run(
        [
            "/usr/sbin/sshd",
            "-p",
            "2222",
            "-o",
            "PidFile=/tmp/sshd.pid",
            "-o",
            "PasswordAuthentication=no",
            "-o",
            "PermitRootLogin=no",
        ]
    )
    log("sshd listening on port 2222")


def _strip_host_permissions(settings_path: Path) -> None:
    """Strip permission policy so host rules can't override sandbox bypass mode.

    deny survives bypass by design; defaultMode in settings.local.json would
    shadow bypassPermissions in settings.json. allow is moot under bypass.
    """
    if not settings_path.exists():
        return
    try:
        settings = json.loads(settings_path.read_text())
    except (OSError, json.JSONDecodeError):
        return
    perms = settings.get("permissions")
    if not isinstance(perms, dict):
        return
    perms.pop("allow", None)
    perms.pop("deny", None)
    perms.pop("defaultMode", None)
    if not perms:
        settings.pop("permissions", None)
    settings_path.write_text(json.dumps(settings, indent=2) + "\n")


def _force_sandbox_claude_defaults(settings_path: Path) -> None:
    """Make the sandbox's Claude default to max effort + bypass permissions.

    The sandbox is a disposable, isolated container — the safety friction that
    justifies permission prompts on the host doesn't apply here, and the Claude
    window boots unattended in tmux. Plugins and theme survive the merge;
    allow/deny lists are stripped separately by _strip_host_permissions.
    """
    try:
        settings = json.loads(settings_path.read_text()) if settings_path.exists() else {}
    except (OSError, json.JSONDecodeError):
        settings = {}

    settings["effortLevel"] = "max"
    # Equivalent of --dangerously-skip-permissions; the second key suppresses
    # the one-time confirmation so the unattended tmux window doesn't block.
    settings.setdefault("permissions", {})["defaultMode"] = "bypassPermissions"
    settings["skipDangerousModePermissionPrompt"] = True
    # Auto-approve the repo's checked-in .mcp.json servers (phrocs, traffic-sim)
    # so they load on the unattended boot without the approval prompt.
    settings["enableAllProjectMcpServers"] = True

    settings_path.write_text(json.dumps(settings, indent=2) + "\n")


def copy_claude_auth(uid: int, gid: int) -> None:
    """Copy Claude Code auth files into the sandbox home."""
    claude_dir = SANDBOX_HOME / ".claude"
    claude_dir.mkdir(parents=True, exist_ok=True)

    for name in (".credentials.json", "settings.json", "settings.local.json", "CLAUDE.md"):
        src = Path(f"/tmp/claude-auth/{name}")
        if src.exists():
            shutil.copy2(src, claude_dir / name)

    _strip_host_permissions(claude_dir / "settings.json")
    _strip_host_permissions(claude_dir / "settings.local.json")
    _force_sandbox_claude_defaults(claude_dir / "settings.json")

    # bin/sandbox ships a finalized .claude.json (installMethod stripped, MCP
    # servers merged under the /workspace project); just place it. is_file()
    # skips the /dev/null mount used when there's nothing to ship.
    src = Path("/tmp/claude-auth.json")
    if src.is_file():
        shutil.copy2(src, SANDBOX_HOME / ".claude.json")

    run(["chown", "-R", f"{uid}:{gid}", str(SANDBOX_HOME)])


def copy_host_gitconfig(uid: int, gid: int) -> None:
    """Copy host ~/.gitconfig and rewrite signing paths for the container.

    Copied (not mounted) so ``git config --global`` works without leaking
    back to the host. Signing paths (user.signingkey, allowedSignersFile)
    are rewritten to their in-container mount locations.
    """
    src = Path("/tmp/host-gitconfig")
    if not src.is_file():
        return
    dst = SANDBOX_HOME / ".gitconfig"
    shutil.copy2(src, dst)
    run(["chown", f"{uid}:{gid}", str(dst)])

    # Rewrite host-absolute paths to in-container mount paths.
    # Compose mounts /dev/null when unconfigured, so is_file() is correct.
    rewrites = {
        "user.signingkey": Path("/tmp/host-git-signingkey"),
        "gpg.ssh.allowedSignersFile": Path("/tmp/host-allowed-signers"),
    }
    for config_key, mount_path in rewrites.items():
        if mount_path.is_file():
            run(["git", "config", "--file", str(dst), config_key, str(mount_path)])


def root_phase() -> None:
    uid = int(os.environ.get("SANDBOX_UID", "1000"))
    gid = int(os.environ.get("SANDBOX_GID", "1000"))

    SANDBOX_HOME.mkdir(parents=True, exist_ok=True)
    Path("/tmp/sandbox-cache").mkdir(parents=True, exist_ok=True)
    # Start interactive shells in the workspace, not the dotfiles-only home dir.
    (SANDBOX_HOME / ".bashrc").write_text(f"cd {WORKSPACE}\n")
    run(["chown", f"{uid}:{gid}", str(SANDBOX_HOME), "/tmp/sandbox-cache"])

    create_sandbox_user(uid, gid)

    # Must run before export_environment snapshots os.environ into /etc/profile.d.
    # The user's sandbox.env secrets are already in os.environ (compose env_file),
    # so export_environment propagates them to every shell and the Claude tree.
    _configure_ssh_agent_env()
    export_environment(uid, gid)

    # cache-init doesn't need SSH, auth, or IDE config.
    if os.environ.get("SANDBOX_MODE") != "cache-init":
        start_sshd(uid, gid)
        configure_user_ssh(uid, gid)
        copy_claude_auth(uid, gid)
        copy_host_gitconfig(uid, gid)

    # Re-exec as the sandbox user.
    os.execvp("gosu", ["gosu", f"{uid}:{gid}", sys.executable, __file__, *sys.argv[1:]])


# ---------------------------------------------------------------------------
# User phase — runs as the sandbox UID
# ---------------------------------------------------------------------------


def _run_captured(label: str, cmd: list[str], **kwargs) -> None:
    """Run a subprocess, suppressing output on success, dumping it on failure."""
    result = subprocess.run(cmd, capture_output=True, text=True, **kwargs)
    if result.returncode == 0:
        return
    info(f"ERROR: {label} failed:")
    for line in (result.stdout or "").strip().splitlines():
        info(f"  {line}")
    for line in (result.stderr or "").strip().splitlines():
        info(f"  {line}")
    raise subprocess.CalledProcessError(result.returncode, cmd)


def install_python_deps() -> None:
    info("Started: uv sync...")
    _write_status("installing python deps")
    _run_captured("uv sync", ["uv", "sync", "--no-editable"])
    info("Finished: uv sync.")
    # Make insightscli available — normally done by flox on-activate.sh.
    insightscli_link = Path("/cache/python/bin/insightscli")
    insightscli_link.unlink(missing_ok=True)
    insightscli_link.symlink_to("/workspace/bin/insightscli")
    phrocs_link = WORKSPACE / "bin/phrocs"
    if not phrocs_link.exists():
        phrocs_link.symlink_to("/usr/local/bin/phrocs")


def install_node_deps() -> None:
    info("Started: pnpm install...")
    _write_status("installing node deps")
    # CI=1 suppresses interactive prompts. --no-frozen-lockfile is needed
    # because the sandbox branch may have different dependencies than the cache.
    _run_captured(
        "pnpm install",
        ["pnpm", "install", "--no-frozen-lockfile"],
        env={**os.environ, "CI": "1"},
    )
    info("Finished: pnpm install.")


def fetch_rust_crates() -> None:
    """Pre-fetch Rust crate sources so concurrent cargo builds don't race."""
    info("Started: cargo fetch...")
    _write_status("fetching rust crates")
    _run_captured("cargo fetch", ["cargo", "fetch"], cwd=str(WORKSPACE / "rust"))
    info("Finished: cargo fetch.")


def ensure_demo_data() -> None:
    """Generate demo data on first boot; skip if already present."""
    result = run_quiet(
        [
            "psql",
            "-h",
            "db",
            "-U",
            "insights",
            "-d",
            "insights",
            "-tAc",
            "SELECT 1 FROM insights_user WHERE email='[email protected]' LIMIT 1",
        ]
    )
    if result.stdout.strip() == b"1":
        info("Demo data already present, skipping generation.")
    else:
        info("Generating demo data (first boot)...")
        run(["python", "manage.py", "generate_demo_data"])


def install_geoip() -> None:
    """Symlink GeoIP database from the Docker image into the worktree.

    The .json sidecar prevents bin/start's download-mmdb from clobbering
    the symlink and re-fetching on every boot.
    """
    share = WORKSPACE / "share"
    share.mkdir(parents=True, exist_ok=True)

    mmdb = share / "GeoLite2-City.mmdb"
    if not mmdb.exists() and not mmdb.is_symlink():
        mmdb.symlink_to("/share/GeoLite2-City.mmdb")

    sidecar = share / "GeoLite2-City.json"
    if not sidecar.exists():
        sidecar.write_text(f'{{ "date": "{time.strftime("%Y-%m-%d")}" }}\n')


def create_kafka_topics() -> None:
    info("Pre-creating Kafka topics...")
    for topic in ("datastore_events_json",):
        if run_quiet(["rpk", "topic", "describe", topic, "--brokers", "kafka:9092"]).returncode != 0:
            run(["rpk", "topic", "create", topic, "--brokers", "kafka:9092", "-p", "1", "-r", "1"])


def generate_mprocs_config() -> None:
    info("Generating mprocs config...")
    config_dir = WORKSPACE / ".insights/.generated"
    config_dir.mkdir(parents=True, exist_ok=True)
    config_file = config_dir / "mprocs.yaml"

    # Seed intents on first boot; on restart, reuse the saved config.
    needs_seed = not config_file.exists()
    if not needs_seed:
        needs_seed = "_insights:" not in config_file.read_text()

    if needs_seed:
        intents = os.environ.get("SANDBOX_INTENTS", "product_analytics")
        info(f"Seeding intents: {intents}")
        # Write seed YAML directly — system python doesn't have PyYAML or
        # insightscli installed, and the structure is simple enough to emit by hand.
        lines = ["_insights:", "  intents:"]
        for intent in intents.split(","):
            lines.append(f"  - {intent.strip()}")
        # Migrations already ran in the entrypoint; skip autostart to avoid
        # redundant Django startup overhead in phrocs.
        lines.append("  skip_autostart:")
        lines.append("  - migrate-postgres")
        lines.append("  - migrate-datastore")
        lines.append("  - migrate-persons-db")
        lines.append("procs: {}")
        config_file.write_text("\n".join(lines) + "\n")

    _run_captured("insightscli dev:generate", ["insightscli", "dev:generate"])


def _read_jetbrains_data_dir_name() -> str:
    """Read the config directory name from the installed IDE's product-info.json."""
    import json

    product_info = Path("/opt/idea/product-info.json")
    data = json.loads(product_info.read_text())
    data_dir_name = data["dataDirectoryName"]
    return data_dir_name


def setup_jetbrains_background() -> None:
    """Register JetBrains IDE backend in a background process."""
    idea_script = Path("/opt/idea/bin/remote-dev-server.sh")
    if not idea_script.exists():
        return

    data_dir_name = _read_jetbrains_data_dir_name()

    pid = os.fork()
    if pid != 0:
        return  # Parent continues

    # The child must never return to the caller. If an exception unwinds past
    # os.fork(), run_setup's "JetBrains is non-fatal" guard catches it and the
    # forked child then falls through and re-runs the rest of run_setup,
    # spawning a duplicate phrocs window. Always terminate the child here.
    try:
        _run_jetbrains_child(idea_script, data_dir_name)
    except Exception:
        traceback.print_exc()
    os._exit(0)


def _run_jetbrains_child(idea_script: Path, data_dir_name: str) -> None:
    # Child: redirect stdio to log so late writes don't clobber the tmux prompt.
    log_path = "/tmp/sandbox-jetbrains.log"
    log_fd = os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o644)
    os.dup2(log_fd, 1)
    os.dup2(log_fd, 2)
    os.close(log_fd)

    os.environ["JAVA_TOOL_OPTIONS"] = f"-Duser.home={SANDBOX_HOME}"

    info(f"Registering {data_dir_name} for Gateway (background)...")
    result = subprocess.run(
        ["remote-dev-server.sh", "registerBackendLocationForGateway"],
        executable=str(idea_script),
        env={**os.environ, "REMOTE_DEV_NON_INTERACTIVE": "1"},
    )
    if result.returncode != 0:
        print(f"[{_ts()}] ERROR: IDE backend registration failed (exit {result.returncode}).")  # noqa: T201

    # Install Python plugin if not already present
    jetbrains_dir = SANDBOX_HOME / ".local/share/JetBrains"
    has_python_plugin = any(jetbrains_dir.rglob("python*")) if jetbrains_dir.exists() else False
    if not has_python_plugin:
        info("Installing Python plugin...")
        result = subprocess.run(
            [
                "remote-dev-server.sh",
                "installPlugins",
                "PythonCore",
                "Pythonid",
                "intellij.python.dap.plugin",
                "com.intellij.python.django",
            ],
            executable=str(idea_script),
            env={**os.environ, "REMOTE_DEV_NON_INTERACTIVE": "1"},
        )
        if result.returncode != 0:
            print(f"[{_ts()}] ERROR: Plugin installation failed (exit {result.returncode}).")  # noqa: T201

    # Configure Python SDK using the IDE's own config directory name
    idea_config = SANDBOX_HOME / ".config/JetBrains" / data_dir_name
    write_file_if_missing(
        idea_config / "options/jdk.table.xml",
        dedent("""\
            <application>
              <component name="ProjectJdkTable">
                <jdk version="2">
                  <name value="Python 3.13 (sandbox)" />
                  <type value="Python SDK" />
                  <homePath value="/cache/python/bin/python3" />
                  <roots>
                    <classPath><root type="composite" /></classPath>
                    <sourcePath><root type="composite" /></sourcePath>
                  </roots>
                  <additional />
                </jdk>
              </component>
            </application>
        """),
    )

    write_file_if_missing(
        WORKSPACE / ".idea/modules.xml",
        dedent("""\
            <?xml version="1.0" encoding="UTF-8"?>
            <project version="4">
              <component name="ProjectModuleManager">
                <modules>
                  <module fileurl="file://$PROJECT_DIR$/.idea/insights.iml" filepath="$PROJECT_DIR$/.idea/insights.iml" />
                </modules>
              </component>
            </project>
        """),
    )

    write_file_if_missing(
        WORKSPACE / ".idea/misc.xml",
        dedent("""\
            <?xml version="1.0" encoding="UTF-8"?>
            <project version="4">
              <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.13 (sandbox)" project-jdk-type="Python SDK" />
              <component name="TestRunnerService">
                <option name="PROJECT_TEST_RUNNER" value="py.test" />
              </component>
            </project>
        """),
    )

    info(f"{data_dir_name} backend ready")


def _setup_user_env() -> None:
    """Shared env setup for user_phase and cache_init_phase."""
    PROGRESS_FILE.write_text("")
    os.environ.update(
        {
            "HOME": str(SANDBOX_HOME),
            # Mirror export_environment so the tmux/claude tree gets VIRTUAL_ENV.
            "VIRTUAL_ENV": "/cache/python",
            "UV_CACHE_DIR": "/cache/uv",
            "UV_LINK_MODE": "copy",
            "XDG_CACHE_HOME": "/tmp/sandbox-cache",
            "COREPACK_ENABLE_AUTO_PIN": "0",
            "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0",
            "CARGO_TARGET_DIR": "/cache/cargo-target",
            "npm_config_store_dir": "/cache/pnpm",
        }
    )
    _configure_ssh_agent_env()
    Path("/cache/cargo-target").mkdir(parents=True, exist_ok=True)
    os.chdir(WORKSPACE)
    # pushInsteadOf (not insteadOf) so anonymous HTTPS fetches still work
    # for cargo deps that clone without an SSH agent during cache-init.
    run(["git", "config", "--global", "[email protected]:.pushInsteadOf", "https://github.com/"])


def _run_parallel(tasks: dict[str, Callable[[], None]]) -> None:
    """Run named tasks in parallel, re-raising the first failure."""
    with ThreadPoolExecutor() as pool:
        futures = {pool.submit(fn): name for name, fn in tasks.items()}
        for future in as_completed(futures):
            name = futures[future]
            try:
                future.result()
            except Exception as e:
                print(f"[{_ts()}] ERROR: {name} failed: {e}", flush=True)  # noqa: T201
                raise


def user_phase() -> None:
    """PID 1: launch tmux (claude + setup windows), then poll-block on the session."""
    _setup_user_env()
    install_geoip()

    _write_status("booting")

    # Window 1: claude (launched immediately, may lack insightscli until uv sync finishes)
    # Window 2: setup (deps, migrations, then spawns phrocs in window 3)
    tmux = ["tmux", "-L", "sandbox"]
    run(
        [
            *tmux,
            "-f",
            "/etc/tmux.sandbox.conf",
            "new-session",
            "-d",
            "-s",
            "insights",
            "-c",
            "/workspace",
            "-n",
            "claude",
            "claude",
        ]
    )
    run(
        [
            *tmux,
            "new-window",
            "-t",
            "insights:",
            "-n",
            "setup",
            "-e",
            "SANDBOX_MODE=setup",
            f"{sys.executable} {__file__}",
        ]
    )
    run([*tmux, "select-window", "-t", "insights:claude"])

    # Block on tmux session. Exit non-zero when it dies so compose's
    # restart: on-failure brings the container back.
    while (
        subprocess.run(
            [*tmux, "has-session", "-t", "insights"],
            capture_output=True,
        ).returncode
        == 0
    ):
        time.sleep(2)
    sys.exit(1)


def _port_open(host: str, port: int) -> bool:
    try:
        with socket.create_connection((host, port), timeout=2):
            return True
    except OSError:
        return False


def _datastore_ready() -> bool:
    """True once Datastore answers a trivial query over HTTP."""
    try:
        # nosemgrep: python.lang.security.audit.dynamic-urllib-use-detected.dynamic-urllib-use-detected -- URL is a hardcoded constant, no user-controlled value
        with urllib.request.urlopen("http://datastore:8123/", data=b"SELECT 1", timeout=3) as resp:
            return resp.read().strip() == b"1"
    except OSError:
        return False


def _restore_status() -> str | None:
    """What the host's detached restore last wrote: "done", a failure reason, or
    None if it hasn't written yet. The setup window waits for "done" (cache
    creates) and aborts on a failure reason instead of timing out.
    """
    if RESTORE_STATUS_FILE.exists():
        return RESTORE_STATUS_FILE.read_text().strip() or None
    return None


def _restore_failed_reason() -> str | None:
    """The failure reason if the restore reported one, else None ("done" is success)."""
    status = _restore_status()
    return status if status and status != "done" else None


def _wait_for(
    label: str,
    ready: Callable[[], bool],
    timeout: int = 300,
    fail_check: Callable[[], str | None] | None = None,
) -> None:
    """Poll `ready` until true, updating the status bar; raise on timeout.

    `fail_check` returns a reason string to abort early. It's checked *after*
    `ready`, so a restore that already signalled success wins over a late failure.
    """
    _write_status(f"waiting for {label}")
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if ready():
            return
        if fail_check and (reason := fail_check()):
            raise RuntimeError(reason)
        time.sleep(2)
    raise RuntimeError(f"{label} not ready within {timeout}s")


def _wait_for_databases() -> None:
    """Block until infra is ready before migrating.

    Postgres restores before its container boots, so a reachable db is already
    restored. Datastore restores onto a *running* server, and the database can
    appear mid-restore — so for a cache-backed create we wait for the "done" status
    the host writes once its synchronous RESTORE has fully landed, not on a probe.
    """
    _wait_for(
        "Postgres",
        lambda: run_quiet(["psql", "-h", "db", "-U", "insights", "-d", "insights", "-tAc", "SELECT 1"]).returncode == 0,
        fail_check=_restore_failed_reason,
    )
    if os.environ.get("SANDBOX_USE_CACHE") == "1":
        _wait_for("Datastore cache restore", lambda: _restore_status() == "done", fail_check=_restore_failed_reason)
    else:
        _wait_for("Datastore", _datastore_ready, fail_check=_restore_failed_reason)
    _wait_for("Kafka", lambda: _port_open("kafka", 9092), fail_check=_restore_failed_reason)


def run_setup() -> None:
    """Tmux window 2: install deps, migrate, seed data, spawn phrocs, exec bash.

    On failure, writes "SETUP FAILED" to the status bar and drops into bash
    so the user can diagnose. Claude keeps running in window 1 either way.
    """
    _setup_user_env()
    # The status file can linger in /tmp across a stop/start; clear it before
    # waiting so a stale "done"/failure from a previous boot can't trip this one.
    # On a cache create the host worker writes the fresh status later.
    RESTORE_STATUS_FILE.unlink(missing_ok=True)
    _write_status("setup starting")

    try:
        # Deps need no database, so install them while the host brings up and
        # restores infra in the background.
        _run_parallel(
            {
                "python deps": install_python_deps,
                "node deps": install_node_deps,
                "rust crates": fetch_rust_crates,
            }
        )

        # Gate DB work on infra being ready (and the cache restore having landed)
        # so we never migrate a half-restored database.
        _wait_for_databases()
        create_kafka_topics()
        _write_status("running migrations")
        run(["python", "manage.py", "sandbox_migrate", "--progress-file", str(PROGRESS_FILE)])
        _write_status("seeding demo data")
        ensure_demo_data()

        # generate_mprocs_config needs the insightscli symlink created by install_python_deps.
        generate_mprocs_config()

        # Forks internally — must run after ThreadPoolExecutor is closed.
        try:
            setup_jetbrains_background()
        except Exception as e:
            # Loud but non-fatal — IDE setup failing should not prevent the sandbox from booting.
            print(f"[{_ts()}] ERROR: JetBrains IDE setup failed: {e}", flush=True)  # noqa: T201

        lock = WORKSPACE / "bin/start.lock"
        lock.unlink(missing_ok=True)

        _write_status("C-b 1 claude  C-b 2 setup  C-b 3 phrocs  C-b d detach  mouse enabled")

        # Spawn phrocs in its own tmux window.
        run(["tmux", "-L", "sandbox", "new-window", "-t", "insights:", "-n", "phrocs", "bin/start --phrocs"])

        print(  # noqa: T201
            "\nSetup complete — see status bar for keybinding hints.\n",
            flush=True,
        )
    except Exception as e:
        traceback.print_exc()
        _write_status("!! SETUP FAILED. Ctrl-b 1 returns to claude")
        print(  # noqa: T201
            f"\n\n!!! Setup failed: {e}"
            "\n!!! Traceback above. This window is now a bash shell; claude is still running in window 1."
            "\n!!! Fix the issue, then `tmux new-window bin/start --phrocs` when ready.\n",
            flush=True,
        )
        # Pull the user's view to this window so the failure can't be missed,
        # but leave claude running in window 1, since they may already be working.
        # select before rename, so the "setup" target still resolves.
        tmux = ["tmux", "-L", "sandbox"]
        subprocess.run([*tmux, "select-window", "-t", "insights:setup"])
        subprocess.run([*tmux, "rename-window", "-t", "insights:setup", "⚠ SETUP FAILED"])

    # Exec into bash so this window stays usable (both success and failure).
    os.execvp("bash", ["bash", "-l"])


# ---------------------------------------------------------------------------
# Cache-init phase — one-shot cache warming for `bin/sandbox create`
# ---------------------------------------------------------------------------


def cache_init_phase() -> None:
    """Warm all shared caches (uv, pnpm store, cargo target) and populate databases.

    Runs inside `compose run` during `bin/sandbox create` when no database cache
    exists. Exits when done so the host script can snapshot Postgres + Datastore.
    """
    _setup_user_env()

    def install_python_and_migrate() -> None:
        install_python_deps()
        info("Running migrations...")
        run(["python", "manage.py", "sandbox_migrate", "--progress-file", str(PROGRESS_FILE)])
        info("Generating demo data (this takes a few minutes)...")
        run(["python", "manage.py", "generate_demo_data", "--skip-flag-sync"])

    def build_rust() -> None:
        info("Pre-building Rust workspace...")
        try:
            run(["cargo", "build", "--workspace"], cwd=str(WORKSPACE / "rust"))
            info("Finished: cargo build.")
        except subprocess.CalledProcessError as e:
            # Non-fatal: even a partial build warms the cargo cache for subsequent
            # boots. Cargo's stderr already streamed to the terminal so the error
            # details are visible above; we just need to preserve the exit code
            # for debugging.
            info(f"Rust pre-build failed (exit {e.returncode}), continuing (cargo cache still warmed).")

    # Python path is usually the critical path (migrations on a fresh DB take
    # several minutes) so we want node + rust overlapping with it.
    _run_parallel(
        {
            "python + migrations + demo data": install_python_and_migrate,
            "node deps": install_node_deps,
            "rust build": build_rust,
        }
    )

    info("Cache init complete.")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    mode = os.environ.get("SANDBOX_MODE")
    uid = int(os.environ.get("SANDBOX_UID", "1000"))
    if os.getuid() == 0 and uid != 0:
        root_phase()
    elif mode == "cache-init":
        cache_init_phase()
    elif mode == "setup":
        run_setup()
    else:
        user_phase()