From 78348f8e6ae4481af61bd405ef791840612ca4ed Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 13:00:43 -0700 Subject: [PATCH 1/3] fix(providers): tear down a Modal sandbox that already finished instead of reporting it up Teardown gets a handle with get_sandbox before it terminates. The Modal providers' get_sandbox read the sandbox's tunnels, which Modal refuses for a sandbox that has finished (killed, out of memory, past its timeout) with ConflictError, so teardown logged "was not terminated" and "may still be up" for a sandbox that was down. get_sandbox now returns a handle without tunnels when the tunnels call conflicts and poll() has an exit code. Modal's terminate() is a no-op on a finished sandbox, so teardown counts it as down. Modal raises ConflictError for other refusals too, so a conflict while poll() says running is raised as before. Co-Authored-By: Claude Opus 5.5 --- .../sandbox_providers/modal_sandbox.py | 16 +++++++-- .../sandbox_providers/modal_vm_sandbox.py | 5 ++- .../sandbox_providers/modal_sandbox_test.py | 33 +++++++++++++++++++ .../modal_vm_sandbox_test.py | 22 +++++++++++++ tst/unit/task/teardown_test.py | 21 ++++++++++++ 5 files changed, 92 insertions(+), 5 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index cc647e8d..6e1e902c 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -90,6 +90,17 @@ def _tunnel_url(t) -> str: return f"https://{t.host}" + (f":{t.port}" if t.port != 443 else "") +async def _tunnel_urls(sb: modal.Sandbox) -> Optional[dict[int, str]]: + """None once the sandbox has finished. Modal raises ``ConflictError`` for other refusals too, hence ``poll()``.""" + try: + tunnels = await sb.tunnels.aio() + except modal.exception.ConflictError: + if await sb.poll.aio() is None: + raise + return None + return {p: _tunnel_url(t) for p, t in tunnels.items()} + + def _modal_network_kwargs(policy: NetworkPolicy) -> dict: """``policy`` as Modal ``_experimental_create`` kwargs. @@ -475,8 +486,9 @@ async def create_sandbox( async def get_sandbox(self, sandbox_id: str) -> Sandbox: client = await self._get_client() sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client) - tunnels = await sb.tunnels.aio() - tunnel_urls = {p: _tunnel_url(t) for p, t in tunnels.items()} + tunnel_urls = await _tunnel_urls(sb) + if tunnel_urls is None: + return self._sandbox_cls(sb, {}) i6pn_address = await _resolve_i6pn_address(sb) return self._sandbox_cls(sb, tunnel_urls, i6pn_address=i6pn_address) diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index d615d39a..ce42e8c6 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -33,6 +33,7 @@ _log_sandbox_started, _resolve_app_base_name, _tunnel_url, + _tunnel_urls, _modal_network_kwargs, ) from agent_env.attribution import Attribution @@ -407,10 +408,8 @@ async def create_sandbox( async def get_sandbox(self, sandbox_id: str) -> ModalVmSandbox: client = await self._get_client() sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client) - tunnels = await sb.tunnels.aio() - tunnel_urls = {p: _tunnel_url(t) for p, t in tunnels.items()} # dockerd is already running on a live sandbox; do not re-run setup here. - return ModalVmSandbox(sb, tunnel_urls) + return ModalVmSandbox(sb, await _tunnel_urls(sb) or {}) # ── Object downloads (internal) ──────────────────────────────────────────────────── diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py index 3b1e9bea..ee86a566 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py +++ b/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,35 @@ 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")] + + +def _refusing_tunnels(*, exit_code: int | None) -> MagicMock: + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + sb.poll.aio = AsyncMock(return_value=exit_code) + sb.terminate.aio = AsyncMock() + sb.exec.aio = AsyncMock() + return sb + + +@pytest.mark.asyncio +async def test_a_sandbox_that_already_finished_gets_a_handle_without_tunnels_whose_terminate_takes_it_down( + _patched_provider, +): + sb = _refusing_tunnels(exit_code=137) + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(return_value=sb) + sandbox = await _patched_provider.get_sandbox("sb-finished") + await sandbox.terminate() + + assert sandbox.tunnel_urls == {} + sb.exec.aio.assert_not_called() + sb.terminate.aio.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_a_conflict_on_a_sandbox_still_running_is_raised(_patched_provider): + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(return_value=_refusing_tunnels(exit_code=None)) + with pytest.raises(modal.exception.ConflictError): + await _patched_provider.get_sandbox("sb-running") diff --git a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py index 514c5e29..7b206360 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/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,24 @@ 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("exit_code", [137, None], ids=["finished", "running"]) +@pytest.mark.asyncio +async def test_a_vm_that_already_finished_gets_a_handle_without_tunnels_and_a_conflict_while_running_is_raised( + exit_code, +): + provider = ModalVmSandboxProvider() + provider._get_client = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + sb.poll.aio = AsyncMock(return_value=exit_code) + with patch("agent_env.providers.sandbox_providers.modal_vm_sandbox.modal.Sandbox.from_id") as from_id: + from_id.aio = AsyncMock(return_value=sb) + if exit_code is None: + with pytest.raises(modal.exception.ConflictError): + await provider.get_sandbox("sb-running") + return + sandbox = await provider.get_sandbox("sb-finished") + + assert sandbox.tunnel_urls == {} diff --git a/tst/unit/task/teardown_test.py b/tst/unit/task/teardown_test.py index 1c7b5aea..13405913 100644 --- a/tst/unit/task/teardown_test.py +++ b/tst/unit/task/teardown_test.py @@ -8,9 +8,13 @@ import tempfile import time +from unittest.mock import AsyncMock, MagicMock + +import modal import pytest from agent_env.env.env import DeployedEnv, DeployedGatewayEnv +from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox, remove_local_work_dir from agent_env.task import teardown from agent_env.task.teardown import RecordedSandbox, TeardownReport, recorded_sandboxes, teardown_run @@ -144,6 +148,23 @@ def test_a_terminate_that_fails_is_reported_not_raised_and_the_others_still_go(t assert "sb-gw" not in context.metadata[TORN_DOWN_KEY] +def test_a_modal_sandbox_that_already_finished_is_counted_down_not_reported_up(monkeypatch): + sb = MagicMock() + sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) + sb.poll.aio = AsyncMock(return_value=137) + sb.terminate.aio = AsyncMock() + monkeypatch.setattr(modal.Sandbox, "from_id", MagicMock(aio=AsyncMock(return_value=sb))) + monkeypatch.setattr(ModalSandboxProvider, "_get_client", AsyncMock(return_value=MagicMock())) + context = TaskStepContext(deployed_sandboxes=[ + DeployedSandbox(sandbox_name="box", sandbox_id="sb-finished", sandbox_mode="container", sandbox_type="modal")]) + + report = _run(context) + + assert report.terminated == (RecordedSandbox("sb-finished", "modal", "sandbox"),) + assert report.still_up == () + assert context.metadata[TORN_DOWN_KEY] == ["sb-finished"] + + def test_cancelling_it_returns_what_it_reached_and_leaves_the_rest(terminated): terminated.behaviors["sb-gw"] = "hang" context = _context() From 0fd8ac4218e2aa24a94aa5fdde7464136bd903ce Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 15:06:13 -0700 Subject: [PATCH 2/3] test: clear the ruff findings the SmolVM provider left on main main's unit job fails at its lint step: the SmolVM checks import the local_stack fixture with a noqa for F401 but take it as a parameter, which ruff reports as F811, and one test has an f-string without placeholders. The checks now import it as ``local_stack as local_stack``, the explicit re-export another test already uses. Co-Authored-By: Claude Opus 5.5 --- tst/manual/smol_vm_agent_check.py | 2 +- tst/manual/smol_vm_gateway_check.py | 2 +- tst/unit/providers/sandbox_providers/smol_vm_provider_test.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tst/manual/smol_vm_agent_check.py b/tst/manual/smol_vm_agent_check.py index 9909d467..eda376aa 100644 --- a/tst/manual/smol_vm_agent_check.py +++ b/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 diff --git a/tst/manual/smol_vm_gateway_check.py b/tst/manual/smol_vm_gateway_check.py index 7b8730d2..a5534e33 100644 --- a/tst/manual/smol_vm_gateway_check.py +++ b/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 diff --git a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py index 85c0a4a6..55b4db63 100644 --- a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py +++ b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py @@ -57,7 +57,7 @@ class FakeMachine: id = "agentenv-test" def endpoint(self, port): - return SimpleNamespace(http_url=f"http://127.0.0.1:44000") + return SimpleNamespace(http_url="http://127.0.0.1:44000") def write_file(self, path, data): created.append((path, data)) From ef41dd97f274d86eec6200c157db5b55aa2867a7 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 14 Oct 2026 15:06:30 -0700 Subject: [PATCH 3/3] fix(providers): teardown terminates a sandbox by id, so one that already finished is counted down Teardown got a handle with get_sandbox before it terminated a sandbox. Modal's get_sandbox reads the sandbox's tunnels, which Modal refuses for a sandbox that has finished, so teardown reported a killed or timed-out sandbox as "may still be up". Making get_sandbox return a handle for a finished sandbox instead (this PR's first version) would have hidden dead containers from env reattach, which relies on get_sandbox raising. SandboxProvider.terminate_sandbox(sandbox_id) terminates a sandbox the provider made, by id; one that has already finished or is gone counts as terminated. The base default is get_sandbox then terminate. Every provider in the repo implements it without get_sandbox's extras: - modal, modal_vm: Sandbox.from_id then terminate (a no-op once finished), with NotFoundError counted as gone. - e2b: the class-level kill by id, which returns False for a sandbox already absent. - tenki: client.get then the session close TenkiSandbox.terminate does, now shared as close_session. - tensorlake: connect then terminate, skipping info(), with SandboxNotFoundError counted as gone. - sail_vm: through get_sandbox, since its terminate also deletes the model-key policy only the live Sailbox's applied policy names; NotFoundError counts as gone. - smol_vm: a machine Machine.list doesn't list is gone; a listed one is connected and deleted, skipping the Docker setup and host configuration get_sandbox runs. - local and chained: the base default. A chain's get_sandbox finds the owner; a per-provider terminate_sandbox loop would let the first provider count another provider's sandbox as gone. Both teardown paths call terminate_sandbox for a record that names its backend. A record that doesn't still goes through the slot provider's live handle, since "not found" from a provider that didn't make the sandbox doesn't mean it's gone. Local work-folder removal keys on the provider's ON_THIS_MACHINE. Co-Authored-By: Claude Opus 5.5 --- .../sandbox_providers/e2b/provider.py | 4 ++ .../sandbox_providers/modal_sandbox.py | 21 ++++---- .../sandbox_providers/modal_vm_sandbox.py | 9 +++- .../sandbox_providers/sail_vm/provider.py | 11 ++++ .../sandbox_providers/sandbox_provider.py | 5 ++ .../sandbox_providers/smol_vm/provider.py | 10 ++++ .../sandbox_providers/tenki/provider.py | 12 ++++- .../sandbox_providers/tenki/sandbox.py | 52 ++++++++++++------- .../sandbox_providers/tensorlake/provider.py | 7 +++ src/agent_env/task/teardown.py | 19 ++++--- .../task_steps/teardown_sandboxes.py | 12 +++-- tst/unit/bundle/run_teardown_test.py | 3 +- tst/unit/cli/eval_run_teardown_test.py | 3 +- .../e2b_sandbox_provider_test.py | 14 +++++ .../sandbox_providers/modal_sandbox_test.py | 35 ++++++------- .../modal_vm_sandbox_test.py | 25 ++++----- .../sail_vm_sandbox_provider_test.py | 13 +++++ .../smol_vm_provider_test.py | 17 ++++++ .../tenki_sandbox_provider_test.py | 25 ++++++++- .../tensorlake_sandbox_provider_test.py | 22 ++++++++ tst/unit/task/teardown_test.py | 28 +++++++++- tst/unit/task_step/test_teardown_sandboxes.py | 17 ++++++ 22 files changed, 280 insertions(+), 84 deletions(-) diff --git a/src/agent_env/providers/sandbox_providers/e2b/provider.py b/src/agent_env/providers/sandbox_providers/e2b/provider.py index 050d74a4..c5c11b4c 100644 --- a/src/agent_env/providers/sandbox_providers/e2b/provider.py +++ b/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) diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index 6e1e902c..cb09bd87 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -90,15 +90,12 @@ def _tunnel_url(t) -> str: return f"https://{t.host}" + (f":{t.port}" if t.port != 443 else "") -async def _tunnel_urls(sb: modal.Sandbox) -> Optional[dict[int, str]]: - """None once the sandbox has finished. Modal raises ``ConflictError`` for other refusals too, hence ``poll()``.""" +async def _terminate_by_id(client, sandbox_id: str) -> None: try: - tunnels = await sb.tunnels.aio() - except modal.exception.ConflictError: - if await sb.poll.aio() is None: - raise - return None - return {p: _tunnel_url(t) for p, t in tunnels.items()} + 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: @@ -486,12 +483,14 @@ async def create_sandbox( async def get_sandbox(self, sandbox_id: str) -> Sandbox: client = await self._get_client() sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client) - tunnel_urls = await _tunnel_urls(sb) - if tunnel_urls is None: - return self._sandbox_cls(sb, {}) + tunnels = await sb.tunnels.aio() + tunnel_urls = {p: _tunnel_url(t) for p, t in tunnels.items()} 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) diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index ce42e8c6..0cec9499 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -32,8 +32,8 @@ _attribution_tags, _log_sandbox_started, _resolve_app_base_name, + _terminate_by_id, _tunnel_url, - _tunnel_urls, _modal_network_kwargs, ) from agent_env.attribution import Attribution @@ -408,8 +408,13 @@ async def create_sandbox( async def get_sandbox(self, sandbox_id: str) -> ModalVmSandbox: client = await self._get_client() sb = await modal.Sandbox.from_id.aio(sandbox_id, client=client) + tunnels = await sb.tunnels.aio() + tunnel_urls = {p: _tunnel_url(t) for p, t in tunnels.items()} # dockerd is already running on a live sandbox; do not re-run setup here. - return ModalVmSandbox(sb, await _tunnel_urls(sb) or {}) + 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) ──────────────────────────────────────────────────── diff --git a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py b/src/agent_env/providers/sandbox_providers/sail_vm/provider.py index 76170d72..24dfd5e1 100644 --- a/src/agent_env/providers/sandbox_providers/sail_vm/provider.py +++ b/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() diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 522886ae..314de40b 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/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 diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/provider.py b/src/agent_env/providers/sandbox_providers/smol_vm/provider.py index 05f9c8f9..6c2c026d 100644 --- a/src/agent_env/providers/sandbox_providers/smol_vm/provider.py +++ b/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(): diff --git a/src/agent_env/providers/sandbox_providers/tenki/provider.py b/src/agent_env/providers/sandbox_providers/tenki/provider.py index bf00a79d..a52d5ce0 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/provider.py +++ b/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) diff --git a/src/agent_env/providers/sandbox_providers/tenki/sandbox.py b/src/agent_env/providers/sandbox_providers/tenki/sandbox.py index 841997b0..5031cfb5 100644 --- a/src/agent_env/providers/sandbox_providers/tenki/sandbox.py +++ b/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" diff --git a/src/agent_env/providers/sandbox_providers/tensorlake/provider.py b/src/agent_env/providers/sandbox_providers/tensorlake/provider.py index 087a97c7..cbd0ed46 100644 --- a/src/agent_env/providers/sandbox_providers/tensorlake/provider.py +++ b/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) diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index 263e2f0e..f01419fe 100644 --- a/src/agent_env/task/teardown.py +++ b/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: diff --git a/src/agent_env/task_step/task_steps/teardown_sandboxes.py b/src/agent_env/task_step/task_steps/teardown_sandboxes.py index 7f8f1539..e53ec765 100644 --- a/src/agent_env/task_step/task_steps/teardown_sandboxes.py +++ b/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() diff --git a/tst/unit/bundle/run_teardown_test.py b/tst/unit/bundle/run_teardown_test.py index 916dbaf2..40af715d 100644 --- a/tst/unit/bundle/run_teardown_test.py +++ b/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) diff --git a/tst/unit/cli/eval_run_teardown_test.py b/tst/unit/cli/eval_run_teardown_test.py index 6a07b939..c9575317 100644 --- a/tst/unit/cli/eval_run_teardown_test.py +++ b/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) diff --git a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py index 7ec2b44d..1cdfc4fe 100644 --- a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py +++ b/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() diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py index ee86a566..82b0b0d7 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_test.py @@ -440,33 +440,32 @@ async def collect(chunks, destination_path): assert written == [(b"LOCAL", "/data/a.json")] -def _refusing_tunnels(*, exit_code: int | None) -> MagicMock: +@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.poll.aio = AsyncMock(return_value=exit_code) sb.terminate.aio = AsyncMock() - sb.exec.aio = AsyncMock() - return sb - - -@pytest.mark.asyncio -async def test_a_sandbox_that_already_finished_gets_a_handle_without_tunnels_whose_terminate_takes_it_down( - _patched_provider, -): - sb = _refusing_tunnels(exit_code=137) with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox.from_id") as from_id: from_id.aio = AsyncMock(return_value=sb) - sandbox = await _patched_provider.get_sandbox("sb-finished") - await sandbox.terminate() + await _patched_provider.terminate_sandbox("sb-finished") - assert sandbox.tunnel_urls == {} - sb.exec.aio.assert_not_called() sb.terminate.aio.assert_awaited_once() + sb.tunnels.aio.assert_not_called() + sb.exec.aio.assert_not_called() @pytest.mark.asyncio -async def test_a_conflict_on_a_sandbox_still_running_is_raised(_patched_provider): +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(return_value=_refusing_tunnels(exit_code=None)) + 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-running") + await _patched_provider.get_sandbox("sb-finished") diff --git a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py index 7b206360..2a8ce590 100644 --- a/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/modal_vm_sandbox_test.py @@ -398,22 +398,17 @@ async def test_apps_are_looked_up_in_the_configured_environment(): assert lookup.aio.call_args.kwargs["environment_name"] == "rl-prod" -@pytest.mark.parametrize("exit_code", [137, None], ids=["finished", "running"]) +@pytest.mark.parametrize("found", [True, False], ids=["live or finished", "gone"]) @pytest.mark.asyncio -async def test_a_vm_that_already_finished_gets_a_handle_without_tunnels_and_a_conflict_while_running_is_raised( - exit_code, -): +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.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) - sb.poll.aio = AsyncMock(return_value=exit_code) - with patch("agent_env.providers.sandbox_providers.modal_vm_sandbox.modal.Sandbox.from_id") as from_id: - from_id.aio = AsyncMock(return_value=sb) - if exit_code is None: - with pytest.raises(modal.exception.ConflictError): - await provider.get_sandbox("sb-running") - return - sandbox = await provider.get_sandbox("sb-finished") - - assert sandbox.tunnel_urls == {} + 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() diff --git a/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py index bce5ff88..fb89a870 100644 --- a/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/sail_vm_sandbox_provider_test.py @@ -407,3 +407,16 @@ def test_a_rejected_key_can_be_retried(): assert _sdk.API_KEY_ENV not in os.environ _sdk.connect("sail-secret", "agent-env", sdk=sdk) assert sdk.reset_transports.call_count == 2 + + +@pytest.mark.asyncio +async def test_terminate_sandbox_terminates_through_the_handle_and_counts_a_sailbox_gone_as_terminated(): + sailbox = _sailbox() + sdk = _fake_sdk(sailbox) + provider = _provider(sdk) + + await provider.terminate_sandbox("sb_1") + sailbox.terminate.aio.assert_awaited_once() + + sdk.Sailbox.get.aio = AsyncMock(side_effect=sdk.NotFoundError("gone")) + await provider.terminate_sandbox("sb_gone") diff --git a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py index 55b4db63..274d504d 100644 --- a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py +++ b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py @@ -298,3 +298,20 @@ async def test_loopback_advertised_grant_uses_object_push(monkeypatch, tmp_path) await sandbox._download_object_to_vm(object_url, "/storage/image.tar.gz") sandbox._write_unsigned_object.assert_awaited_once_with(store, object_url, "/storage/image.tar.gz") + + +@pytest.mark.parametrize("listed", [True, False], ids=["listed", "gone"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_deletes_a_listed_machine_without_setting_it_up(monkeypatch, listed): + deleted = [] + connect = MagicMock(return_value=SimpleNamespace(delete=lambda: deleted.append(True))) + machines = [SimpleNamespace(name="agentenv-vm-1")] if listed else [] + monkeypatch.setattr(module, "Machine", SimpleNamespace(list=lambda conn: machines, connect=connect)) + setup = AsyncMock() + monkeypatch.setattr(module.SmolVmSandboxProvider, "_setup_docker", setup) + + await module.SmolVmSandboxProvider().terminate_sandbox("agentenv-vm-1") + + assert deleted == ([True] if listed else []) + assert connect.called is listed + setup.assert_not_called() diff --git a/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py index 9830f0b5..8a905ec6 100644 --- a/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/tenki_sandbox_provider_test.py @@ -21,7 +21,7 @@ vm_shape, ) from agent_env.providers.sandbox_providers.tenki.sandbox import DOCKER_BOOTSTRAP, TenkiSandbox -from tenki import SandboxError +from tenki import SandboxError, SessionNotFoundError from tst.unit.providers.sandbox_providers.tenki_fakes import FakeSession, fake_client_cls, result @@ -383,3 +383,26 @@ async def test_get_sandbox_restores_ports_policy_and_lifetime(client_cls): mode=ALLOWLIST, allow_hosts=("pypi.org",), allow_cidrs=("10.0.0.0/8",), ) assert 1700 <= sandbox.command_timeout <= 1800 + + +@pytest.mark.asyncio +async def test_terminate_sandbox_closes_the_session_by_id(client_cls): + provider = _provider(client_cls) + created = await provider.create_vm(setup_for_gateway=False) + session = client_cls.instances[0].sessions[created.sandbox_id] + + await provider.terminate_sandbox(created.sandbox_id) + + assert (session.closed, session.info.state) == (1, "TERMINATED") + + +@pytest.mark.asyncio +async def test_terminate_sandbox_counts_a_session_tenki_no_longer_knows_as_gone(client_cls, monkeypatch): + provider = _provider(client_cls) + client = provider._client() + + async def gone(sandbox_id): + raise SessionNotFoundError("gone") + + monkeypatch.setattr(client, "get", gone) + await provider.terminate_sandbox("sess-gone") diff --git a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py index 788a49d9..703fd224 100644 --- a/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/tensorlake_sandbox_provider_test.py @@ -451,3 +451,25 @@ async def test_get_sandbox_with_an_unknown_policy_fails_closed(): sandbox = await provider.get_sandbox("tl-test") assert sandbox.network_policy is None + + +@pytest.mark.asyncio +async def test_terminate_sandbox_connects_by_id_and_terminates_without_reading_its_info(): + raw = _raw_sandbox() + provider, sdk = _provider(raw) + + await provider.terminate_sandbox("tl-test") + + sdk.connect.assert_awaited_once() + raw.terminate.assert_awaited_once() + raw.info.assert_not_called() + + +@pytest.mark.parametrize("where", ["connect", "terminate"]) +@pytest.mark.asyncio +async def test_terminate_sandbox_counts_a_sandbox_tensorlake_no_longer_knows_as_gone(where): + raw = _raw_sandbox() + provider, sdk = _provider(raw) + (sdk.connect if where == "connect" else raw.terminate).side_effect = SandboxNotFoundError("tl-test") + + await provider.terminate_sandbox("tl-test") diff --git a/tst/unit/task/teardown_test.py b/tst/unit/task/teardown_test.py index 13405913..7ceb5c77 100644 --- a/tst/unit/task/teardown_test.py +++ b/tst/unit/task/teardown_test.py @@ -39,6 +39,8 @@ async def terminate(self): class _FakeProvider: + ON_THIS_MACHINE = False + def __init__(self, name, log): self.name, self.log = name, log @@ -47,6 +49,9 @@ async def get_sandbox(self, sandbox_id): raise RuntimeError(f"no sandbox {sandbox_id}") return _FakeSandbox(self, sandbox_id) + async def terminate_sandbox(self, sandbox_id): + await (await self.get_sandbox(sandbox_id)).terminate() + class _Log(list): """The (provider, sandbox_id) terminate log, and how given sandboxes misbehave.""" @@ -148,10 +153,30 @@ def test_a_terminate_that_fails_is_reported_not_raised_and_the_others_still_go(t assert "sb-gw" not in context.metadata[TORN_DOWN_KEY] +def test_a_record_that_names_its_backend_is_terminated_by_id_and_one_that_doesnt_through_a_handle( + terminated, monkeypatch, +): + by_id = [] + + async def terminate_sandbox(self, sandbox_id): + by_id.append((self.name, sandbox_id)) + + monkeypatch.setattr(_FakeProvider, "terminate_sandbox", terminate_sandbox) + context = TaskStepContext(deployed_sandboxes=[ + DeployedSandbox(sandbox_name="named", sandbox_id="sb-named", sandbox_mode="vm", sandbox_type="modal"), + DeployedSandbox(sandbox_name="unnamed", sandbox_id="sb-unnamed", sandbox_mode="vm"), + ]) + + report = _run(context) + + assert by_id == [("modal", "sb-named")] + assert terminated == [("default-sandbox", "sb-unnamed")] + assert {sandbox.sandbox_id for sandbox in report.terminated} == {"sb-named", "sb-unnamed"} + + def test_a_modal_sandbox_that_already_finished_is_counted_down_not_reported_up(monkeypatch): sb = MagicMock() sb.tunnels.aio = AsyncMock(side_effect=modal.exception.ConflictError("Sandbox has already finished with status terminated")) - sb.poll.aio = AsyncMock(return_value=137) sb.terminate.aio = AsyncMock() monkeypatch.setattr(modal.Sandbox, "from_id", MagicMock(aio=AsyncMock(return_value=sb))) monkeypatch.setattr(ModalSandboxProvider, "_get_client", AsyncMock(return_value=MagicMock())) @@ -163,6 +188,7 @@ def test_a_modal_sandbox_that_already_finished_is_counted_down_not_reported_up(m assert report.terminated == (RecordedSandbox("sb-finished", "modal", "sandbox"),) assert report.still_up == () assert context.metadata[TORN_DOWN_KEY] == ["sb-finished"] + sb.tunnels.aio.assert_not_called() def test_cancelling_it_returns_what_it_reached_and_leaves_the_rest(terminated): diff --git a/tst/unit/task_step/test_teardown_sandboxes.py b/tst/unit/task_step/test_teardown_sandboxes.py index de32bcda..63617d85 100644 --- a/tst/unit/task_step/test_teardown_sandboxes.py +++ b/tst/unit/task_step/test_teardown_sandboxes.py @@ -34,6 +34,9 @@ def __init__(self, name, terminated, failing): async def get_sandbox(self, sandbox_id): return _FakeSandbox(self, sandbox_id) + async def terminate_sandbox(self, sandbox_id): + await (await self.get_sandbox(sandbox_id)).terminate() + class _Log(list): """The (provider, sandbox_id) terminate log, plus the sandbox ids whose terminate should fail.""" @@ -141,6 +144,20 @@ async def test_an_env_tears_down_every_sandbox_behind_it_on_its_own_backend(term ] +@pytest.mark.asyncio +async def test_a_typed_sandbox_is_terminated_by_id_without_a_handle(terminated, monkeypatch): + async def no_handle(self, sandbox_id): + raise AssertionError(f"asked {self.name} for a handle to {sandbox_id}") + + async def by_id(self, sandbox_id): + self.terminated.append((self.name, sandbox_id)) + + monkeypatch.setattr(_FakeProvider, "get_sandbox", no_handle) + monkeypatch.setattr(_FakeProvider, "terminate_sandbox", by_id) + await _step(agent_names=["solver"]).execute(_context()) + assert terminated == [("modal", "sb-solver")] + + @pytest.mark.asyncio async def test_an_untyped_sandbox_uses_its_slot_default(terminated): await _step(sandbox_names=["gpu-box"]).execute(_context())