fix(providers): teardown terminates a sandbox by id, so one that already finished is counted down

src/agent_env/providers/sandbox_providers/e2b/provider.py

@@ -381,3 +381,7 @@ async def get_sandbox(self, sandbox_id: str) -> E2BSandbox:
             )
             sandbox.network_policy = None
         return sandbox
+
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        if not await self._get_sandbox_cls().kill(sandbox_id, api_key=self._api_key):
+            logger.info("E2B sandbox %s was already absent during terminate", sandbox_id)

src/agent_env/providers/sandbox_providers/modal_sandbox.py

@@ -90,6 +90,14 @@ def _tunnel_url(t) -> str:
     return f"https://{t.host}" + (f":{t.port}" if t.port != 443 else "")
 
 
+async def _terminate_by_id(client, sandbox_id: str) -> None:
+    try:
+        sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client)
+        await sb.terminate.aio()
+    except modal.exception.NotFoundError:
+        logger.info("Modal sandbox %s was already gone at terminate", sandbox_id)
+
+
 def _modal_network_kwargs(policy: NetworkPolicy) -> dict:
     """``policy`` as Modal ``_experimental_create`` kwargs.
 
@@ -480,6 +488,9 @@ async def get_sandbox(self, sandbox_id: str) -> Sandbox:
         i6pn_address = await _resolve_i6pn_address(sb)
         return self._sandbox_cls(sb, tunnel_urls, i6pn_address=i6pn_address)
 
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        await _terminate_by_id(await self._get_client(), sandbox_id)
+
     async def _registry_image(self, image_name: str) -> modal.Image:
         """``image_name`` pulled from its registry, with the image store's credentials for it, if any."""
         image_store = get_config().get_image_store_at(image_name)

src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py

@@ -32,6 +32,7 @@
     _attribution_tags,
     _log_sandbox_started,
     _resolve_app_base_name,
+    _terminate_by_id,
     _tunnel_url,
     _modal_network_kwargs,
 )
@@ -412,6 +413,9 @@ async def get_sandbox(self, sandbox_id: str) -> ModalVmSandbox:
         # dockerd is already running on a live sandbox; do not re-run setup here.
         return ModalVmSandbox(sb, tunnel_urls)
 
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        await _terminate_by_id(await self._get_client(), sandbox_id)
+
 
 # ── Object downloads (internal) ────────────────────────────────────────────────────
 # Object downloads run as one aria2c per object: a pool of HTTP range connections with small pieces,

src/agent_env/providers/sandbox_providers/sail_vm/provider.py

@@ -429,3 +429,14 @@ async def get_sandbox(self, sandbox_id: str) -> SailVmSandbox:
                 sandbox_id, applied,
             )
         return sandbox
+
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        """Through the handle: terminating also deletes the Sailbox's model-key policy, which only its applied
+        policy names."""
+        sdk, _ = await self._connect()
+        try:
+            sandbox = await self.get_sandbox(sandbox_id)
+        except sdk.NotFoundError:
+            logger.info("Sailbox %s was already gone at terminate", sandbox_id)
+            return
+        await sandbox.terminate()

src/agent_env/providers/sandbox_providers/sandbox_provider.py

@@ -162,6 +162,11 @@ def _container_args(self, sandbox: VmSandbox, *, image_name: str, port: int, env
     async def get_sandbox(self, sandbox_id: str) -> Sandbox:
         raise NotImplementedError(f"{type(self).__name__} does not support get_sandbox")
 
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        """Terminate a sandbox this provider made, by id; one that has already finished or is gone is terminated.
+        A backend that can terminate without a live handle overrides this, so teardown doesn't need one."""
+        await (await self.get_sandbox(sandbox_id)).terminate()
+
     async def close(self) -> None:
         pass
 

src/agent_env/providers/sandbox_providers/smol_vm/provider.py

@@ -183,6 +183,16 @@ async def get_sandbox(self, sandbox_id: str) -> SmolVmSandbox:
         sandbox.network_policy = self.effective_network_policy(None)
         return sandbox
 
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        """Connecting starts a stopped machine, and the SDK deletes only a connected one, so a listed machine is
+        connected first; one not listed is gone."""
+        conn = ConnectOptions(target="local", wait_for_ports=False)
+        if sandbox_id not in {machine.name for machine in await asyncio.to_thread(Machine.list, conn)}:
+            logger.info("SmolVM machine %s was already gone at terminate", sandbox_id)
+            return
+        machine = await asyncio.to_thread(Machine.connect, sandbox_id, conn)
+        await asyncio.to_thread(machine.delete)
+
     async def _configure_host_access(self, sandbox: SmolVmSandbox) -> None:
         code, stdout, stderr = await sandbox.exec_with_output("getent", "ahostsv4", "host.smolvm.internal")
         if code or not stdout.strip():

src/agent_env/providers/sandbox_providers/tenki/provider.py

@@ -18,13 +18,13 @@
 from typing import Any, ClassVar, Self
 from urllib.parse import urlparse
 
-from tenki import AsyncClient
+from tenki import AsyncClient, SessionNotFoundError
 
 from agent_env.attribution import Attribution
 from agent_env.config.errors import ConfigError
 from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy, NetworkPolicyUnsupportedError
 from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider, apply_default_attribution
-from agent_env.providers.sandbox_providers.tenki.sandbox import ENDED_STATES, TenkiSandbox, policy_from_info
+from agent_env.providers.sandbox_providers.tenki.sandbox import ENDED_STATES, TenkiSandbox, close_session, policy_from_info
 
 logger = logging.getLogger(__name__)
 
@@ -413,3 +413,11 @@ async def get_sandbox(self, sandbox_id: str) -> TenkiSandbox:
             network_policy=policy_from_info(info),
             command_timeout=_remaining_lifetime(info.timeout_at),
         )
+
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        try:
+            session = await self._client().get(sandbox_id)
+        except SessionNotFoundError:
+            logger.info("Tenki session %s was already gone at terminate", sandbox_id)
+            return
+        await close_session(session)

src/agent_env/providers/sandbox_providers/tenki/sandbox.py

@@ -174,30 +174,13 @@ def __init__(
         self.network_policy = network_policy
 
     async def terminate(self) -> None:
-        """Terminate the VM. One another handle already terminated, while this handle's cached state still says
-        running, is not an error."""
-        try:
-            await self._session.close_if_open()
-        except SessionNotFoundError:
-            logger.info("Tenki session %s was already gone at terminate", self.sandbox_id)
-        except (InvalidStateError, SessionTerminatedError):
-            if await self._current_state() not in ENDED_STATES:
-                raise
-            logger.info("Tenki session %s was already terminating at terminate", self.sandbox_id)
-        await self.release()
+        await close_session(self._session)
 
     async def release(self) -> None:
-        """Stop the session's data-plane credential renewal, which close_if_open leaves to the caller for a session
-        that is already ending."""
-        reset = getattr(self._session, "_reset_data_plane", None)
-        if reset is not None:
-            await reset()
+        await release_session(self._session)
 
     async def _current_state(self) -> str:
-        try:
-            return (await self._session.refresh()).state
-        except SessionNotFoundError:
-            return "GONE"
+        return await _session_state(self._session)
 
     async def exec(self, *command: str) -> _CompletedProcess:
         """Run argv to completion as root, so a leading ``sudo`` is dropped. A lost stream is exit -1, which
@@ -329,3 +312,32 @@ async def _load_tarballs(self, artifacts: list) -> None:
         signed_urls = await self._signed_image_urls(artifacts)
         reachable = [url if self._can_download(url) else None for url in signed_urls]
         await self._load_docker_images(artifacts, reachable)
+
+
+async def close_session(session: Any) -> None:
+    """Terminate a session's VM. One another handle already terminated, while this one's cached state still says
+    running, is not an error."""
+    try:
+        await session.close_if_open()
+    except SessionNotFoundError:
+        logger.info("Tenki session %s was already gone at terminate", session.id)
+    except (InvalidStateError, SessionTerminatedError):
+        if await _session_state(session) not in ENDED_STATES:
+            raise
+        logger.info("Tenki session %s was already terminating at terminate", session.id)
+    await release_session(session)
+
+
+async def release_session(session: Any) -> None:
+    """Stop the session's data-plane credential renewal, which close_if_open leaves to the caller for a session
+    that is already ending."""
+    reset = getattr(session, "_reset_data_plane", None)
+    if reset is not None:
+        await reset()
+
+
+async def _session_state(session: Any) -> str:
+    try:
+        return (await session.refresh()).state
+    except SessionNotFoundError:
+        return "GONE"

src/agent_env/providers/sandbox_providers/tensorlake/provider.py

@@ -336,3 +336,10 @@ async def get_sandbox(self, sandbox_id: str) -> TensorlakeSandbox:
                 "Reconnected to Tensorlake sandbox %s, but image loading will fail closed: %s", sandbox_id, exc,
             )
         return sandbox
+
+    async def terminate_sandbox(self, sandbox_id: str) -> None:
+        try:
+            raw_sandbox = await self._sandbox_cls.connect(sandbox_id, **self._client_kwargs())
+            await raw_sandbox.terminate()
+        except SandboxNotFoundError:
+            logger.info("Tensorlake sandbox %s was already absent during terminate", sandbox_id)

src/agent_env/task/teardown.py

@@ -98,14 +98,14 @@ async def take_down(sandbox: RecordedSandbox) -> None:
             return
         try:
             async with asyncio.timeout(timeout):
-                handle = await _terminate(sandbox)
+                on_this_machine = await _terminate(sandbox)
         except Exception as e:
             reason = f"timed out after {timeout:g}s" if isinstance(e, TimeoutError) else _described(e)
             logger.warning("teardown: %s (%s) was not terminated: %s", sandbox.sandbox_id, kind(sandbox), reason)
             failed.append((sandbox, reason))
             return
         context.metadata[TORN_DOWN_KEY] = [*context.metadata.get(TORN_DOWN_KEY, []), sandbox.sandbox_id]
-        if handle.ON_THIS_MACHINE and (why := await _remove_folder(sandbox)):
+        if on_this_machine and (why := await _remove_folder(sandbox)):
             failed.append((sandbox, why))
         else:
             terminated.append(sandbox)
@@ -133,12 +133,17 @@ async def _remove_folder(sandbox: RecordedSandbox) -> str | None:
     return None
 
 
-async def _terminate(sandbox: RecordedSandbox):
-    provider = (build_sandbox_provider(sandbox.sandbox_type) if sandbox.sandbox_type
-                else _DEFAULT_PROVIDERS[sandbox.slot]())
-    handle = await provider.get_sandbox(sandbox.sandbox_id)
+async def _terminate(sandbox: RecordedSandbox) -> bool:
+    """Terminate ``sandbox`` and say whether it ran on this machine: by id through the backend that made it, or, for
+    a record that doesn't name one, through a live handle, since "not found" from a provider that didn't make the
+    sandbox doesn't mean it's gone."""
+    if sandbox.sandbox_type:
+        provider = build_sandbox_provider(sandbox.sandbox_type)
+        await provider.terminate_sandbox(sandbox.sandbox_id)
+        return provider.ON_THIS_MACHINE
+    handle = await _DEFAULT_PROVIDERS[sandbox.slot]().get_sandbox(sandbox.sandbox_id)
     await handle.terminate()
-    return handle
+    return handle.ON_THIS_MACHINE
 
 
 def _described(error: Exception) -> str:

src/agent_env/task_step/task_steps/teardown_sandboxes.py

@@ -166,8 +166,10 @@ def _is_backend(name: str) -> bool:
 
 
 async def _terminate(sandbox_id: str, sandbox_type: Optional[str], default_slot: str) -> None:
-    provider: SandboxProvider = (
-        build_sandbox_provider(sandbox_type) if sandbox_type else _DEFAULT_PROVIDERS[default_slot]()
-    )
-    sandbox = await provider.get_sandbox(sandbox_id)
-    await sandbox.terminate()
+    """By id through the backend that made it; a record that doesn't name one goes through a live handle, since
+    "not found" from a provider that didn't make the sandbox doesn't mean it's gone."""
+    if sandbox_type:
+        await build_sandbox_provider(sandbox_type).terminate_sandbox(sandbox_id)
+        return
+    provider: SandboxProvider = _DEFAULT_PROVIDERS[default_slot]()
+    await (await provider.get_sandbox(sandbox_id)).terminate()

tst/manual/smol_vm_agent_check.py

@@ -13,7 +13,7 @@
 from agent_env.task_step.context import TaskStepContext
 from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep
 from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep
-from tst.integration.store.local_parallel_deploy_test import local_stack  # noqa: F401
+from tst.integration.store.local_parallel_deploy_test import local_stack as local_stack
 from tst.util.a2a_test_agent import build_test_agent_image, put_test_agent
 
 

tst/manual/smol_vm_gateway_check.py

@@ -18,7 +18,7 @@
 from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep
 from tst.integration.env.envs.multi_env_local_test import _items_universe
 from tst.integration.env.envs.server_env_local_test import _put_items_image
-from tst.integration.store.local_parallel_deploy_test import _seed_bootstrap_envs, local_stack  # noqa: F401
+from tst.integration.store.local_parallel_deploy_test import _seed_bootstrap_envs, local_stack as local_stack
 from tst.util.a2a_test_agent import put_test_agent
 
 

tst/unit/bundle/run_teardown_test.py

@@ -22,6 +22,7 @@
 from agent_env.env.env import DeployedGatewayEnv
 from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext
 from agent_env.task_step.task_step import TaskStep
+from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider
 from tst.unit.bundle._support import layout, local_store
 
 pytestmark = pytest.mark.usefixtures("sigint_handled")
@@ -91,7 +92,7 @@ def __init__(self):
 def terminated(monkeypatch):
     log = _Terminated()
 
-    class _Provider:
+    class _Provider(SandboxProvider):
         async def get_sandbox(self, sandbox_id):
             return _FakeSandbox(log, sandbox_id)
 

tst/unit/cli/eval_run_teardown_test.py

@@ -15,6 +15,7 @@
 from agent_env.task import Task, teardown
 from agent_env.task.teardown import TeardownReport
 from agent_env.task_step.context import DeployedSandbox
+from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider
 
 eval_run = importlib.import_module("agent_env.cli.eval.run")  # the package's ``run`` is the click command
 pytestmark = pytest.mark.usefixtures("sigint_handled")
@@ -81,7 +82,7 @@ async def terminate(self):
                 await asyncio.sleep(0.2)
             terminated.append(self.sandbox_id)
 
-    class _Provider:
+    class _Provider(SandboxProvider):
         async def get_sandbox(self, sandbox_id):
             return _Sandbox(sandbox_id)
 

tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py

@@ -328,3 +328,17 @@ async def test_adapter_construction_failure_kills_raw_sandbox():
         await provider.create_vm(exposed_ports=[8080], setup_for_gateway=False)
 
     raw.kill.assert_awaited_once_with()
+
+
+@pytest.mark.parametrize("killed", [True, False], ids=["killed", "already absent"])
+@pytest.mark.asyncio
+async def test_terminate_sandbox_kills_by_id_without_connecting(monkeypatch, killed):
+    kill = AsyncMock(return_value=killed)
+    monkeypatch.setattr(_AsyncSandboxSdk, "kill", kill, raising=False)
+    provider = E2BSandboxProvider(api_key="e2b-secret", base_template="agent-env-v1", template_resolver=_Resolver(),
+                                  sandbox_cls=_AsyncSandboxSdk)
+
+    await provider.terminate_sandbox("e2b-sandbox-1")
+
+    kill.assert_awaited_once_with("e2b-sandbox-1", api_key="e2b-secret")
+    _AsyncSandboxSdk.connect.assert_not_called()

tst/unit/providers/sandbox_providers/modal_sandbox_test.py

@@ -5,6 +5,7 @@
 import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import modal
 import pytest
 
 from agent_env.config import get_config
@@ -437,3 +438,34 @@ async def collect(chunks, destination_path):
     sandbox._write_stream_via_exec = collect
     await sandbox.write_file_from_object(url, "/data/a.json")
     assert written == [(b"LOCAL", "/data/a.json")]
+
+
+@pytest.mark.asyncio
+async def test_terminate_sandbox_terminates_by_id_without_reading_its_tunnels(_patched_provider):
+    sb = MagicMock()
+    sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated"))
+    sb.terminate.aio = AsyncMock()
+    with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id:
+        from_id.aio = AsyncMock(return_value=sb)
+        await _patched_provider.terminate_sandbox("sb-finished")
+
+    sb.terminate.aio.assert_awaited_once()
+    sb.tunnels.aio.assert_not_called()
+    sb.exec.aio.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_terminate_sandbox_counts_a_sandbox_modal_no_longer_knows_as_gone(_patched_provider):
+    with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id:
+        from_id.aio = AsyncMock(side_effect=modal.exception.NotFoundError("Sandbox not found"))
+        await _patched_provider.terminate_sandbox("sb-gone")
+
+
+@pytest.mark.asyncio
+async def test_get_sandbox_still_raises_for_a_finished_sandbox_so_a_reattach_sees_it_dead(_patched_provider):
+    sb = MagicMock()
+    sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated"))
+    with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id:
+        from_id.aio = AsyncMock(return_value=sb)
+        with pytest.raises(modal.exception.ConflictError):
+            await _patched_provider.get_sandbox("sb-finished")

tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py

@@ -9,6 +9,7 @@
 
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import modal
 import pytest
 
 from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandbox, ModalVmSandboxProvider
@@ -395,3 +396,19 @@ async def test_apps_are_looked_up_in_the_configured_environment():
         lookup.aio = AsyncMock(return_value=MagicMock())
         await provider._get_app("agent-env-p")
     assert lookup.aio.call_args.kwargs["environment_name"] == "rl-prod"
+
+
+@pytest.mark.parametrize("found", [True, False], ids=["live or finished", "gone"])
+@pytest.mark.asyncio
+async def test_terminate_sandbox_terminates_a_vm_by_id_and_counts_one_modal_no_longer_knows_as_gone(found):
+    provider = ModalVmSandboxProvider()
+    provider._get_client = AsyncMock(return_value=MagicMock())  # type: ignore[method-assign]
+    sb = MagicMock()
+    sb.terminate.aio = AsyncMock()
+    with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id:
+        from_id.aio = (AsyncMock(return_value=sb) if found
+                       else AsyncMock(side_effect=modal.exception.NotFoundError("Sandbox not found")))
+        await provider.terminate_sandbox("sb-vm")
+
+    assert sb.terminate.aio.await_count == (1 if found else 0)
+    sb.tunnels.aio.assert_not_called()