bot-bottle/bot_bottle/backend/smolmachines/sidecar_bundle.py

"""Per-bottle sidecar bundle bringup for the smolmachines backend.

The sidecar bundle runs as its own smolVM. The agent VM reaches
bundle daemons through host-loopback ports published by that sidecar
VM and wrapped by per-bottle address-bound forwarders."""

from __future__ import annotations

import os
import socket
from dataclasses import dataclass, field
from pathlib import Path
from typing import Sequence

from ...log import warn
from ..docker import util as docker_mod
from ..docker.sidecar_bundle import (
    SIDECAR_BUNDLE_DOCKERFILE,
    SIDECAR_BUNDLE_IMAGE,
)
from . import smolvm as _smolvm


_REPO_DIR = str(Path(__file__).resolve().parent.parent.parent.parent)


def bundle_network_name(slug: str) -> str:
    """`bot-bottle-bundle-<slug>` — distinct from the docker
    backend's `bot-bottle-net-<slug>` so a smolmachines bottle
    and a docker bottle for the same agent don't collide on
    network name."""
    return f"bot-bottle-bundle-{slug}"


def bundle_container_name(slug: str) -> str:
    """`bot-bottle-sidecars-<slug>` — same name shape the docker
    backend uses for the bundle (PRD 0024 chunk 5). The dashboard's
    prefix-based discovery covers both backends with one filter."""
    return f"bot-bottle-sidecars-{slug}"


def bundle_machine_name(slug: str) -> str:
    """Sidecar smolVM machine name."""
    return bundle_container_name(slug)


@dataclass(frozen=True)
class BundleLaunchSpec:
    """Everything `start_bundle` needs to bring up one bundle
    container. Populated by chunk-2d's launch flow from the inner
    Plans the prepare step already produces."""

    slug: str
    network_name: str
    subnet: str
    gateway: str
    bundle_ip: str
    image: str = SIDECAR_BUNDLE_IMAGE
    # Daemon subset CSV for BOT_BOTTLE_SIDECAR_DAEMONS. The
    # supervisor inside the bundle reads it to skip
    # bottle-irrelevant daemons (e.g. supervise=False bottles).
    daemons_csv: str = "egress"
    # Plain "KEY=VALUE" strings + "KEY" bare names. Bare names inherit
    # from the host env passed to the sidecar VM launch.
    environment: Sequence[str] = field(default_factory=tuple)
    # (host_path, container_path, read_only) bind mounts.
    volumes: Sequence[tuple[str, str, bool]] = field(default_factory=tuple)
    # Container ports to publish on `publish_host_ip`, random
    # host-side port per entry. The smolvm guest's TSI talks via
    # macOS networking, so docker container IPs (192.168.x.x in
    # the daemon's bridge) aren't directly reachable from the
    # guest — host-loopback port-forwards are. Egress's port
    # is bundle-internal and never published.
    ports_to_publish: Sequence[int] = field(default_factory=tuple)
    # Loopback IP to bind published ports against. Per-bottle
    # loopback aliases (`127.0.0.16` etc., added via sudo
    # ifconfig lo0 alias) narrow the TSI allowlist so a bottle
    # can't reach other bottles' (or other host services') ports
    # via 127.0.0.1.
    publish_host_ip: str = "127.0.0.1"


@dataclass(frozen=True)
class BundleVmLaunch:
    raw_ports: dict[int, int]


def ensure_bundle_image(image: str = SIDECAR_BUNDLE_IMAGE) -> None:
    """Build the sidecar bundle image before `docker run`.

    The Docker backend gets this for free from compose's `build:`
    stanza. smolmachines starts the bundle with plain `docker run`,
    so without an explicit build a first launch tries to pull the
    local-only `bot-bottle-sidecars:latest` tag from a registry.
    """
    docker_mod.build_image(
        image,
        _REPO_DIR,
        dockerfile=SIDECAR_BUNDLE_DOCKERFILE,
    )


def allocate_raw_host_ports(container_ports: Sequence[int]) -> dict[int, int]:
    """Reserve candidate host loopback ports for smolVM `-p HOST:GUEST`.

    The sockets are closed before smolVM binds them, so this remains
    best-effort. smolVM failing to bind a selected port is fatal in
    the caller's launch path."""
    out: dict[int, int] = {}
    for container_port in container_ports:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.bind(("127.0.0.1", 0))
            out[container_port] = int(sock.getsockname()[1])
    return out


def start_bundle_vm(
    spec: BundleLaunchSpec,
    *,
    from_path: Path,
    host_env: dict[str, str] | None = None,
) -> BundleVmLaunch:
    """Create and start the sidecar bundle as a smolVM.

    smolVM's own published ports are raw host-loopback ports. The
    launch flow wraps them with per-bottle address-bound forwarders
    before exposing anything to the agent VM."""
    raw_ports = allocate_raw_host_ports(spec.ports_to_publish)
    effective_host_env = host_env if host_env is not None else os.environ
    env: dict[str, str] = {"BOT_BOTTLE_SIDECAR_DAEMONS": spec.daemons_csv}
    for entry in spec.environment:
        name, sep, value = entry.partition("=")
        if sep:
            env[name] = value
        elif entry in effective_host_env:
            env[entry] = effective_host_env[entry]
    name = bundle_machine_name(spec.slug)
    # Destroy any leftover machine from a previous run that didn't
    # clean up (e.g. crash, interrupted teardown).
    try:
        _smolvm.machine_stop(name)
    except _smolvm.SmolvmError:
        pass
    try:
        _smolvm.machine_delete(name)
    except _smolvm.SmolvmError:
        pass
    _smolvm.machine_create(
        name,
        from_path=from_path,
        net=True,
        env=env,
        volumes=spec.volumes,
        ports=tuple(
            (host_port, guest_port)
            for guest_port, host_port in raw_ports.items()
        ),
    )
    _smolvm.machine_start(name)
    _smolvm.wait_exec_ready(name)
    return BundleVmLaunch(raw_ports=raw_ports)


def stop_bundle_vm(slug: str) -> None:
    """Best-effort sidecar VM teardown."""
    name = bundle_machine_name(slug)
    try:
        _smolvm.machine_stop(name)
    except _smolvm.SmolvmError as exc:
        warn(f"smolvm machine stop {name} failed: {exc}")
    try:
        _smolvm.machine_delete(name)
    except _smolvm.SmolvmError as exc:
        warn(f"smolvm machine delete {name} failed: {exc}")