fix(local-sandbox): hold each container a local run starts to the sandbox's cpu and memory

.agentenv/config.example.toml

@@ -11,7 +11,7 @@ image    = "local"   # local OCI registry (registry:2 at localhost:5000)
 secret   = "local"   # process env vars, or an optional flat YAML file
 
 [sandbox]
-default       = "local"   # envs run as containers on the host Docker daemon
+default       = "local"   # envs run as containers on the host Docker daemon, each held to the cpu and memory it was deployed with
 agent_default = "local"   # agents too
 
 [agents]

src/agent_env/a2a_agent/a2a_agent.py

@@ -13,7 +13,12 @@
 
 from agent_env.config import get_config
 from agent_env.entity_refs import EntityRef
-from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError, port_bindings
+from agent_env.providers.sandbox_providers.sandbox import (
+    ContainerLimits,
+    NetworkPolicy,
+    NetworkPolicyUnsupportedError,
+    port_bindings,
+)
 from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
 from agent_env.a2a_agent.object_transfer import (
     TRANSFER_TIMEOUT_SECONDS,
@@ -509,9 +514,11 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[
         agent_env = dict(merged_env) if trust_dir is None else {**LOCAL_TRUST_ENV, **merged_env}
         setup_script = ""
         network_flag = "".join(f"--add-host {entry} \\\n    " for entry in self._sandbox.extra_hosts)
+        limits = self._sandbox.container_limits
+        limit_flag = "".join(f"{shlex.quote(arg)} " for arg in limits.docker_args) + "\\\n    " if limits else ""
         if enable_docker:
             logger.info("enable_docker: starting rootless Docker-in-Docker sidecar for the agent (no host socket)")
-            setup_script = self._dind_setup_script()
+            setup_script = self._dind_setup_script(limits)
             network_flag += f"--network {_DIND_NETWORK} \\\n    "
             agent_env["DOCKER_HOST"] = f"tcp://{_DIND_CONTAINER}:{_DIND_PORT}"
 
@@ -527,7 +534,7 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[
 {setup_script}docker {"run -d" if trust_dir is None else "create"} \\
     --name {self._sandbox.container_name} \\
     {publish} \\
-    {network_flag}{env_str} \\
+    {network_flag}{limit_flag}{env_str} \\
     {image_name} > /dev/null
 {"sleep 2" if trust_dir is None else ""}
 """
@@ -536,20 +543,23 @@ async def _run_container(self, image_name: str, a2a_port: int, merged_env: dict[
             await start_trusting(self._sandbox, self._sandbox.container_name, trust_dir)
             await asyncio.sleep(2)
 
-    def _dind_setup_script(self) -> str:
+    def _dind_setup_script(self, limits: ContainerLimits | None = None) -> str:
         # Constants become shell vars so the body stays a raw string (no f-string
         # brace escaping). TLS off is safe: the tcp endpoint is only reachable on
-        # the VM-local bridge.
+        # the VM-local bridge. DIND_LIMITS is expanded unquoted, so it splits into
+        # flags, or into nothing when the sandbox holds its containers to nothing.
         env = (
             f"DIND_IMAGE={_DIND_IMAGE}\n"
             f"DIND_CONTAINER={_DIND_CONTAINER}\n"
             f"DIND_NETWORK={_DIND_NETWORK}\n"
             f"DIND_PORT={_DIND_PORT}\n"
+            f"DIND_LIMITS={shlex.quote(' '.join(limits.docker_args)) if limits else ''}\n"
         )
         return env + r"""docker network create "$DIND_NETWORK" >/dev/null 2>&1 || true
 docker rm -f "$DIND_CONTAINER" >/dev/null 2>&1 || true
 docker run -d \
     --name "$DIND_CONTAINER" \
+    $DIND_LIMITS \
     --privileged \
     --network "$DIND_NETWORK" \
     -e DOCKER_TLS_CERTDIR="" \

src/agent_env/env/envs/multi_env.py

@@ -612,6 +612,7 @@ async def _load_from_snapshot(self, snapshot) -> None:
             host_ips=self._sandbox.host_ips,
             extra_hosts=self._sandbox.extra_hosts,
             mcp_server_name=self._mcp_server_name or self.name,
+            container_limits=self._sandbox.container_limits,
         )
         compose_content = compose_content.replace(
             "    volumes:\n      - ./init-schemas.sql:/docker-entrypoint-initdb.d/init-schemas.sql\n", "",

src/agent_env/providers/env_providers/env_gateway_provider.py

@@ -34,7 +34,7 @@
 )
 from agent_env.providers.env_providers.env_provider import _SandboxEnvironmentProvider, _size_kwargs, _tool_names
 from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider
-from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, port_bindings
+from agent_env.providers.sandbox_providers.sandbox import ContainerLimits, Sandbox, VmSandbox, port_bindings
 from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SandboxProvider
 
 if TYPE_CHECKING:
@@ -213,6 +213,7 @@ def create_docker_compose(
         host_ips: tuple[str, ...] = (),
         extra_hosts: tuple[str, ...] = (),
         mcp_server_name: str | None = None,
+        container_limits: ContainerLimits | None = None,
     ) -> str:
         """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine.
 
@@ -224,6 +225,8 @@ def create_docker_compose(
             gateway_image: Docker image for the gateway.
             gateway_port: Port to expose for the gateway.
             website_configs: Optional list of website configurations (frontend + backend pairs).
+            container_limits: What every service's container is held to, for a sandbox that is not a VM
+                bounding them all (the local one); None adds nothing.
 
         Returns:
             docker-compose.yml content as a string.
@@ -504,6 +507,8 @@ def create_docker_compose(
             "  env-network:",
             "    driver: bridge",
         ])
+        if container_limits is not None:
+            lines = _with_container_limits(lines, container_limits)
         return "\n".join(lines)
 
     async def deploy(
@@ -710,6 +715,7 @@ async def _stand_up(_spec):
             host_ips=sandbox.host_ips,
             extra_hosts=sandbox.extra_hosts,
             mcp_server_name=mcp_server_name,
+            container_limits=sandbox.container_limits,
         )
         logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}")
 
@@ -1304,6 +1310,20 @@ def _gateway_topology(env: Env) -> _GatewayTopology:
 _SECRET_ENV_KEY_RE = re.compile(r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY)[A-Za-z0-9_]*)=.*$")
 
 
+def _with_container_limits(lines: list[str], limits: ContainerLimits) -> list[str]:
+    """``lines`` with ``limits`` under every service, each a two-space key in the top-level ``services`` section."""
+    limited: list[str] = []
+    section = None
+    for line in "\n".join(lines).split("\n"):  # an entry may hold several lines
+        limited.append(line)
+        if line and not line[0].isspace():
+            section = line.rstrip().removesuffix(":")
+        elif (section == "services" and len(line) > 2 and line.startswith("  ") and not line[2].isspace()
+              and line.rstrip().endswith(":")):
+            limited.extend(f"    {key}" for key in limits.compose_keys)
+    return limited
+
+
 def _redact_compose_secrets(compose_content: str) -> str:
     """Mask secret-valued env lines so the logged compose never carries a live secret."""
     return "\n".join(

src/agent_env/providers/sandbox_providers/local_sandbox.py

@@ -10,6 +10,7 @@
 import contextlib
 import functools
 import glob
+import json
 import logging
 import os
 import platform
@@ -28,7 +29,7 @@
 
 from agent_env import config
 from agent_env.attribution import Attribution
-from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, NetworkPolicy, VmSandbox
+from agent_env.providers.sandbox_providers.sandbox import SANDBOX_LABEL, ContainerLimits, NetworkPolicy, VmSandbox
 from agent_env.providers.sandbox_providers.sandbox_provider import (
     SANDBOX_MODE_CONTAINER,
     SANDBOX_MODE_VM,
@@ -74,6 +75,40 @@ def _runs_in_container(cmd: list[str]) -> bool:
 # Marker dropped in the work dir when this sandbox runs a container, so a later get_sandbox()
 # (post-run teardown reconstructs the sandbox from disk) knows it owns that container.
 _CONTAINER_MODE_MARKER = ".agent-container-mode"
+# What a sandbox's containers are held to, kept in its work dir for a handle get_sandbox() rebuilds.
+_LIMITS_FILE = ".container-limits.json"
+
+
+def _engine_cpus() -> int | None:
+    """How many CPUs the Docker engine has now (its VM's, on Docker Desktop or Rancher Desktop), or None when Docker
+    can't be asked. Not cached: the engine can be resized while a process runs."""
+    try:
+        return int(_docker("info", "--format", "{{.NCPU}}"))
+    except (OSError, RuntimeError, ValueError, subprocess.SubprocessError):
+        return None
+
+
+def _container_limits(cpu: float, memory: int) -> ContainerLimits:
+    """``cpu`` and ``memory`` as what a local sandbox's containers are held to. Docker refuses a container more CPUs
+    than its engine has, so cpu is held to that, with a warning; memory beyond the engine's is accepted."""
+    engine_cpus = _engine_cpus() if cpu > 1 else None  # every engine has at least one CPU
+    if engine_cpus is not None and cpu > engine_cpus:
+        logger.warning(
+            "Asked for %g CPUs, but the Docker engine has %d; holding this sandbox's containers to %d", cpu, engine_cpus, engine_cpus,
+        )
+        cpu = float(engine_cpus)
+    return ContainerLimits(cpus=cpu, memory_mib=memory)
+
+
+def _read_limits(work_dir: Path) -> ContainerLimits | None:
+    path = work_dir / _LIMITS_FILE
+    if not path.exists():
+        return None  # a sandbox created before limits were recorded
+    try:
+        return ContainerLimits.from_dict(json.loads(path.read_text()))
+    except (OSError, ValueError, KeyError, TypeError) as e:
+        logger.warning("Ignoring unreadable container limits at %s: %s", path, e)
+        return None
 
 
 def _local_sandbox_root_path() -> Path:
@@ -126,7 +161,7 @@ class LocalSandbox(VmSandbox):
     extra_hosts = _EXTRA_HOSTS
 
     def __init__(self, exposed_ports: list[int] | None = None, work_dir: Path | None = None, sandbox_id: str | None = None,
-                 port_map: dict[int, int] | None = None):
+                 port_map: dict[int, int] | None = None, container_limits: ContainerLimits | None = None):
         self.sandbox_id = sandbox_id or f"local-{uuid4().hex[:8]}"
         # Container port -> host port. Deployments share one host, so the reserved ports
         # can only be published once; the provider supplies a map to spread them. Callers
@@ -144,6 +179,17 @@ def __init__(self, exposed_ports: list[int] | None = None, work_dir: Path | None
         self._work_dir = work_dir or Path(
             tempfile.mkdtemp(prefix=f"agent-env-{self.sandbox_id}-", dir=_local_sandbox_root())
         )
+        if container_limits is not None:
+            self.container_limits = container_limits
+            # Unrecorded, a handle a later step rebuilds would start containers unheld, so a failed write fails here.
+            try:
+                (self._work_dir / _LIMITS_FILE).write_text(json.dumps(container_limits.to_dict()))
+            except OSError:
+                if work_dir is None:  # no handle is returned to tear down the folder just made
+                    shutil.rmtree(self._work_dir, ignore_errors=True)
+                raise
+        elif work_dir is not None:
+            self.container_limits = _read_limits(work_dir)
 
     def host_port(self, port: int) -> int:
         """The allocated host port for a published container port (identity if unmapped)."""
@@ -397,6 +443,7 @@ async def create_vm(
         # is I/O, and callers that only need a handle must not pay for it.
         sandbox = LocalSandbox(
             port_map={port: _free_host_port() for port in (exposed_ports or [])},
+            container_limits=await asyncio.to_thread(_container_limits, cpu, memory),
         )
         sandbox.network_policy = self.effective_network_policy(network_policy)
         return sandbox

src/agent_env/providers/sandbox_providers/sandbox.py

@@ -134,6 +134,35 @@ def from_dict(cls, data: dict) -> "NetworkPolicy":
         )
 
 
+@dataclass(frozen=True)
+class ContainerLimits:
+    """The cpu (cores) and memory (MiB) one container started on a sandbox's host is held to."""
+
+    cpus: float
+    memory_mib: int
+
+    def __post_init__(self) -> None:
+        if self.cpus <= 0 or self.memory_mib <= 0:
+            raise ValueError(f"cpus and memory_mib must be positive; got cpus={self.cpus!r}, memory_mib={self.memory_mib!r}")
+
+    @property
+    def docker_args(self) -> tuple[str, ...]:
+        """The ``docker run`` / ``docker create`` flags. Swap is held to the memory limit, so a container can't page past it."""
+        return ("--cpus", f"{self.cpus:g}", "--memory", f"{self.memory_mib}m", "--memory-swap", f"{self.memory_mib}m")
+
+    @property
+    def compose_keys(self) -> tuple[str, ...]:
+        """The same limits as compose service keys, which ``docker compose up`` applies outside swarm mode."""
+        return (f"cpus: {self.cpus:g}", f"mem_limit: {self.memory_mib}m", f"memswap_limit: {self.memory_mib}m")
+
+    def to_dict(self) -> dict:
+        return {"cpus": self.cpus, "memory_mib": self.memory_mib}
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ContainerLimits":
+        return cls(cpus=float(data["cpus"]), memory_mib=int(data["memory_mib"]))
+
+
 class Sandbox(ABC):
     """Universal sandbox contract — anything that can host a process and expose ports."""
 
@@ -148,6 +177,8 @@ class Sandbox(ABC):
     host_ips: tuple[str, ...] = ()
     # ``name:address`` entries the containers started on this sandbox add to their hosts file.
     extra_hosts: tuple[str, ...] = ()
+    # What each container started on this sandbox is held to; None where the sandbox itself (a VM) is the bound.
+    container_limits: ContainerLimits | None = None
 
     _VM_READY_TIMEOUT = 1200      # wait_for_vm wall-clock budget (s)
     _VM_READY_POLL_INTERVAL = 30  # sparse polling (s)

src/agent_env/providers/sandbox_providers/sandbox_provider.py

@@ -139,9 +139,11 @@ def _container_args(self, sandbox: VmSandbox, *, image_name: str, port: int, env
         """The ``docker run`` / ``docker create`` arguments for the sandbox's container."""
         env_flags = " \\\n    ".join(f"-e {shlex.quote(k)}={shlex.quote(v)}" for k, v in env.items())
         extra_args = f"{self.EXTRA_CONTAINER_RUN_ARGS} " if self.EXTRA_CONTAINER_RUN_ARGS else ""
+        limits = sandbox.container_limits
+        limit_args = "".join(f"{shlex.quote(arg)} " for arg in limits.docker_args) if limits else ""
         publish = " ".join(f"-p {spec}" for spec in port_bindings(sandbox.host_ips, sandbox.host_port(port), port))
         return (
-            f"--name {shlex.quote(sandbox.container_name)} {publish} {extra_args}\\\n    "
+            f"--name {shlex.quote(sandbox.container_name)} {publish} {extra_args}{limit_args}\\\n    "
             f"{env_flags} \\\n    {shlex.quote(image_name)}"
         )
 

src/agent_env/task_step/task_steps/run_docker_container.py

@@ -219,10 +219,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
         cap_flags = " ".join(f"--cap-add {shlex.quote(c)}" for c in self.cap_add)
         volume_flags = " ".join(f"-v {shlex.quote(v)}" for v in self.volumes)
         host_flags = " ".join(f"--add-host {shlex.quote(entry)}" for entry in sandbox.extra_hosts)
+        limits = sandbox.container_limits
+        limit_flags = " ".join(shlex.quote(arg) for arg in limits.docker_args) if limits else ""
         extra_flags = " ".join(f for f in (
             "--privileged" if self.privileged else "",
             f"--shm-size {shlex.quote(self.shm_size)}" if self.shm_size else "",
-            device_flags, cap_flags, volume_flags, host_flags,
+            device_flags, cap_flags, volume_flags, host_flags, limit_flags,
         ) if f)
         entrypoint_flag = ""
         if self.command_override:

tst/integration/task_step/test_container_limits_local.py

@@ -0,0 +1,85 @@
+"""The containers a local sandbox starts are held to the cpu and memory it was created with.
+
+A VM's size bounds its containers; a local sandbox shares this host's Docker, so Docker holds each one.
+Needs Docker and a pull of a small public image.
+"""
+
+import shutil
+import subprocess
+import uuid
+
+import pytest
+
+from agent_env.artifact import FileArtifact, FileArtifactUniverse
+from agent_env.artifact.store import reset_artifact_store
+from agent_env.config import configure, reset_config
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandboxProvider
+from agent_env.task.teardown import teardown_run
+from agent_env.task_step.context import TaskStepContext
+from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep
+from agent_env.task_step.task_steps.run_docker_container import RunDockerContainerTaskStep
+from tst.util.capabilities import missing_capability_reason
+
+pytestmark = pytest.mark.integration
+
+_IMAGE = "mirror.gcr.io/library/nginx:1.27-bookworm"
+
+
+@pytest.fixture
+def local_backends(monkeypatch, tmp_path):
+    if shutil.which("docker") is None or subprocess.run(["docker", "info"], capture_output=True).returncode != 0:
+        pytest.skip(missing_capability_reason("docker_daemon"))
+    monkeypatch.chdir(tmp_path)
+    monkeypatch.setenv("AGENT_ENV_DOCUMENT_STORE", "local")
+    monkeypatch.setenv("AGENT_ENV_OBJECT_STORE", "local")
+    monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path / "sandboxes"))
+    configure()
+    reset_artifact_store()
+    try:
+        yield tmp_path
+    finally:
+        reset_artifact_store()
+        reset_config()
+
+
+def _limits(container: str) -> tuple[int, int, int]:
+    """``(NanoCpus, Memory, MemorySwap)`` Docker holds ``container`` to."""
+    out = subprocess.run(
+        ["docker", "inspect", container, "--format", "{{.HostConfig.NanoCpus}} {{.HostConfig.Memory}} {{.HostConfig.MemorySwap}}"],
+        capture_output=True, text=True, check=True,
+    ).stdout.split()
+    return int(out[0]), int(out[1]), int(out[2])
+
+
+@pytest.mark.asyncio
+async def test_a_container_a_step_starts_on_a_local_sandbox_is_held_to_its_cpu_and_memory(local_backends):
+    suffix = uuid.uuid4().hex[:8]
+    (local_backends / "Dockerfile").write_text(f"FROM {_IMAGE}\n")
+    build_context = FileArtifactUniverse.put(id=f"lim-ctx-{suffix}", file_artifacts={
+        "Dockerfile": FileArtifact.put(id=f"lim-df-{suffix}", description="Dockerfile",
+                                       file_path=str(local_backends / "Dockerfile")),
+    })
+    container = f"lim-{suffix}"
+    context = TaskStepContext(instance_id=f"lim-{suffix}")
+    try:
+        context = await DeploySandboxTaskStep(
+            id="box", version=None, sandbox_name="box", sandbox_mode="vm", sandbox_type="local",
+            cpu=0.5, memory_mb=256,
+        ).execute(context)
+        context = await RunDockerContainerTaskStep(
+            id="ctr", version=None, sandbox_name="box", docker_context_artifact_id=build_context.id,
+            container_name=container,
+        ).execute(context)
+
+        assert _limits(container) == (500_000_000, 256 * 2**20, 256 * 2**20)
+    finally:
+        await teardown_run(context)
+
+
+@pytest.mark.asyncio
+async def test_a_container_mode_local_sandbox_is_held_to_its_cpu_and_memory(local_backends):
+    sandbox = await LocalSandboxProvider().create_sandbox(image_name=_IMAGE, port=80, env={}, cpu=0.5, memory=256)
+    try:
+        assert _limits(sandbox.container_name) == (500_000_000, 256 * 2**20, 256 * 2**20)
+    finally:
+        await sandbox.terminate()