fix(sandbox): steps that target a reattached local agent run in its container

src/agent_env/providers/sandbox_providers/local_sandbox.py

@@ -52,7 +52,7 @@ def _runs_in_container(cmd: list[str]) -> bool:
 _REAP_SECONDS = 5
 
 # 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) can restore container mode and clean up.
+# (post-run teardown reconstructs the sandbox from disk) knows it owns that container.
 _CONTAINER_MODE_MARKER = ".agent-container-mode"
 
 
@@ -158,6 +158,11 @@ def container_name(self) -> str:
         container and lets teardown remove only the one this sandbox created."""
         return f"agent-{self.sandbox_id}"
 
+    @property
+    def owns_container(self) -> bool:
+        """Whether create_container started ``container_name`` for this sandbox, not an agent placed on it."""
+        return (self._work_dir / _CONTAINER_MODE_MARKER).exists()
+
     @property
     def work_dir(self) -> Path:
         return self._work_dir
@@ -173,8 +178,8 @@ def _rewrite_app_script(self, script: str) -> str:
     async def terminate(self) -> None:
         """Tear down whatever this sandbox is running.
 
-        A container-mode sandbox (the A2A agent) owns the ``self.container_name`` container that
-        create_container started — remove it. A VM-mode sandbox (env/gateway) runs a docker compose
+        A sandbox create_container made (the A2A agent; VM mode once reattached) owns the container
+        ``self.container_name`` it started — remove it. A VM-mode sandbox (env/gateway) runs a docker compose
         stack out of its work dir — ``docker compose down`` it — and may carry an agent placed on it,
         which runs as the same ``self.container_name`` container. Each path only touches resources
         this sandbox created: the container name is per-sandbox (LocalSandbox.container_name), so a
@@ -394,10 +399,8 @@ async def get_sandbox(self, sandbox_id: str) -> LocalSandbox:
         work_dir = LocalSandbox.find_work_dir(sandbox_id)
         if work_dir is None:
             raise RuntimeError(f"Local sandbox work directory not found for sandbox_id={sandbox_id!r}")
-        sandbox = LocalSandbox(sandbox_id=sandbox_id, work_dir=work_dir)
-        if (work_dir / _CONTAINER_MODE_MARKER).exists():
-            sandbox.mode = SANDBOX_MODE_CONTAINER
-        return sandbox
+        # VM mode even for a container it owns: exec runs on this host, so steps must `docker exec` into it.
+        return LocalSandbox(sandbox_id=sandbox_id, work_dir=work_dir)
 
     @classmethod
     def shares_network_with(cls, sandbox_type: Optional[str]) -> bool:

src/agent_env/task_step/task_steps/collect_artifacts.py

@@ -437,6 +437,11 @@ async def _discover_container(self, sandbox) -> str:
         # Prefer this sandbox's own container name, fall back to the a2a-agent-* prefix.
         if sandbox.container_name in running:
             return sandbox.container_name
+        if getattr(sandbox, "owns_container", False):
+            # Its Docker host is shared, so any other agent container there is another run's.
+            raise RuntimeError(
+                f"Agent container {sandbox.container_name!r} is not running. Running containers: {running}."
+            )
         fallback = [n for n in running if n.startswith("a2a-agent-")]
         if fallback:
             if len(fallback) > 1:

src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py

@@ -297,6 +297,11 @@ async def _discover_container(self, sandbox) -> str:
         running = [n.strip() for n in stdout.splitlines() if n.strip()]
         if sandbox.container_name in running:
             return sandbox.container_name
+        if getattr(sandbox, "owns_container", False):
+            # Its Docker host is shared, so any other agent container there is another run's.
+            raise RuntimeError(
+                f"Agent container {sandbox.container_name!r} is not running. Running containers: {running}."
+            )
         fallback = [n for n in running if n.startswith("a2a-agent-")]
         if fallback:
             if len(fallback) > 1:

tst/integration/task_step/test_verify_sandbox_local.py

@@ -1,26 +1,34 @@
 """A task scores against a bare local sandbox with no Docker, model or network.
 
 A local VM-mode sandbox is a host work dir driven by subprocesses, so this needs `bash`
-and `python3` on PATH and nothing else.
+and `python3` on PATH and nothing else. The agent tests start a real container in place of
+an agent's, so they need Docker too, but no model.
 """
 
 import shutil
+import subprocess
 import uuid
 
 import pytest
+import pytest_asyncio
 
 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 LocalSandbox
+from agent_env.config import configure, get_config, reset_config
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider
 from agent_env.task import Task
+from agent_env.task_step.context import DeployedAgent, TaskStepContext
+from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep
 from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep
 from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep
 from agent_env.task_step.task_steps.verifiers.verify_sandbox import VerifySandboxTaskStep
+from tst.util.capabilities import missing_capability_reason
 
 pytestmark = pytest.mark.integration
 
 _CHECK_PY = "import pathlib, sys\nsys.exit(0 if 'hello' in pathlib.Path('hello.txt').read_text() else 1)\n"
+# Stands in for an agent's image: a server that stays up, on a Debian userland (bash, GNU find), for any arch.
+_AGENT_IMAGE = "mirror.gcr.io/library/nginx:1.27-bookworm"
 
 
 @pytest.fixture
@@ -40,6 +48,21 @@ def local_backends(monkeypatch, tmp_path):
         reset_config()
 
 
+@pytest_asyncio.fixture
+async def local_agent(local_backends):
+    """A running local container sandbox recorded as the deployed agent `agent`."""
+    if shutil.which("docker") is None or subprocess.run(["docker", "info"], capture_output=True).returncode != 0:
+        pytest.skip(missing_capability_reason("docker_daemon"))
+    sandbox = await LocalSandboxProvider().create_container(image_name=_AGENT_IMAGE, port=80, env={})
+    try:
+        yield sandbox, TaskStepContext(instance_id=f"agent-{uuid.uuid4().hex[:8]}", deployed_agents=[
+            DeployedAgent(agent_name="agent", api_url=sandbox.tunnel_urls[80], sandbox_id=sandbox.sandbox_id,
+                          sandbox_type="local"),
+        ])
+    finally:
+        await sandbox.terminate()
+
+
 def _greeting_universe(src_dir, suffix):
     src_dir.mkdir()
     (src_dir / "hello.txt").write_text("hello, world\n")
@@ -112,3 +135,40 @@ async def test_file_artifacts_stage_beside_each_other(local_backends):
     assert ctx.metadata["verifications"]["hello"]["score"] == 1.0
     loaded = ctx.metadata["loaded_file_artifact_universes"]
     assert [(entry["artifact_type"], entry["files"]) for entry in loaded] == [("file", ["hello.txt"]), ("file", ["check.py"])]
+
+
+@pytest.mark.asyncio
+async def test_verify_sandbox_probes_a_local_agent_inside_its_container(local_agent):
+    sandbox, context = local_agent
+    _write_in_container(sandbox, "/app/out/report.txt", "done")
+    step = VerifySandboxTaskStep(
+        id="verify", version=None, agent_name="agent", base_dir="/app/out", verifier_id="agent",
+        criteria=[
+            {"type": "probe_file_exists", "criterion": "wrote it", "paths": ["report.txt"]},
+            {"type": "bash_cmd_succeeds", "criterion": "runs in the container", "bash_cmd": 'test "$(uname -s)" = Linux'},
+        ],
+    )
+
+    ctx = await step.execute(context)
+
+    assert [r["result"] for r in ctx.metadata["verifications"]["agent"]["results"]] == [True, True]
+    assert list(sandbox.work_dir.iterdir()) == [sandbox.work_dir / ".agent-container-mode"]
+
+
+@pytest.mark.asyncio
+async def test_collect_artifacts_reads_a_local_agents_file_from_its_container(local_agent):
+    sandbox, context = local_agent
+    token = uuid.uuid4().hex
+    _write_in_container(sandbox, "/app/artifact/report.txt", token)
+    step = CollectArtifactsTaskStep(id="collect", version=None, agent_name="agent", artifact_paths=["report.txt"])
+
+    ctx = await step.execute(context)
+
+    url = ctx.metadata["collected_artifacts"]["collect"]["artifacts"]["report.txt"]
+    assert get_config().get_object_store().get(url) == f"{token}\n".encode()
+
+
+def _write_in_container(sandbox, path, text):
+    """Write ``text`` at ``path`` inside the sandbox's container, as the agent would."""
+    script = f"mkdir -p $(dirname {path}) && echo {text} > {path}"
+    subprocess.run(["docker", "exec", sandbox.container_name, "sh", "-c", script], check=True, capture_output=True)

tst/unit/providers/sandbox_providers/local_sandbox_test.py

@@ -11,7 +11,7 @@
 from agent_env.a2a_agent.a2a_agent import A2AAgent
 from agent_env.config import reset_config, set_object_store
 from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider
-from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER
+from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_CONTAINER, SANDBOX_MODE_VM
 from agent_env.store import LocalFilesystemObjectStore
 
 
@@ -244,15 +244,18 @@ async def test_terminate_vm_without_compose_removes_only_an_agent_placed_on_it(t
 
 
 @pytest.mark.asyncio
-async def test_get_sandbox_restores_container_mode_from_marker(tmp_path: Path, monkeypatch):
-    """The container-mode marker persisted at create time is restored so teardown cleans up."""
+async def test_a_reattached_container_stays_vm_mode_and_owned(tmp_path: Path, monkeypatch):
+    """Its exec runs on this host, so a reattached agent's steps must take the VM path into its container."""
     monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path))
-    work_dir = tmp_path / "agent-env-local-agent1-abc123"
-    work_dir.mkdir()
-    (work_dir / ".agent-container-mode").write_text("agent-local-agent1")
+    for name in ("agent1", "vm1"):
+        (tmp_path / f"agent-env-local-{name}-abc123").mkdir()
+    (tmp_path / "agent-env-local-agent1-abc123" / ".agent-container-mode").write_text("agent-local-agent1")
 
-    sandbox = await LocalSandboxProvider().get_sandbox("local-agent1")
-    assert sandbox.mode == SANDBOX_MODE_CONTAINER
+    agent = await LocalSandboxProvider().get_sandbox("local-agent1")
+    vm = await LocalSandboxProvider().get_sandbox("local-vm1")
+
+    assert (agent.mode, agent.owns_container) == (SANDBOX_MODE_VM, True)
+    assert (vm.mode, vm.owns_container) == (SANDBOX_MODE_VM, False)
 
 
 @pytest.mark.asyncio
@@ -285,8 +288,10 @@ async def create_vm(self, **kwargs):
 
     created = await getattr(_Provider(), create)(image_name="img:v1", port=8000, env={})
     rebuilt = await LocalSandboxProvider().get_sandbox(created.sandbox_id)
+    del scripts[:]
     await rebuilt.terminate()
-    assert rebuilt.mode == SANDBOX_MODE_CONTAINER and scripts[-1] == f"docker rm -f {created.container_name} >/dev/null 2>&1 || true"
+    assert created.owns_container and rebuilt.owns_container and rebuilt.mode == SANDBOX_MODE_VM
+    assert scripts == [f"docker rm -f {created.container_name} >/dev/null 2>&1 || true"]
 
 
 @pytest.mark.asyncio

tst/unit/task/teardown_test.py

@@ -190,6 +190,26 @@ def test_a_local_sandbox_is_brought_down_and_its_work_folder_removed(sandbox_roo
     assert list(sandbox_root.iterdir()) == []
 
 
+def test_a_reattached_local_agent_has_its_container_and_work_folder_removed(sandbox_root, monkeypatch):
+    agent = LocalSandbox()
+    (agent.work_dir / ".agent-container-mode").write_text(agent.container_name)
+    scripts = []
+
+    async def record(self, script, *, max_retries=0):
+        scripts.append(script)
+        return ""
+
+    monkeypatch.setattr(LocalSandbox, "exec_script", record)
+    context = TaskStepContext(deployed_agents=[
+        DeployedAgent(agent_name="solver", api_url="http://agent", sandbox_id=agent.sandbox_id, sandbox_type="local")])
+
+    report = _run(context)
+
+    assert report.terminated == (RecordedSandbox(agent.sandbox_id, "local", "agent"),)
+    assert scripts == [f"docker rm -f {agent.container_name} >/dev/null 2>&1 || true"]
+    assert list(sandbox_root.iterdir()) == []
+
+
 def test_a_slow_folder_removal_does_not_count_against_the_terminate_timeout(sandbox_root, monkeypatch):
     sandbox = LocalSandbox()
     context = TaskStepContext(deployed_sandboxes=[

tst/unit/task_step/test_collect_artifacts.py

@@ -18,7 +18,8 @@
 import pytest
 
 from agent_env.task_step.context import DeployedAgent, PromptResponse, TaskStepContext
-from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep, _is_url_entry
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider
+from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep, _exec_args, _is_url_entry
 from agent_env.env.env import DeployedEnv, DeployedGatewayEnv, EnvCapabilityUnsupported
 from agent_env.env.gateway.constants import EXT_STEP_URI, GATEWAY_EXTENSIONS, WELL_KNOWN_PATH
 from tst.unit.event_loop_probe import on_event_loop
@@ -464,6 +465,27 @@ def test_unreachable_sandbox_is_clear_error(self):
             except RuntimeError as e:
                 assert "unreachable" in str(e)
 
+    def test_a_reattached_local_agent_is_read_inside_its_own_container(self, tmp_path, monkeypatch):
+        sandbox, calls = _reattached_local_agent(tmp_path, monkeypatch, running="a2a-agent-other\nagent-local-agent1\n")
+        step = CollectArtifactsTaskStep(id="collect", version=None, agent_name="solver")
+
+        container = _run(step._discover_container(sandbox))
+        _run(step._list_base_directory(sandbox, container))
+
+        assert calls == [
+            ("sudo", "docker", "ps", "--format", "{{.Names}}"),
+            ("sudo", "docker", "exec", "agent-local-agent1", "find", "/app/artifact", "-type", "f", "-printf", "%P\n"),
+        ]
+        assert _exec_args(sandbox, container, ("bash", "-c", "base64 < /app/artifact/a.txt")) == (
+            "sudo", "docker", "exec", "agent-local-agent1", "bash", "-c", "base64 < /app/artifact/a.txt")
+
+    def test_a_reattached_local_agent_never_borrows_another_runs_container(self, tmp_path, monkeypatch):
+        sandbox, calls = _reattached_local_agent(tmp_path, monkeypatch, running="a2a-agent-other\n")
+
+        with pytest.raises(RuntimeError, match="'agent-local-agent1' is not running"):
+            _run(CollectArtifactsTaskStep(id="collect", version=None)._discover_container(sandbox))
+        assert calls == [("sudo", "docker", "ps", "--format", "{{.Names}}")]
+
 
 def _manifest_ctx(artifacts: dict) -> TaskStepContext:
     # Manifest lives on the producing step's response (keyed by step_id), which collect reads.
@@ -756,3 +778,18 @@ def handle(request: httpx.Request) -> httpx.Response:
 
     monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle)))
     return sent
+
+
+def _reattached_local_agent(tmp_path, monkeypatch, *, running: str):
+    """A local agent sandbox rebuilt from its work dir, answering `docker ps` with ``running`` and recording each command."""
+    monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path))
+    (tmp_path / "agent-env-local-agent1-abc123").mkdir()
+    (tmp_path / "agent-env-local-agent1-abc123" / ".agent-container-mode").write_text("agent-local-agent1")
+    calls: list[tuple[str, ...]] = []
+
+    async def exec_with_output(self, *args):
+        calls.append(args)
+        return 0, running if "ps" in args else "", ""
+
+    monkeypatch.setattr(LocalSandbox, "exec_with_output", exec_with_output)
+    return _run(LocalSandboxProvider().get_sandbox("local-agent1")), calls

tst/unit/task_step/test_run_code.py

@@ -308,6 +308,19 @@ async def test_agent_target_honors_the_agents_sandbox_type():
     assert ctx.metadata["script_results"]["filter"] == {"keep": True}
 
 
+@pytest.mark.asyncio
+async def test_a_reattached_local_agent_runs_the_script_inside_its_container(tmp_path, monkeypatch):
+    monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path))
+    (tmp_path / "agent-env-local-agent1-abc123").mkdir()
+    (tmp_path / "agent-env-local-agent1-abc123" / ".agent-container-mode").write_text("agent-local-agent1")
+    ctx = _make_context()
+    ctx.deployed_agents[0].sandbox_id, ctx.deployed_agents[0].sandbox_type = "local-agent1", "local"
+
+    target = await _step()._agent_target(ctx)
+
+    assert target.prefix == ("sudo", "docker", "exec", "-u", "0", "agent-local-agent1")
+
+
 @pytest.mark.asyncio
 async def test_host_installed_agent_on_a_vm_runs_on_the_host():
     step = _step(result_id="filter")

tst/unit/task_step/verifiers/test_verify_sandbox.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, LocalSandboxProvider
 from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext
 from agent_env.task_step.task_step import TaskStep
 from agent_env.task_step.task_steps.verifiers import verify_sandbox
@@ -164,6 +165,32 @@ async def test_agent_source_probes_inside_the_agent_container(provide):
     assert ctx.metadata["verifications"]["hello"]["score"] == 1.0
 
 
+@pytest.mark.asyncio
+async def test_a_reattached_local_agent_is_probed_inside_its_own_container(provide, tmp_path, monkeypatch):
+    sandbox, calls = await _reattached_local_agent(tmp_path, monkeypatch, running="a2a-agent-other\nagent-local-agent1\n")
+    provide(sandbox)
+
+    await _step(agent_name="solver").execute(_local_agent_context())
+
+    assert calls == [
+        ("sudo", "docker", "ps", "--format", "{{.Names}}"),
+        ("sudo", "docker", "exec", "agent-local-agent1", "test", "-e", "/app/greeting/hello.txt"),
+        ("sudo", "docker", "exec", "agent-local-agent1", "test", "-e", "/etc/hosts"),
+        ("sudo", "docker", "exec", "agent-local-agent1", "cat", "/app/greeting/hello.txt"),
+        ("sudo", "docker", "exec", "-w", "/app/greeting", "agent-local-agent1", "bash", "-c", "python3 check.py"),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_a_reattached_local_agent_whose_container_is_gone_fails_rather_than_borrow_another(provide, tmp_path, monkeypatch):
+    sandbox, calls = await _reattached_local_agent(tmp_path, monkeypatch, running="a2a-agent-other\n")
+    provide(sandbox)
+
+    with pytest.raises(RuntimeError, match="'agent-local-agent1' is not running"):
+        await _step(agent_name="solver").execute(_local_agent_context())
+    assert calls == [("sudo", "docker", "ps", "--format", "{{.Names}}")]
+
+
 @pytest.mark.asyncio
 async def test_unknown_sandbox_name_names_the_deployed_ones():
     with pytest.raises(RuntimeError, match=r"Sandbox 'nope' not found in context.deployed_sandboxes \(deployed: \['box'\]\)"):
@@ -239,3 +266,24 @@ async def exec_with_output(self, *args):
         await _step(agent_name="solver").execute(context)
 
     assert "Connected to sandbox sb-1 (mode=vm)" in caplog.text
+
+
+def _local_agent_context():
+    return TaskStepContext(deployed_agents=[
+        DeployedAgent(agent_name="solver", api_url="http://agent", sandbox_id="local-agent1", sandbox_type="local"),
+    ])
+
+
+async def _reattached_local_agent(tmp_path, monkeypatch, *, running: str):
+    """A local agent sandbox rebuilt from its work dir, answering `docker ps` with ``running`` and recording each command."""
+    monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path))
+    (tmp_path / "agent-env-local-agent1-abc123").mkdir()
+    (tmp_path / "agent-env-local-agent1-abc123" / ".agent-container-mode").write_text("agent-local-agent1")
+    calls: list[tuple[str, ...]] = []
+
+    async def exec_with_output(self, *args):
+        calls.append(args)
+        return 0, running if "ps" in args else "hello world\n", ""
+
+    monkeypatch.setattr(LocalSandbox, "exec_with_output", exec_with_output)
+    return await LocalSandboxProvider().get_sandbox("local-agent1"), calls