refactor(env): run stock MCR images, prepare the sandbox on the fly

.github/workflows/test.yml

@@ -141,19 +141,6 @@ jobs:
         with:
           python-version: "3.12"
 
-      - name: Set up Docker Buildx
-        if: matrix.test-type != 'unit'
-        uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1
-
-      - name: Cache Docker layers
-        if: matrix.test-type != 'unit'
-        uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
-        with:
-          path: /tmp/.buildx-cache
-          key: ${{ runner.os }}-buildx-${{ hashFiles('src/microbots/environment/local_docker/image_builder/Dockerfile') }}
-          restore-keys: |
-            ${{ runner.os }}-buildx-
-
       - name: Cache pip dependencies
         uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
@@ -194,19 +181,10 @@ jobs:
         if: matrix.test-type == 'unit'
         run: npm test --prefix azure-pipelines/MicrobotsLogAnalyzerTask
 
-      - name: Build Docker images for integration tests
+      - name: Pull the stock sandbox image for integration tests
         if: matrix.test-type != 'unit'
-        run: |
-          # Build the shell server image needed for Docker tests
-          docker buildx build \
-            --cache-from type=local,src=/tmp/.buildx-cache \
-            --cache-to type=local,dest=/tmp/.buildx-cache-new,mode=max \
-            --load \
-            -f src/microbots/environment/local_docker/image_builder/Dockerfile \
-            -t kavyasree261002/shell_server:latest .
-          # Prevent cache from growing indefinitely
-          rm -rf /tmp/.buildx-cache
-          mv /tmp/.buildx-cache-new /tmp/.buildx-cache
+        # Sandboxes run this unmodified MCR image and are prepared on the fly.
+        run: docker pull mcr.microsoft.com/devcontainers/python:3.11
 
       - name: Check disk space before ollama installation
         if: matrix.test-type == 'ollama_local'

docs/api-reference/microbots/environment/local_docker/images.md

@@ -0,0 +1 @@
+::: microbots.environment.local_docker.images

docs/api-reference/microbots/environment/local_docker/shell_server/ShellCommunicator.md

@@ -0,0 +1 @@
+::: microbots.environment.local_docker.shell_server.ShellCommunicator

docs/api-reference/microbots/environment/local_docker/shell_server/dockerShell.md

@@ -0,0 +1 @@
+::: microbots.environment.local_docker.shell_server.dockerShell

docs/getting-started/microbots-logging.md

@@ -38,7 +38,8 @@ Working directory creation, the host-to-container volume mapping, container star
 ```log
 INFO:microbots.environment.local_docker.LocalDockerEnvironment:🗂️  Created working directory at /home/user/MICROBOTS_WORKDIR_<id>
 INFO:microbots.environment.local_docker.LocalDockerEnvironment:📦 Volume mapping: /home/user/microbots-introduction/code → /ro/code
-INFO:microbots.environment.local_docker.LocalDockerEnvironment:🚀 Started container <container-id> with image microbots/shell_server:latest on host port <host-port>
+INFO:microbots.environment.local_docker.LocalDockerEnvironment:🚀 Started container <container-id> with image mcr.microsoft.com/devcontainers/python:3.11 on host port <host-port>
+INFO:microbots.environment.local_docker.LocalDockerEnvironment:🛠️  Prepared mcr.microsoft.com/devcontainers/python:3.11 for the shell server in 9.0s
 INFO:microbots.environment.local_docker.LocalDockerEnvironment:🔒 Set up overlay mount for read-only directory at /workdir/code
 INFO:microbots.environment.local_docker.LocalDockerEnvironment:✅ Successfully copied /home/user/microbots-introduction/code/build.log to container:/var/log
 ```

mkdocs.yml

@@ -106,9 +106,10 @@ nav:
       - Environment: api-reference/microbots/environment/Environment.md
       - local_docker:
         - LocalDockerEnvironment: api-reference/microbots/environment/local_docker/LocalDockerEnvironment.md
-        - image_builder:
-          - dockerShell: api-reference/microbots/environment/local_docker/image_builder/dockerShell.md
-          - ShellCommunicator: api-reference/microbots/environment/local_docker/image_builder/ShellCommunicator.md
+        - images: api-reference/microbots/environment/local_docker/images.md
+        - shell_server:
+          - dockerShell: api-reference/microbots/environment/local_docker/shell_server/dockerShell.md
+          - ShellCommunicator: api-reference/microbots/environment/local_docker/shell_server/ShellCommunicator.md
     - extras:
       - mount: api-reference/microbots/extras/mount.md
     - llm:

src/microbots/environment/local_docker/LocalDockerEnvironment.py

@@ -1,7 +1,9 @@
+import io
 import logging
 import os
 import shlex
 import subprocess
+import tarfile
 import time
 from pathlib import Path
 from typing import Optional
@@ -10,18 +12,63 @@
 import requests
 
 from microbots.environment.Environment import CmdReturn, Environment
+from microbots.environment.local_docker import shell_server
+from microbots.environment.local_docker.images import DEFAULT_IMAGE
 from microbots.constants import DOCKER_WORKING_DIR, WORKING_DIR, PermissionLabels
 from microbots.extras.mount import Mount
 
 logger = logging.getLogger(__name__)
 
+# Where the shell server and its venv live inside the sandbox (root-owned).
+MICROBOTS_HOME = "/opt/microbots"
+_SHELL_SERVER_DIR = Path(shell_server.__file__).parent
+_SHELL_SERVER_FILES = ("dockerShell.py", "ShellCommunicator.py", "prepare.sh")
+_SHELL_SERVER_LOG = "/var/log/microbots/shell_server.log"
+# Preparation may install packages and pip-install the server's dependencies.
+_SHELL_SERVER_READY_TIMEOUT = 120
+
+
+def _shell_server_archive() -> bytes:
+    """Tar ``shell_server/`` for ``put_archive``, root-owned so the bot cannot modify it."""
+    buffer = io.BytesIO()
+    with tarfile.open(fileobj=buffer, mode="w") as tar:
+        entries = [("shell_server", _SHELL_SERVER_DIR)]
+        entries += [(f"shell_server/{name}", _SHELL_SERVER_DIR / name) for name in _SHELL_SERVER_FILES]
+        for arcname, path in entries:
+            info = tar.gettarinfo(str(path), arcname=arcname)
+            info.uid = info.gid = 0
+            info.uname = info.gname = "root"
+            info.mode = 0o755 if info.isdir() or arcname.endswith(".sh") else 0o644
+            if info.isfile():
+                with open(path, "rb") as f:
+                    tar.addfile(info, f)
+            else:
+                tar.addfile(info)
+    return buffer.getvalue()
+
 
 class LocalDockerEnvironment(Environment):
+    """Sandboxed Docker container that hosts the bot's shell.
+
+    Parameters
+    ----------
+    port : int
+        Host loopback port mapped to the in-container shell server.
+    folder_to_mount : Optional[Mount]
+        Host folder to expose inside the sandbox.
+    image : str
+        Container image to run, for example ``images.DEBIAN_IMAGE``.  Stock
+        images are prepared on the fly: the shell server is copied in and
+        its prerequisites are installed when the environment starts.  Its
+        distribution is detected from the container's ``/etc/os-release``
+        when needed.  Defaults to ``DEFAULT_IMAGE``.
+    """
+
     def __init__(
         self,
         port: int,
         folder_to_mount: Optional[Mount] = None,
-        image: str = "kavyasree261002/shell_server:latest",
+        image: str = DEFAULT_IMAGE,
     ):
 
         self.image = image
@@ -85,27 +132,31 @@ def start(self):
 
         self.container = self.client.containers.run(
             self.image,
+            # Keep the stock image idle; the shell server is started over the
+            # control plane once the image has been prepared. init forwards
+            # SIGTERM so the container stops promptly.
+            entrypoint=["sleep", "infinity"],
+            init=True,
             volumes=volumes_config,
             ports=port_mapping,
             detach=True,
-            working_dir="/app",
             # SYS_ADMIN is the narrowest grant that allows the overlayfs mount;
             # the docker-default AppArmor profile denies mount regardless of caps.
             cap_add=["SYS_ADMIN"],
             security_opt=["no-new-privileges:true", "apparmor=unconfined"],
-            environment={
-                "BOT_PORT": str(self.container_port),
-                "BOT_WORKDIR": DOCKER_WORKING_DIR,
-                **self._host_identity(),
-            },
         )
         logger.info(
             "🚀 Started container %s with image %s on host port %s",
             self.container.id[:12],
             self.image,
             self.port,
         )
-        time.sleep(2)  # Give some time for the server to start
+        try:
+            self._prepare_shell_server()
+            self._launch_shell_server()
+        except Exception:
+            self.stop()
+            raise
 
         if self.folder_to_mount and self.folder_to_mount.permission == PermissionLabels.READ_ONLY:
             self._setup_overlay_mount()
@@ -115,6 +166,79 @@ def start(self):
         else:
             self.execute("cd /")
 
+    def _shell_server_environment(self) -> dict:
+        """Environment for the shell server process inside the container."""
+        return {
+            "BOT_PORT": str(self.container_port),
+            "BOT_WORKDIR": DOCKER_WORKING_DIR,
+            **self._host_identity(),
+        }
+
+    def _prepare_shell_server(self):
+        """Copy the shell server into the container and prepare the image for it.
+
+        Runs over the root control plane: ``prepare.sh`` installs missing
+        prerequisites, creates the ``agent`` account and the server's venv.
+
+        Raises
+        ------
+        RuntimeError
+            If the files cannot be copied or preparation fails.
+        """
+        started = time.perf_counter()
+        ret = self.execute_privileged(f"mkdir -p {MICROBOTS_HOME}")
+        if ret.return_code != 0:
+            raise RuntimeError(f"Failed to create {MICROBOTS_HOME} in the sandbox: {ret.stderr}")
+
+        if not self.container.put_archive(MICROBOTS_HOME, _shell_server_archive()):
+            raise RuntimeError("Failed to copy the shell server into the sandbox")
+
+        identity = " ".join(f"{k}={v}" for k, v in self._host_identity().items())
+        ret = self.execute_privileged(
+            f"{identity} sh {MICROBOTS_HOME}/shell_server/prepare.sh".strip()
+        )
+        if ret.return_code != 0:
+            raise RuntimeError(
+                f"Failed to prepare image '{self.image}' for the shell server:\n"
+                f"{ret.stdout}\n{ret.stderr}"
+            )
+        logger.info(
+            "🛠️  Prepared %s for the shell server in %.1fs", self.image, time.perf_counter() - started
+        )
+
+    def _launch_shell_server(self):
+        """Start the shell server as root (it drops to ``agent`` itself) and
+        wait until it answers on the published port.
+
+        Raises
+        ------
+        RuntimeError
+            If the server does not become ready in time.
+        """
+        self.container.exec_run(
+            [
+                "sh",
+                "-c",
+                f"cd {MICROBOTS_HOME}/shell_server && exec {MICROBOTS_HOME}/venv/bin/python "
+                f"dockerShell.py >> {_SHELL_SERVER_LOG} 2>&1",
+            ],
+            user="root",
+            detach=True,
+            environment=self._shell_server_environment(),
+        )
+        deadline = time.monotonic() + _SHELL_SERVER_READY_TIMEOUT
+        while time.monotonic() < deadline:
+            try:
+                # Any HTTP answer (FastAPI rejects GET with 405) means it is up.
+                requests.get(f"http://localhost:{self.port}/", timeout=2)
+                return
+            except requests.exceptions.RequestException:
+                time.sleep(0.5)
+        log = self.execute_privileged(f"tail -n 50 {_SHELL_SERVER_LOG}").stdout
+        raise RuntimeError(
+            f"Shell server did not start within {_SHELL_SERVER_READY_TIMEOUT}s:\n{log}"
+        )
+
     @staticmethod
     def _host_identity() -> dict:
         """Run the bot's shell under the host uid so anything it writes to the

src/microbots/environment/local_docker/images.py

@@ -0,0 +1,18 @@
+"""Stock container images for LocalDockerEnvironment sandboxes.
+
+Sandboxes run unmodified images from the Microsoft Container Registry.  When
+an environment starts, Microbots copies its shell server into the container
+and prepares the image on the fly (``shell_server/prepare.sh``): it installs
+anything the server needs that is missing, creates the unprivileged ``agent``
+account and starts the server.
+
+Any image works if it provides ``sh`` and ``bash`` plus either ``python3`` or
+a package manager (``apt-get``, ``tdnf``, ``dnf`` or ``yum``) to install it.
+The sandbox's distribution is read from the container's ``/etc/os-release``
+when needed; it is never configured separately.
+"""
+
+# Debian with Python and a full developer toolset (git, curl, build tools).
+DEBIAN_IMAGE = "mcr.microsoft.com/devcontainers/python:3.11"
+
+DEFAULT_IMAGE = DEBIAN_IMAGE

src/microbots/environment/local_docker/shell_server/__init__.py

@@ -0,0 +1,12 @@
+"""Shell server that runs inside every LocalDockerEnvironment sandbox.
+
+LocalDockerEnvironment copies this directory into a stock container image
+when the environment starts, runs ``prepare.sh`` as root to install missing
+prerequisites and create the unprivileged ``agent`` account, then starts
+``dockerShell.py``.  It exposes a FastAPI endpoint that relays commands to a
+persistent bash session managed by ``ShellCommunicator.py``.
+
+The two Python files run inside the container as top-level scripts, so they
+must only use the Python standard library plus fastapi, uvicorn and pydantic
+(installed by ``prepare.sh`` into a root-owned venv).
+"""

src/microbots/environment/local_docker/image_builder/dockerShell.py → src/microbots/environment/local_docker/shell_server/dockerShell.py

@@ -23,7 +23,7 @@
 
 logger = logging.getLogger(__name__)
 
-# Must match the account created in the Dockerfile.
+# Must match the account created by prepare.sh.
 AGENT_USER = "agent"
 DEFAULT_AGENT_UID = 1001
 DEFAULT_AGENT_GID = 1001